Cut Cross Entropy on JAX/Tunix TPU: Loss-Logits Memory Wall and Transfer Evidence

This report makes one narrow, reproducible claim: on LoRA SFT with JAX/Tunix on Cloud TPU v5e, Cut Cross Entropy removes a real loss-logits memory wall when that wall is active. Gemma3 270M is the exhaustive base case. Gemma3 1B and Gemma4 E2B are transfer checks under the same four-chip FSDP setup. Gemma3 4B and Gemma4 E4B are focused larger-model checks under an eight-chip FSDP setup, followed by a 12B/27B boundary check that deliberately probes where CCE stops being the dominant lever. Qwen3 0.6B is a focused non-Gemma transfer check.

Executive Summary

The short version is: CCE is not a universal memory optimizer and it is not mainly a speed optimization. It is a targeted replacement for the dense loss-logits path. When that path is the memory wall, CCE opens batch/context shapes that Default CE cannot compile. When another buffer becomes the wall, the visible gain is smaller.

The memory metric used throughout the new 270M plots is max per-chip XLA buffer-assignment planned HBM. This is the number that decides whether a TPU program fits on a chip. Runtime HBM snapshots are retained when the worker reported them, but they are not the primary frontier axis.

Claim Boundary

The report supports five claims:

1. CCE increases the feasible batch/context frontier when the dense loss logits tensor is the active memory wall.
2. At matched shapes, CCE can substantially reduce planned HBM while keeping train/eval loss in the same band.
3. The drop-in Tunix patch works on both single-chip and four-chip Gemma3 270M LoRA jobs, but throughput depends on mesh layout and chunk policy.
4. The same frontier pattern transfers to Gemma3 1B, Gemma4 E2B, Gemma3 4B, Gemma4 E4B, and a focused non-Gemma Qwen3 0.6B check under the tested setups.
5. The Gemma3 12B/27B boundary check is a negative transfer result: CCE still saves HBM, but the saved amount is not enough to move the fit boundary on the tested eight-chip setup.

It does not claim that CCE is always faster, that CCE removes every memory wall, or that the OPUS100 samples establish translation quality. Those rows are sanity checks for training behavior and generation-path integrity.

Why This Should Work

Default language-model cross entropy usually materializes a full logits tensor with shape roughly:

batch_size * context_length * vocabulary_size

That tensor grows directly with the two knobs we care about for SFT capacity: batch and context length. CCE computes the same cross-entropy objective by streaming over token/vocab chunks, avoiding the dense full-vocab logits tensor in the loss path.

The Tunix integration is a drop-in patch for the LoRA training path. During training, it intercepts the model output before the tied LM-head decode and computes CCE against the frozen LM head. During generation, the original decode path is restored before sampling. A regression test now guards that restore path because generation must not see the hidden-state intercept.

Experiment Scope

The one-chip rerun produced 307 result rows: frontier sweeps, pressure points, rank sensitivity, chunk tuning, one-step parity, and OPUS100 training runs. The four-chip follow-up added repeated mesh timings, frontier sweeps, HLO text scans, mixed-mesh chunk tuning, and a 1,000-step OPUS100 parity run.

The transfer checks were split into three tiers. The 1B/E2B tier reused the four-chip FSDP setup and retained short OPUS100 sanity rows. The 4B/E4B tier used an eight-chip FSDP setup and focused on frontier and chunk behavior. The 12B/27B tier reused the same eight-chip setup as a boundary check: it asks whether CCE alone still moves the fit frontier once the base model is much larger.

Those transfer rows are deliberately smaller than the 270M evidence package. Their purpose is to test whether the same memory/frontier story survives model scale and Gemma generation, and then to mark where CCE alone stops being the active boundary mover.

1. Frontier: CCE Moves the Fit Boundary

The frontier sweep is deliberately synthetic. It is not intended to prove translation quality. Its job is to isolate the batch * context * vocab memory pressure and ask whether the same model, TPU, and LoRA configuration can run when only the loss implementation changes.

