Skip to content

Running SFT on TPU: a known-good starting point

Olmo 3.5 small (the 4T EMO base) on Dolci-Think with H045's recipe, on one TPU v4 4x4x8 slice (128 chips) of the gke-tpu-v4 cluster, through MaxText + Tunix.

Measured cost and outcome: 11,768 steps of 16 × 65,536 tokens at a median 3.94 s/step, 13.5 h wall clock from submit to the last step. The loss matches H045's GPU runs (parity windows in open-instruct#1933). The GPU recipe takes 5 h 54 m on 2×8 H100, so the TPU path is slower per run. Its advantage is capacity: four of these runs fit at once without touching Beaker.

The recipe is scripts/tpu/maxtext/k8s/olmo35-sft-4x4x8/. Its kustomization.yaml header records why each setting is what it is; read it before editing. Olmo 3 7B has its own overlay, k8s/instruct-sft-4x4x8/; its base k8s/olmo3-sft/ targets a 2x2x4 slice, which no longer exists. scripts/tpu/maxtext/README.md lists the MaxText patches.

Example runs

Run What it is
olmo35-sft-h045-tpu-full-r2 This recipe, verbatim: the TPU replica of the ledger's H045 anchor (arm anchor-emo-olmo35-tpu). 0 restarts, final CE 0.9245. HF export: /weka/oe-adapt-default/abhishekr/tpu-posttrain/hf/olmo35-sft-h045-tpu-full-r2-step11768.
H067, H072 (olmo-post-training-ledger) The same recipe with other data: arms, seeds, resumes and their checks are recorded in those hypothesis files and in open-instruct#1976 and #1977.

Before you start

  • IAM on the ai2-olmo GCP project, e.g. roles/container.developer.
  • A network the cluster accepts. allenai/olmo-gcp-infra's README lists them: the office IPs, its Tailscale exit nodes, and Beaker sessions on ai2/hammond. From anywhere else, use the DNS endpoint (--dns-endpoint). Then:
gcloud container clusters get-credentials gke-tpu-v4 --zone us-central2-b --project ai2-olmo
kubectl -n olmo get workloads

kubectl needs gke-gcloud-auth-plugin (gcloud components install gke-gcloud-auth-plugin, or a binary in ~/.local/bin). - gcloud logins expire about daily, and kubectl then fails quietly. Watchers go blind without an error. Agents use the service account tpu-posttrain-agent@ai2-olmo.iam.gserviceaccount.com: export CLOUDSDK_CORE_ACCOUNT to it in every process, background watchers included. Move ~/.kube/gke_gcloud_auth_plugin_cache aside when switching identity, or kubectl keeps the cached one. People log in with their own account. - Agents. Claude Code and Codex in this repo load a summary of this guide's rules from the tpu skill. Update it when you change a rule here. - Where things live. Everything is in us-central2. Data, tokenizers and checkpoints are in gs://ai2-olmo-dev/post-training/, and the nodes mount it at /mnt/post-training through gcsfuse. Images are in us-docker.pkg.dev/ai2-olmo/olmo/maxtext-posttrain.

What the cluster is

gke-tpu-v4 holds 4,096 reserved v4 chips at 4 per host. Since 2026-10-21 they are cut into seven 4x4x8 slices (128 chips) and three 8x8x16 (1,024), with 128 chips held back so GKE can auto-repair failed nodes; the 8x8x8 is gone. The layout changes as the cluster's admins reslice it, so check it before planning: kubectl get nodes -L cloud.google.com/gke-nodepool --no-headers | awk '{print $NF}' | sort | uniq -c (32 hosts per 4x4x8, 256 per 8x8x16). A slice runs one job at a time. Kueue admits jobs from four queues in one cohort:

Queue Who Preempted by
hero the main pretraining run; owns all the quota nothing
dev tests hero
post-training, data-adaptation us; share what is left hero, dev

Use post-training. Its jobs run on borrowed quota and are preempted whenever hero or dev needs the slice, so checkpoint often enough to lose little. Never hold a slice in a way that delays hero.

