training_scripts/README.md

Training scripts

Full reproduction pipeline for training the 4-layer distilled Voxtral draft in this repository. See the main model card for the end-to-end recipe. This document describes each script individually.

All scripts are self-contained and driven by command-line arguments. Defaults assume /mnt/data/ for working data and /mnt/drafts/ for output checkpoints, but every path is overridable via a flag.

Two training paths

The scripts support two Trainium software stacks:

1. torch-neuronx XLA (SDK 2.31 DLAMI) -- the default. Uses torch_xla with xm.mark_step(). Runs in the standard SDK 2.31 DLAMI's pytorch_2_9 venv. Faster steady-state per-step, produced the weights in this repo.
2. PyTorch Native Beta 4 (SageMaker Training DLC) -- eager mode + torch.compile(backend='neuron'). Runs in a separate venv created by beta35-eval's deploy-beta35.sh --variant beta4 script. Same convergence, slightly slower per-step, simpler port from vanilla PyTorch training code, and the direction AWS is investing in for future Trainium training.

Both paths produce the same trained weights (within run-to-run seed variance). The _beta4 variants of the sanity and training scripts are the Beta 4 ports.

Pipeline diagram

LibriSpeech .flac -----> precompute_pseudo_labels.py ---->  {clip}.pt per clip
                                          |                   (input_ids,
   audio_encoder (CPU) + Voxtral greedy --+                    audio_embeds,
                                                               teacher_token_ids)
                                                                      |
Trainium training loop:                                               v
    for step in range(total_steps):                                   |
        batch = load({clip}.pt files, pad to max_seq_len)  <----------+
        inputs_embeds = splice audio_embeds into text embeds
        logits = student(inputs_embeds)
        loss = CE(logits, teacher_token_ids)  # teacher forcing
        backward + step
                                                                      |
                                                                      v
                                       /mnt/drafts/best  <---- best checkpoint by
                                                              val top-1 agreement
                                                              on held-out audio

Scripts

Script	Purpose	Stack
sanity_voxtral_xla.py	Verify Voxtral loads on trn2 via torch_xla. Expected to fail on the audio path (documents the trn2 sort op limitation).	XLA
sanity_voxtral_xla_v2.py	Verified working XLA path: audio encoder on CPU, LLM decoder on trn2 via torch_xla.	XLA
sanity_voxtral_beta4.py	Beta 4 equivalent of sanity_voxtral_xla_v2.py. Uses torch.device("neuron:0") eager mode + torch.compile(backend='neuron'). Prints wall time for both eager and compiled variants.	Beta 4
build_draft.py	Take a Voxtral checkpoint and prune to N decoder layers. Saves to --out-dir. Used at training-start to create the initial student.	Any
build_librispeech_manifest.py	Build a CSV manifest from a single LibriSpeech split (dev-clean etc.).	Any
build_librispeech_combined_manifest.py	Build a manifest combining multiple LibriSpeech splits, with de-duplication against a previously-precomputed pseudo-label directory.	Any
precompute_pseudo_labels.py	For each audio clip in a manifest: run Voxtral greedy on CPU, save (input_ids, audio_embeds, teacher_token_ids) as <stem>.pt.	Any (CPU only)
train_distill.py (v1)	Simple training loop, no validation. Kept for reference; use v2 instead.	XLA
train_distill_v2.py (v2)	Training with periodic validation on a held-out set + best-checkpoint tracking. This is the script that produced this repo's weights.	XLA
train_distill_v2_beta4.py	Beta 4 port of v2. Identical arguments + adds --compile flag to wrap the LLM decoder in torch.compile(backend='neuron').	Beta 4
validate_draft.py	Standalone quality check for a trained draft. Runs greedy generation with the draft alone, then computes teacher-forced per-position top-1 agreement between draft and target.	Any (CPU)

Key design choices

Layer selection: [0, 10, 20, 29]