The left panel is the practical result: maximum completed context length by batch size. The right panel shows the underlying XLA planned HBM pressure for representative batch sizes.

The important pattern is not that CCE makes every shape fit. It moves the wall. For small batches, both variants can still be limited by other model or activation buffers. As batch grows, the loss-logits tensor becomes the visible wall, and CCE opens shapes that Default CE cannot compile.

The pass/fail map shows that the sweep did not stop at the first Default CE failure. CCE was pushed until it reached its own boundary.

2. Pressure Points: Same Shape, Different Memory Wall

The pressure-point rows keep the sweep readable by focusing on a few representative shapes.

This is the cleanest systems result: CCE turns multiple Default CE OOM rows into completed TPU train steps, and the remaining CCE failures occur at much higher pressure points.

3. Rank and Chunk Checks

CCE should mostly track batch * context * vocab, not LoRA rank. The rank sweep used ranks 4, 16, and 64 across b8/b16/b32/b64 and L512/L1024/L2048/L4096. The frontier pattern was unchanged across ranks: for example, b16 moved from L512 to L1024 and b32 moved from no Default CE fit at L512 to CCE fitting through L1024.

The chunk sweep asks a different question: once CCE is enabled, which token/vocab chunk sizes are faster at a fixed shape? On b16/L512, all tested chunk pairs used roughly the same XLA planned HBM, but step time varied.

The best b16/L512 chunk rows were:

The conservative package knobs are 128/8192 because they are stable across shapes. The one-chip sweep suggests that larger token chunks can be faster for this particular 270M b16/L512 case, but chunk tuning should be shape-aware. The four-chip follow-up below shows why: the best chunk policy can change when the mesh changes.

4. Multi-Chip Mesh Generalization and the Mixed-Mesh Outlier

The main rerun is intentionally single-chip so that the memory story is easy to read. We then checked whether the patch still works when Tunix/JAX shards the same Gemma3 270M LoRA job over four TPU chips.

The result is not a single-chip artifact. CCE worked in the FSDP-only, mixed-FSDP/TP, and TP-only layouts. At b16/L512, all three meshes reduced max per-chip XLA planned HBM, but throughput behavior was mesh-specific:

That table separates three facts that are easy to blur together:

1. CCE still saves memory in every tested mesh.
2. TP-heavy layouts can be slow even without CCE on this small model.
3. The dramatic CCE throughput outlier is local to the mixed fsdp=2,tp=2 layout with the conservative default chunk settings.

The broader four-chip frontier sweep still shows the desired capacity behavior:

To inspect the fsdp=2,tp=2 outlier, we kept full train-step XLA dumps and scanned the optimized HLO text. This scan is intentionally a coarse compiler artifact check, not a full profiler. It did not show a simple collective-count explosion large enough to explain the 97x slowdown by itself.

The more actionable finding came from chunk tuning. The package default (token_chunk=128, vocab_chunk=8192) is conservative and portable, but it is too fine-grained for this mixed mesh. On fsdp=2,tp=2, b16/L512, increasing the chunk size kept XLA HBM at 2.65 GiB/chip while dramatically reducing step time:

The same grid also gives a one-factor-at-a-time view. At b16/L512, holding token_chunk=128 fixed and increasing vocab_chunk made the run faster; holding vocab_chunk=8192 fixed and increasing token_chunk showed the same pattern. XLA planned HBM stayed at 2.65 GiB/chip for every row in these two slices.

So the mixed-mesh slowdown was not caused by only one bad axis. The problematic case was excessive chunk granularity: many small token chunks multiplied by many small vocab chunks. The larger 512/65536 preset reduces both loop dimensions, which is why it fixes the outlier so strongly.