The recipe needs exactly a 4x4x8: its mesh puts one 64K row on 8 chips, so 128 chips hold its 16 rows. The measurements below on the former 8x8x8 ran the same mesh at 64 rows per step, which is a different recipe (H078's 4x batch). On the former 8x8x8, the measured 16-row configuration used olmo35-cp-ep, splitting each row 32 ways and changing the MoE load-balance term. No validated 16-row configuration for the 8x8x16 is documented here. Measured meshes are in Changing the recipe.

Cost per slice

From TensorBoard perf/*. Throughput is tokens per chip-second: tokens per device per step divided by the mean step time over steps 31-100, so a stalled step counts at its full cost.

slice config rows/step median s/step tokens/s/chip steps over 3x median
4x4x8 the recipe (A0') 16 3.88 2,055 0
8x8x8 H078's config, with the data-order shim (H088 opt-a) 64 4.28 1,386 4.3% (13% over H078's whole runs)
8x8x8 same, no shim and no batch hashes (H088 opt-b) 64 4.23 1,887 0
8x8x8 opt-b + remat_policy=save_dot_except_mlp (H088 opt-c) 64 4.18 1,900 0; a tie, keep full

Startup before step 1 (pods, restore, an 11-12 min compile) adds 20-35 min on either slice, paid again at every resume.

  • Use a 4x4x8 for this recipe. Among the measured configurations it used the fewest chip-hours per token; no throughput is measured here for the 8x8x16, which is hero's and borrowed only. The 8x8x8 was split into 4x4x8s on 2026-10-21, so the 8x8x8 results below are historical. The 8x8x8 overlay is k8s/olmo35-sft-8x8x8/: the same token budget at 64 rows and 2,942 steps. Its data order differs from the 4x4x8 recipe's, so an arm on it needed a control on the same slice and config. Without the shim the 8x8x8 is projected at ~1.09x the chip-hours of the 4x4x8, finishing a run ~3.6x sooner (H088, allenai/open-instruct#2020). This is a steady-state projection, not a measured cost. Both rates come from stall-free 70-step windows on different days (4x4x8 10-14, 8x8x8 10-18). The 8x8x8's stall rate over a whole run without the shim is not measured. The 8x8x8 could not run the 16-row recipe at that rate (see Changing the recipe).
  • Historical 8x8x8 input stalls. H088 removed the data-order shim (MAXTEXT_R2_DATA_EMULATION) and MAXTEXT_LOG_BATCH_HASH=1 together, so their separate contributions are not isolated. With both, one step in eight stalled for 20-35 s, at the same steps in every run, and H078's runs cost 1.8-1.9x the chip-hours of a 4x4x8. In xprof the stall is one host's input pipeline (grain process_prefetch waiting 29 s) while every TPU waits for its batch in jit__device_put_reshard. The HF input pipeline allows at most one grain worker, so more workers are not a fix.
  • Keep remat_policy=full. Saving dot outputs except the MLP's gained 0.7%, with the same loss.
  • Outside stalls a step is ~4.0 s of device compute and ~0.2-0.5 s idle. FSDP all-gather waits are ~16% of the compute, and num_vocab_tiling=16 costs ~4% (the 205 MB embedding is all-gathered and its gradient all-reduced 96 times a step).
  • Both slices run at ~4% of v4's 275 TFLOP/s bf16 peak by MaxText's own FLOP count.
  • jax memory_stats() peaks (0.9-2.3 GiB) leave out XLA temporaries: the same 8x8x8 config needed 32.17 GB at compile with ragged_buffer_factor=4.0. Read HBM headroom from the compile error or xprof, not from those lines.

1. Image

us-docker.pkg.dev/ai2-olmo/olmo/maxtext-posttrain:51f8cc4b-p907d86f3-rfb46e215e

This tag is r2's image (51f8cc4b-p907d86f3, built from image-olmo35/) plus image-olmo35-rf/. That patch adds a fix for resuming from a checkpoint (without it, restore runs out of HBM) and batch-hash logging. Training numerics are unchanged: a replay produced bit-identical CE and 9,952 of 9,952 equal batch hashes. Use it for every new run. The Olmo 3 image is built by image/build.sh <maxtext commit>.

To build a new image, run build.sh from any host with gcloud (a Beaker session works). Cloud Build carries the context base64-encoded inside the build config. Builds use Cloud Build's default service account. Set BUILD_SERVICE_ACCOUNT to a full projects/<project>/serviceAccounts/<email> resource only when the caller is authorized to use that account.

2. Base checkpoint

The loader reads an Orbax parameter tree (…/0/items) from GCS.

  • Olmo 3.5 (OLMo-core distcp). This runs in two stages, on CPU:
python scripts/tpu/maxtext/olmo35/dump_olmocore_masters.py <step dir> <masters dir>              # needs olmo_core
python scripts/tpu/maxtext/olmo35/convert_olmocore_to_maxtext.py <step dir> <masters dir> <out>   # needs maxtext
gcloud storage cp -r <out>/0 gs://ai2-olmo-dev/post-training/checkpoints/<name>/

It keeps OLMo-core's padded 100,352 vocabulary, whose padded rows were trained. The 4T EMO base is already at gs://ai2-olmo-dev/post-training/checkpoints/olmo35-small-4t-emo-step5961/0/items. - Olmo 3 (HF). Use MaxText's converter with model_name=olmo3-7b. olmo3-7b-pt has no YaRN, so it is wrong for the 1025 base.

JAX_PLATFORMS=cpu python -m maxtext.checkpoint_conversion.to_maxtext model_name=olmo3-7b base_output_directory=<out> \
  scan_layers=True use_multimodal=false hardware=cpu skip_jax_distributed_system=true \
  checkpoint_storage_use_zarr3=1 checkpoint_storage_use_ocdbt=1 --hf_model_path=allenai/Olmo-3-1025-7B --save_dtype=bfloat16

Check any new conversion numerically against the source model before training on it. Key and shape checks pass on broken conversions.

3. Data and tokenizer

MaxText streams parquet files. Its shuffle only permutes files and mixes a small buffer, and hub shards are grouped by source, so shuffle globally first. olmo35/build_full_dolci_think.py (Olmo 3.5) and data/build_full_dolci.py (Olmo 3) write N shuffled shards from one dataset revision and use the standard HF cache settings. Pin the input with --revision <dataset commit>, or use --data-files '<local parquet glob>' for local files; the output directory must be new or empty. Upload the shards to gs://ai2-olmo-dev/post-training/data/<name>/. Ledger experiments shard each arm and seed into 512 files with experiments/h067/shard.py in olmo-post-training-ledger.

The tokenizer is a directory on the bucket. For Olmo 3.5, olmo35/tok_oi35.py saves H045's tokenizer and template (allenai/dolma2-tokenizer-olmo35 @8b971706, olmo35). For Olmo 3, save PreTrainedTokenizerFast.from_pretrained("allenai/olmo-3-tokenizer-instruct-dev") (see Traps).

4. Launch

Copy k8s/olmo35-sft-4x4x8/ to a new directory. Then rename the JobSet and both ConfigMaps (in jobset.yaml and kustomization.yaml) to a new unique name, and edit the env block:

Variable Meaning
RUN_NAME, OUTPUT_DIR Checkpoints go to $OUTPUT_DIR/$RUN_NAME/checkpoints/<step>/.
LOAD_PARAMETERS_PATH, TOKENIZER_PATH, TRAIN_FILES Step 2 and step 3 outputs (/mnt/post-training/... = gs://ai2-olmo-dev/post-training/...).
TOTAL_DEVICES, GLOBAL_BATCH, PER_DEVICE_BATCH 128 chips × 0.125 = 16 sequences per step, with no gradient accumulation. The launcher derives accumulation and fails if the numbers don't divide.
SEQ_LEN 65,536. Tokens per step = GLOBAL_BATCH × SEQ_LEN.
STEPS, LR_SCHEDULE_STEPS Steps trained, and the horizon the LR decays over (warmup WARMUP_FRAC of it). Set STEPS below the horizon to run the start of a longer recipe, e.g. for a loss-parity check.
NUM_EPOCH Passes the stream holds. STEPS decides what is trained; a pass shorter than STEPS repeats data.
EXTRA_ARGS MaxText flags. The mesh is custom_mesh_and_rule=olmo35-ep-as-cp ici_fsdp_parallelism=16 ici_expert_parallelism=8: batch over FSDP, and the expert axis doubles as context parallelism outside the MoE. checkpoint_period and max_num_checkpoints_to_keep set what is saved.
images: newTag The image tag from step 1. 
kubectl --context gke_ai2-olmo_us-central2-b_gke-tpu-v4 apply -k <your overlay dir>
kubectl -n olmo get workloads              # ADMITTED True = it has a slice
kubectl -n olmo get jobset <name>          # TERMINALSTATE, RESTARTS, SUSPENDED (a hero preemption suspends it)
kubectl -n olmo get pods -l jobset.sigs.k8s.io/jobset-name=<name>

To stop a run, deactivate its Kueue workload: kubectl -n olmo patch $(kubectl -n olmo get workloads -o name | grep "/jobset-<name>-") --type merge -p '{"spec":{"active":false}}'. Setting the JobSet's spec.suspend does nothing while Kueue has it admitted: Kueue resets it, and a watcher that relied on it left a run training on the 8x8x8 for 4.5 h after its stop rule fired.

Startup takes ~35 min: the 32 pods start a few minutes after admission, then the model compiles and the first step runs. k8s/watch_jobset.sh <jobset> <run> [max minutes] saves every pod's log when the JobSet ends. Pod logs vanish once the JobSet is deleted or restarted, so save them before either.

5. Read the loss

Read the step loss only from the LOSS-PARTS step=N ce=… moe_lb=… z=… tokens=… line, which one pod logs. Each pod's completed step: N, … loss: line is that pod's own value: on one step the 32 pods ranged 1.28–1.68. Compare ce= with H045's train/CE loss, never Tunix's loss, which adds the MoE load-balancing and z-losses.

A healthy Olmo 3.5 run: REQUIRED_RBF (the MoE buffer factor that would have dropped no token) stays below the configured ragged_buffer_factor=2.5. r2's maximum was 2.24.

REQUIRED_RBF is logged every step by JAX process 0 only, and that is often not pod slice-0-0: find the pod whose log has it. Kubelet keeps about 700 steps of each pod's log, and the agent service account cannot read Cloud Logging. So if a rule needs every step's value, stream that pod with kubectl logs -f from step 1. A watcher polling kubectl logs every 5 minutes lost 14 steps at one rotation (H078 bb4-lr1e4-s0), and the run had to be relaunched.

6. Checkpoints, preemption and resume

  • enable_checkpointing=true is required even if you want no checkpoints. Without it, jax.distributed is never initialised on multi-host TPU and Tunix crashes.
  • Tunix always saves the final step. An Olmo 3.5 checkpoint with optimizer state is ~137 GB. Agents delete nothing on GCS; ask the run's owner.
  • After a preemption, JobSet restarts the run from the last checkpoint, on the same overlay. A resume replays the data pipeline to the restored step at ~1.4 s per skipped batch, which is hours late in a run. Resume only on the restore-fix image (step 1).

7. Export to HF and evaluate

gcloud storage cp -r gs://ai2-olmo-dev/post-training/runs/<run>/checkpoints/<step>/model_params <local dir>
python scripts/tpu/maxtext/olmo35/export_maxtext_to_hf.py <local dir>/model_params <reference HF dir> <out>

The reference directory supplies tensor names, config.json, modeling code, tokenizer and chat template. Use an HF export of the same model with the matching tokenizer and chat template, and choose <out> under your own output location. One existing reference export is H045's: /weka/oe-training-default/ai2-llm/checkpoints/abhishekr/hero-sft-4t/runs/h045-anchor-emo-olmo35-s34521/hf_step11768. Its config.json is the HF model config, unlike the run directory's OLMo-core training config. The exporter is bit-exact on all 23,441 tensors. For Olmo 3, use export_hf.sh, which also takes the base model's config.json, because MaxText's own is wrong.

Evaluate the HF export with olmo-eval on Beaker GPUs, as for a GPU run (see the GPU guide). HF Olmo 3.5 needs fla, so every pass runs on GPU. Where exports and checkpoints go, and what to keep, is the ledger's storage skill (olmo-post-training-ledger skills/storage/SKILL.md).

Traps

Each of these trained silently on wrong data or a wrong model before it was found.

  • Dropped EOS targets. Olmo tokenizers have bos == eos, and MaxText's target shift ignored BOS targets, so SFT never learned to stop while the loss looked fine. Fixed in the image. Check P(EOS) after an answer on any new pipeline.
  • YaRN. Olmo 3 applies YaRN to full-attention layers only. Unpatched MaxText also applies it to sliding-window layers, which costs +0.1 nats/token at the start. Fixed in the image.
  • Vocab padding. 100,278 isn't divisible by tensor parallelism. Zero-padding to 100,352 is wrong (+3.2 nats/token). Replicate the vocab axis instead: logical_axis_rules=[['vocab',[]],['activation_vocab',[]]].
  • Olmo 3 tokenizer. allenai/olmo-3-tokenizer-instruct-dev loads as GPT2Tokenizer under transformers 5 and drops the pre-tokenizer. Train with a PreTrainedTokenizerFast save.
  • KDA miscompile on v4. Slicing a reshape inside the sub-block inverse dropped a column offset on TPU only (CPU was correct). Fixed in maxtext 9f55e7b6 (in the image). Kernel rewrites need a TPU-vs-exact-scan check, not only a CPU one.
  • MoE buffer. Fully dropless runs out of HBM at 64K. ragged_buffer_factor=2.5 drops nothing as long as REQUIRED_RBF stays below it.
  • Restore OOM. r2's image replicates the expert Adam moments on restore and runs out of HBM. Use the restore-fix image.
  • Data order depends on slice shape. Batches are packed per data host, so a different host count gives a different, still valid, data order. Keep every arm of a comparison, and its control, on the same slice shape.
  • Context parallelism with load balancing gives a wrong forward. With ici_context_parallelism > 1 and MaxText's default context_parallel_load_balance=True, dense Olmo 3 7B's step-1 CE on a fixed packed batch was off by +0.39 (cp4) and +1.98 (cp8) against an HF forward; cp4 with context_parallel_load_balance=False matches it (0.919794 vs 0.919799). The suspected cause, from the code, is that the SFT input path never applies the load-balancing reorder that attention then assumes; upstream confirmation is pending. Turn load balancing off, and check step-1 LOSS-PARTS CE against an HF forward of the same batch before trusting any new mesh. The recipe's olmo35-ep-as-cp uses the expert axis rather than ici_context_parallelism; on Olmo 3.5 small its step-1 CE matched an HF forward of the same 64 rows (1.077855 vs 1.077850).
  • Zero steps, exit 0. train_sft exits 0 with no steps when the data iterator fails. The launcher's step guard catches it; keep that guard.
  • KDA all-gathers on a new mesh. Under olmo35-cp-ep on the 8x8x8, the KDA chunk scan's inputs and states were sharded on the chunk axis, so XLA all-gathered them inside every scan iteration: 97.6% of a 475 s step. Logical constraints in _delta_rule_chunked_subblock (abhishekraok/maxtext c5f1251d) bring the step to 3.33 s. CPU runs never show in-loop gathers, so profile one TPU step with xprof before trusting a new mesh.
  • Compile-time HBM. At 64 rows on the 8x8x8, ragged_buffer_factor=4.0 failed to compile (HLO temporaries 32.17 GB vs 31.75 GB of HBM); 3.0 compiles, and the runs peaked at 2.39-2.84.
  • GCS egress quota. Two big slices loading the base checkpoint at once hit google_egress_bandwidth (us-central2) and failed at startup. Stagger big-slice launches.
  • The control host is a single point of failure. Watchers and log tailers die with the host they run on (a Beaker session lost its node on 2026-10-16). Keep any record a rule depends on in a source that survives it.
  • hostNetwork. The image's libtpu forms a multi-host slice only on the host network. Keep hostNetwork: true and dnsPolicy: ClusterFirstWithHostNet.
  • pgrep -f / pkill -f inside bash -c match their own command line: loops never end and kills hit the caller. Use PIDs.

Changing the recipe

  • Other slice sizes and meshes. olmo35-ep-as-cp puts one 64K row on 8 chips (KDA's 8 heads). Measured on Olmo 3.5 small with 65,536-token rows:
slice mesh rows/step s/step tokens/s/chip status
4x4x8 olmo35-ep-as-cp, fsdp 16 x expert 8 16 3.94 2,080 r2 and A0'; the recipe
8x8x8 olmo35-ep-as-cp, fsdp 64 x expert 8 64 4.28 1,910 H078's 4x batch; ragged_buffer_factor=3.0; 1,887 tokens/s/chip without the shim, see Cost
8x8x8 olmo35-cp-ep + c5f1251d 16 3.33 620 runs, but numerics differ from the 4x4x8 (#1933)
8x8x16 olmo35-cp-ep, fsdp 64 x context 2 x expert 8 64 - - never reached step 1 (egress quota, then a TPU core halt)

Each mesh was picked to keep r2's per-device layout, not tuned. No sweep of the parallelism split, remat policy or per-device batch has been run, so these are known-good points, not optima. Any other slice or mesh needs its own same-slice control. - Comparing against a GPU result. TPU and GPU runs of the same recipe differ in data order and numerics. Compare TPU arms with a TPU control on the same slice shape.