• Includes layer 0 (initial embedding processing) and layer 29 (final projection) for stability.
• Even spacing in between (10, 20) samples different depths of the target's representation.
• 4 layers is ~13% of the target's decoder FLOPs (4/30).
• Not systematically swept. Alternatives like [0, 7, 15, 22, 29] (5 layers) may give a better acceptance/cost tradeoff and are trivial to try by changing --keep-layers.

Pseudo-labels via greedy teacher

• Runs Voxtral end-to-end on each audio clip with greedy decoding.
• Saves target's transcript token IDs; student learns to reproduce them exactly under teacher forcing (cross-entropy loss).
• Alternatives considered:
  • KL on soft logits: potentially better quality but requires storing full 131072-vocab distributions per position (~256 MB per 128-token clip, 500+ GB total). Not run.
  • Ground-truth transcripts (LibriSpeech gold): risks the student learning phrasings different from Voxtral's, hurting downstream spec-decode acceptance.
  • Chose hard-label distillation (teacher's argmax) as the simplest reasonable option.

Frozen audio path

• Audio encoder + projector + embed_tokens + final norm + lm_head are byte-identical between student and target -- no reason to train them.
• Only the 4 decoder layers get gradients (~428M trainable of the ~3B target total).

Batch size 2

• Chose conservatively for a first run. Larger batches (4, 8) are likely feasible and would give more stable gradients. Untested; open for future work.
• Batch size 1 also works but shows higher loss variance.

Cosine LR schedule

• Standard for short distillation runs.
• 100-step linear warmup + cosine decay to 10% of peak LR over remaining steps.
• Peak LR: 1e-4. Reasonable for LLM fine-tuning; not swept.

Trainium-specific workaround (XLA path only)

Root cause -- XLA path only. On the torch-neuronx XLA path, Voxtral's end-to-end forward compile fails with [NCC_EVRF029] Operation sort is not supported on trn2. The op being sorted is s32[1176576] -- a 1D int32 index tensor of size 1 * 383 * 3072 -- which is the flattened form of Voxtral's inputs_embeds [1, 383, 3072]. The sort comes from torch_xla's lowering of masked_scatter in VoxtralModel.forward, specifically the splicing of audio embeddings into inputs_embeds at the audio-token positions. It is not the audio encoder itself. The audio encoder (audio_tower + multi_modal_projector) compiles cleanly on Neuron in both XLA and Beta 4 eager modes.

Verified 2026-07-31:

• XLA model.model.audio_tower(input_features) -> compiles + returns [1, 1500, 1280] correctly.
• XLA inputs_embeds.masked_scatter(mask, audio_embeds) on [1, 383, 3072] -> fails with sort error.
• Beta 4 eager inputs_embeds.masked_scatter(...) on the same tensor -> works.
• Beta 4 eager model(**inputs) end-to-end on neuron:0 -> works, ~1150 ms per forward.

What the workaround does (used by every script here): precompute audio embeddings on CPU once per training clip, splice them into inputs_embeds also on CPU in the DataLoader collate function, and only send the audio-conditioned inputs_embeds to Neuron. The audio encoder + projector never runs on Neuron in the training loop.

1. precompute_pseudo_labels.py runs model.model.get_audio_features(input_features) on CPU. This produces audio_embeds of shape [375, 3072] per clip.
2. train_distill_v2.py (and its Beta 4 variant) splices audio_embeds into inputs_embeds at the audio-token positions in the collate function (on CPU, per batch).
3. The Trainium graph only sees inputs_embeds (already audio-conditioned), not input_features. No audio encoder involved.

Since the audio path is frozen (never trained), this precomputation only runs once per training clip -- it's cheaper than re-running the frozen encoder every training step regardless of whether it would compile.

Beta 4 note: you could rewrite train_distill_v2_beta4.py to compute audio_embeds on-device inside the training loop and skip the precompute step. It works (see the sanity test), but it's slower per step (1150 ms including audio encoder) than the precompute approach (200-590 ms per step, audio-encoder cost paid once offline). The precompute approach is retained in train_distill_v2_beta4.py for parity with the XLA script and for training efficiency.

Observations from the reference run (XLA path, 3000 steps)

Final training summary (2000 LibriSpeech clips, batch=2):

val step  200: top-1 agreement = 12.82%
val step  400: top-1 agreement = 19.23%
val step  600: top-1 agreement = 25.05%
val step  800: top-1 agreement = 31.36%
val step 1000: top-1 agreement = 36.59%
val step 1200: top-1 agreement = 41.62%
val step 1400: top-1 agreement = 49.01%
val step 1600: top-1 agreement = 51.78%
val step 1800: top-1 agreement = 53.75%
val step 2000: top-1 agreement = 57.40%
val step 2200: top-1 agreement = 62.43%
val step 2400: top-1 agreement = 64.40%
val step 2600: top-1 agreement = 61.93%
val step 2800: top-1 agreement = 66.67%  <-- best
val step 3000: top-1 agreement = 65.78%

• Curve was still climbing at step 2400; plateaued at 62-67% for the last 600 steps. Training data appears to be the bottleneck at this scale.
• Data-scaling law observed in earlier runs: 4x more training data (500 -> 2000 clips) produced 1.7x higher final val top-1 (39% -> 67%). Scaling to 10k+ LibriSpeech clips (e.g., train-clean-100 = ~28k clips) should push val top-1 into the 75-85% range.
• Compile bursts happen periodically as new sequence shapes are encountered by the XLA HLO compiler. After the first ~60 s compile the training loop runs at ~130 ms/step steady-state. Save-checkpoint triggers a flush every N steps but is quick.

Beta 4 comparison (500 steps, --compile)

Same hyperparameters, same data, same seed. 500-step run:

Step	XLA v2	Beta 4 compiled
100	8.28%	6.61%
200	9.17%	9.76%
300	10.06%	10.16%
400	14.40%	15.88%
500	18.64%	18.34%

Convergence is essentially identical. The two paths produce interchangeable weights.

Per-step wall clock (steady state, batch=2, max_seq_len=512):

Path	Compile time	Steady-state step
torch-neuronx XLA	~60 s	~130 ms
Beta 4 --compile	~49 s	~200-590 ms (some recompiles on val batches)
Beta 4 eager (no --compile)	~107 s	~3700 ms (28x slower)

XLA is currently the fastest path for this scale. Beta 4 with --compile is 1.5-2x slower total wall time. Beta 4 eager is not competitive.

Recommendation: use train_distill_v2.py (XLA) unless your environment mandates the Beta 4 stack (some SageMaker training clusters, TorchTitan integration, or if you want to try torch.compile fusion on more advanced setups). The Beta 4 script is a drop-in replacement with the same arguments plus --compile.

Known gotchas

1. First iteration compile is slow (both paths): 60-150 s per new tensor shape encountered. Subsequent iterations hit the cache. Keep max_seq_len constant across all batches for best cache hit rate.
2. Checkpoint save is expensive: ~30-40 s per save (writes 3.6 GB safetensors + re-loads the full target for state-dict merging via load_state_dict(..., strict=False)). Use --save-every 500 or larger to minimize interruptions.
3. CPU pseudo-labeling is slow but embarrassingly parallel: ~5-10 s per clip on a single CPU. 2000 clips takes ~2 hours on one instance. If precompute is the bottleneck, run it on a large-core CPU instance (c7i.16xlarge is much faster than trn2's ~48 CPUs).
4. Voxtral's audio path is byte-identical between student and target, so the .pt files produced by precompute_pseudo_labels.py are reusable across draft architectures. If you train a 6-layer variant later, the precompute step doesn't need to re-run.
5. Beta 4-specific: the first log-print says step time is very high (e.g., "step 1 | step time 24768ms"). That's the cumulative-average step time carrying the compile cost. Look at the deltas between successive elapsed timestamps to see the true per-step wall time.
6. Beta 4-specific: torch.compile(dynamic=False) locks shapes. If your batch shapes vary (e.g., variable-length audio), you'll trigger recompiles and lose the fusion advantage. Pad every batch to a single fixed shape.