The clean interpretation is therefore: CCE is compatible with FSDP/TP sharding, and the large mixed-mesh slowdown is mostly a chunk-policy problem rather than a correctness or memory failure. The package keeps the conservative 128/8192 default, but adds an opt-in TUNIX_ACCEL_CE_PRESET=tpu_large_chunks shortcut for the validated 512/65536 TPU setting.

5. Transfer Checks: Scaling the Claim

After the 270M base case was complete, we asked whether the same systems claim survives larger Gemma models. The transfer checks are intentionally narrower than the 270M package: they test whether CCE still moves the fit frontier under the chosen mesh, not whether every 270M ablation repeats at every model size.

Four-Chip 1B/E2B Checks

Gemma3 1B and Gemma4 E2B used v5litepod-4, four chips, with fsdp=4,tp=1 as the primary mesh.

The transfer result is directionally the same as the 270M result: CCE extends the maximum passing context at the same batch size, and the benefit is most visible when Default CE hits the loss-logits wall.

For matched OPUS100 sanity runs, CCE again reduced max per-chip XLA planned HBM while leaving train/eval loss in the same band. These are 1,000-step sanity rows only; they are not translation-quality claims.

The same transfer run also reinforces the chunk-policy lesson from the 270M mixed-mesh outlier. At FSDP-only 1B/E2B shapes, chunk tuning has modest upside; at the 270M mixed mesh, the conservative 128/8192 setting was the dominant throughput problem. The E2B probe shows the other side of the story: TP-heavy meshes can be impractically slow even when the model fits in memory, so larger models should prefer FSDP-only unless there is a separate reason to pay for TP.

Eight-Chip 4B/E4B Checks

We then ran a focused larger-model transfer check on v5litepod-8, fsdp=8,tp=1. This was not another exhaustive sweep. It kept the same question: when the model is larger and the TPU has more chips, do we still see CCE opening shapes that Default CE cannot compile?

The answer is yes, with the expected boundary: CCE helps when the failed shape is near the loss-logits wall. It does not magically make every batch/context pair fit.

The pressure rows show the same story in the raw memory metric used throughout the report: max per-chip XLA buffer-assignment planned HBM.

The Gemma3 4B and Gemma4 E4B rows also show why CCE savings should not be predicted from parameter count alone. At the same b4/L512 shape, Gemma3 4B completed under both variants and went from 9.58 to 8.58 GiB/chip. Gemma4 E4B failed under Default CE with a 20.30 GiB/chip compile plan, then completed under CCE at 14.53 GiB/chip. Both are "4B-class" models, but their XLA graphs do not put the same fraction of peak memory in the Default CE loss path. The larger E4B gain is therefore not a generic "4B model" rule; it is a graph-specific result where the failed Default CE compile plan was dominated more heavily by loss-logits buffers.

The larger-model chunk tuning result is also useful, but more modest than the 270M mixed-mesh outlier. At these FSDP-only pressure points, chunk choice mostly changes step time by a few percent while planned HBM stays in the same band.

The important practical rule is: use CCE to move the memory wall, then tune token_chunk and vocab_chunk on the exact model/mesh/shape you intend to run. The memory savings transfer, but the fastest chunk policy does not have to be universal.

Eight-Chip 12B/27B Boundary Rerun

The next question was whether the 4B/E4B pattern continues automatically at 12B and 27B. We reran a focused boundary grid on the same single-host Cloud TPU v5litepod-8, eight chips, fsdp=8,tp=1, LoRA rank 16, synthetic two-step probe. The result is a useful negative boundary: CCE still reduces planned HBM, but it no longer creates a new passing shape.

The boundary memory plot shows why. Near the pass/fail edge, CCE saves memory, but the saved amount is too small to cross the 16 GiB/chip fit line.

This does not contradict the 4B/E4B result. It narrows the claim. CCE is most visible when the dense loss-logits path is the active peak-memory wall. In the 12B/27B single-host eight-chip setup, the model graph is already close to or over the per-chip HBM limit before the loss path can become the decisive buffer. CCE still trims the graph, but it is a partial lever rather than a frontier-changing one.

After fixing the runner to call jax.distributed.initialize() before Tunix model setup, we ran actual multi-host Tunix smoke tests for the larger models. These are not frontier sweeps; they only prove that the same drop-in CCE path survives real distributed Tunix launch on larger TPU slices.

The distributed logs confirm process_count=4, global_devices=16 for 12B and process_count=8, global_devices=32 for 27B. The runtime peak in this table is the artifact host's local four-chip aggregate, not the whole pod.

Non-Gemma Qwen3 0.6B Check

The Gemma rows test model scale inside one family. To check whether the patch is only a Gemma adapter trick, we also ran a focused Qwen3 0.6B LoRA sweep using the HuggingFace model path Qwen/Qwen3-0.6B, Cloud TPU v5litepod-1, one chip, and fsdp=1,tp=1.

The Qwen result is smaller than the Gemma3 270M package, but it answers the right transfer question. At b4, Default CE fit L512 and failed L1024; CCE fit L1024. At the matched b4/L512 shape, planned HBM dropped from 8.40 to 7.08 GiB/chip. At the newly opened b4/L1024 shape, Default CE planned 16.15 GiB/chip and failed, while CCE completed at 14.13 GiB/chip.

Rank and chunk behavior also looked familiar. For b4/L1024, LoRA rank 16 and 64 failed under Default CE and completed under CCE. Vocab chunk choice barely moved planned HBM, while step time varied enough to justify shape-specific chunk tuning.

The negative row is also useful: b1/L4096 remained OOM under aggressive CCE chunk settings, with plans around 20.5 GiB/chip. In other words, Qwen3 0.6B shows the same bounded story as Gemma: CCE moves the frontier when the loss logits are the active wall, but once another buffer dominates, smaller chunks alone do not solve the shape.

6. Real EN-FR Training Parity

The systems sweeps prove memory behavior, not usefulness. For a training sanity check, we ran OPUS100 EN-FR LoRA SFT.

Same-shape A/B:

The loss curves follow the same trajectory. The normalized metric panel shows the tradeoff clearly: much lower planned HBM, similar train/eval loss, and slower same-shape steps.

BLEU and chrF here are deliberately labeled as sanity metrics. They use 16 evaluation examples, so they are useful for catching obvious generation breakage, not for making a translation-quality claim.

We also repeated a shorter 1,000-step OPUS100 run on the four-chip FSDP-only mesh to check that sharded training behaves normally:

This is not a replacement for the longer one-chip 5,000-step parity result. Its purpose is narrower: the same CCE patch trains normally under the four-chip FSDP sharding setup while preserving the expected memory/speed tradeoff.

We also ran a CCE capacity row: b64/L512 for 1,250 steps, matching the rough token budget of b16/L512 for 5,000 steps. It fit at 14.13 GiB/chip and finished in 19.2 minutes. That row did not create a speed win in this setup; it mainly shows that CCE can spend the saved memory on a much larger batch.

7. Generation Sanity Samples

The restored generation path produced normal text for both variants. The outputs below are not polished translation results; the point is narrower: after the same training recipe, CCE and Default CE are in the same qualitative band and generation no longer sees the hidden-state intercept used by the loss patch.

Sample 1

• EN: They met without me.
• Reference: Ils se sont rencontrés sans moi.
• Default CE: - Ils ne sont pas venus pour moi.
• CCE: Ils ne sont pas venus sans que je ne soit en aware.

Sample 2

• EN: Pay cash week .
• Reference: J'ai payé en liquide pour la semaine.
• Default CE: Pay week .
• CCE: Paye de la semaine.

Sample 3

• EN: The certifrcate can be printed in one or more of the languages of the Convention andshould be completed in one of these languages.
• Reference: Le certificat peut être imprimé dans une ou plusieurs langues de la convention et doit être complété dans l'une de ces langues.
• Default CE: La certificat peut être imprimé dans un des langues de la Convention et complété dans une des langue
• CCE: La certificatabilité peut être imprimée dans une ou des langues de la Convention et peut être compli

Sample 4

• EN: The other issue, which is of enormous significance, is the right to the veto.
• Reference: L'autre question, d'une portée énorme, est le droit de veto.
• Default CE: L'autre aspect important du droit à veto, qui est de grande importance, est le droit à l'apanche.
• CCE: Encore heureux, on questionne aussi le droit à la proposition, à cette aussi, du autre facteur qui rend cette session si importante : l'énoncé de fait du veto.

Sample 5

• EN: Therefore, let me summarize, at least from the vantage point of one Finance Minister.
• Reference: Par conséquent, permettez-moi d'en résumer la teneur, à tout le moins en tant que ministre des finances.
• Default CE: Par conséquent, je veux vous rappeler, d'un point de vue financier.
• CCE: Par conséquent, je veux te rappeler au moins l'angle de l'examinateur.

The complete 32-row side-by-side sample file is retained as 01-CCE/data/gemma3_270m_full_cce/generation_samples.jsonl.

8. Interpretation and Limitations

CCE is a strong memory lever when the loss logits are a dominant tensor. That is why the effect becomes clearer as batch and context grow. It is not a universal memory solution: after CCE removes the dense logits tensor, the next wall can be activations, attention, model state, sharding layout, or compile-time buffer planning.

The tradeoff is also real. CCE streams reductions instead of doing one dense logits computation, so same-shape step time can increase. On the 270M b16/L512 quality run, the slowdown was about 1.85x. The value of CCE is therefore not "always faster." The value is "more trainable shapes before OOM," and sometimes that extra capacity may be worth the slower step.

The mesh check adds one more caveat: memory savings transfer across FSDP/TP layouts, but throughput does not transfer automatically. TP-heavy layouts can make a small model communication-bound, and conservative chunk settings can make mixed FSDP/TP layouts much slower than necessary. The right comparison is therefore always same-model, same-mesh, Default CE versus CCE, with chunk policy recorded as part of the system configuration.

9. Reproducibility Artifacts

The retained files are intentionally compact: enough to redraw the report and audit every claim, without keeping extracted TPU dumps or checkpoints.

Single-chip evidence package:

• 01-CCE/data/gemma3_270m_full_cce/run_manifest.csv
• 01-CCE/data/gemma3_270m_full_cce/all_runs.csv
• 01-CCE/data/gemma3_270m_full_cce/frontier_summary.csv
• 01-CCE/data/gemma3_270m_full_cce/pressure_points.csv
• 01-CCE/data/gemma3_270m_full_cce/rank_frontier_summary.csv
• 01-CCE/data/gemma3_270m_full_cce/chunk_tuning.csv
• 01-CCE/data/gemma3_270m_full_cce/training_history.csv
• 01-CCE/data/gemma3_270m_full_cce/training_summary.csv
• 01-CCE/data/gemma3_270m_full_cce/generation_metrics.csv
• 01-CCE/data/gemma3_270m_full_cce/generation_samples.jsonl
• 01-CCE/data/gemma3_270m_full_cce/oom_events.csv

Four-chip mesh evidence package:

• 01-CCE/data/gemma3_270m_mesh_cce/mesh_summary.csv
• 01-CCE/data/gemma3_270m_mesh_cce/matched_memory.csv
• 01-CCE/data/gemma3_270m_mesh_cce_repeat/repeat_summary.csv
• 01-CCE/data/gemma3_270m_4chip_frontier/frontier_summary.csv
• 01-CCE/data/gemma3_270m_outlier_hlo/hlo_op_counts.csv
• 01-CCE/data/gemma3_270m_4chip_chunk/chunk_summary.csv
• 01-CCE/data/gemma3_270m_4chip_chunk/chunk_axis_ablation.csv
• 01-CCE/data/gemma3_270m_4chip_quality/training_summary.csv

Gemma3 1B / Gemma4 E2B transfer package:

• 01-CCE/data/gemma_1b_e2b_cce_transfer/run_manifest.csv
• 01-CCE/data/gemma_1b_e2b_cce_transfer/transfer_runs.csv
• 01-CCE/data/gemma_1b_e2b_cce_transfer/frontier_summary.csv
• 01-CCE/data/gemma_1b_e2b_cce_transfer/matched_metrics.csv
• 01-CCE/data/gemma_1b_e2b_cce_transfer/chunk_summary.csv
• 01-CCE/data/gemma_1b_e2b_cce_transfer/training_summary.csv
• 01-CCE/data/gemma_1b_e2b_cce_transfer/training_history.csv
• 01-CCE/assets/gemma_cce_transfer_frontier.png
• 01-CCE/assets/gemma_cce_transfer_quality.png
• 01-CCE/assets/gemma_cce_transfer_chunk_mesh.png

Gemma3 4B / Gemma4 E4B focused transfer package:

• 01-CCE/data/gemma_4b_e4b_cce_transfer/run_manifest.csv
• 01-CCE/data/gemma_4b_e4b_cce_transfer/transfer_runs.csv
• 01-CCE/data/gemma_4b_e4b_cce_transfer/frontier_summary.csv
• 01-CCE/data/gemma_4b_e4b_cce_transfer/matched_metrics.csv
• 01-CCE/data/gemma_4b_e4b_cce_transfer/pressure_points.csv
• 01-CCE/data/gemma_4b_e4b_cce_transfer/chunk_summary.csv
• 01-CCE/assets/gemma_cce_large_transfer_frontier.png
• 01-CCE/assets/gemma_cce_large_transfer_pressure.png
• 01-CCE/assets/gemma_cce_large_transfer_chunk_tuning.png

Gemma3 12B / 27B focused boundary package:

• 01-CCE/data/gemma3_12b_27b_cce_focused/run_manifest.csv
• 01-CCE/data/gemma3_12b_27b_cce_focused/frontier_runs.csv
• 01-CCE/data/gemma3_12b_27b_cce_focused/frontier_summary.csv
• 01-CCE/data/gemma3_12b_27b_cce_focused/boundary_hbm_points.csv
• 01-CCE/data/gemma3_12b_27b_cce_focused/matched_metrics.csv
• 01-CCE/assets/gemma3_12b_27b_cce_focused_frontier.png
• 01-CCE/assets/gemma3_12b_27b_cce_focused_hbm.png

Qwen3 0.6B focused non-Gemma transfer package:

• 01-CCE/data/qwen3_0p6b_cce_transfer/run_manifest.csv
• 01-CCE/data/qwen3_0p6b_cce_transfer/frontier_summary.csv
• 01-CCE/data/qwen3_0p6b_cce_transfer/matched_boundary.csv
• 01-CCE/data/qwen3_0p6b_cce_transfer/rank_sensitivity.csv
• 01-CCE/data/qwen3_0p6b_cce_transfer/chunk_tuning.csv
• 01-CCE/data/qwen3_0p6b_cce_transfer/aggressive_l4096_negative.csv
• 01-CCE/assets/qwen3_0p6b_cce_transfer.png
• 01-CCE/assets/qwen3_0p6b_cce_chunk_tuning.png

Lower-level audit files are also kept beside those summaries, including frontier_runs.csv, rank_sensitivity.csv, mesh_runs.csv, repeat_runs.csv, and per-run manifests. Compressed raw worker artifacts are retained under the corresponding raw_artifacts/ directories. Extracted raw/ directories are reproducible from those tarballs and should not be committed.

The Qwen3 CCE transfer tables are compact reconstructions from retained worker log values. The TPU became SSH-unstable during raw artifact recovery and was deleted to avoid leaving an unattended VM running.