jburtoft

Voxtral-Mini-3B-2507-draft-4layer

README

License: apache-2.0

What it is

Architecture: identical to Voxtral-Mini-3B-2507 in every component except the LLM decoder is pruned from 30 layers to 4 layers (kept: indices [0, 10, 20, 29]).
Shared verbatim with target (frozen, not trained): audio encoder (32 layers, Whisper-derived), multi_modal_projector, embed_tokens, final norm, lm_head. Same tokenizer, same audio path, same output vocab.
Trained: only the 4 kept Llama decoder layers (~428M trainable params of the ~3B total).
Compute: 4-layer decoder is ~13% of the target's decoder FLOPs. Practical latency ratio measured on GPU L40S: draft is ~4.4x faster per generated token than the target.

Results

Reference measurements on a held-out corpus of 18 short-to-medium English audio clips (0.4-27 seconds, TED-derived speech), never seen during training:

Mean top-1 argmax agreement with target: 66.67%. Equivalent to the draft's per-position "would this token match the target's greedy pick?" rate under teacher forcing.
Analytical GPU speedup at K=2 candidates/iteration: 1.416x mean (1.381x on clips >= 15 s).
Byte-identical greedy correctness: yes, by construction of the greedy speculative-decoding algorithm.

Per-K analytical speedup on L40S (target: 24.2 ms/tok; draft: 5.5 ms/tok):

Table with columns: K, mean accepted per iter, mean speedup, speedup (long clips >= 15s)
K	mean accepted per iter	mean speedup	speedup (long clips >= 15s)
2	1.13	1.416x	1.381x
3	1.37	1.364x	1.351x
4	1.57	1.300x	1.280x
6	1.80	1.148x	1.122x

K=2 is optimal on GPU L40S given the ~4.4x cost ratio; larger K adds draft cost that isn't repaid by the incremental acceptance gain. On different hardware (e.g., AWS Neuron/Trainium with a fused-K draft NEFF), the optimal K will differ.

Intended use

python
from transformers import VoxtralForConditionalGeneration, AutoProcessor
import torch

target = VoxtralForConditionalGeneration.from_pretrained(
    "mistralai/Voxtral-Mini-3B-2507",
    dtype=torch.bfloat16,
).cuda().eval()
draft = VoxtralForConditionalGeneration.from_pretrained(
    "jburtoft/Voxtral-Mini-3B-2507-draft-4layer",
    dtype=torch.bfloat16,
).cuda().eval()
processor = AutoProcessor.from_pretrained("mistralai/Voxtral-Mini-3B-2507")

inputs = processor.apply_transcription_request(
    language="en",
    audio="my_clip.wav",
    model_id="mistralai/Voxtral-Mini-3B-2507",
).to("cuda", dtype=torch.bfloat16)

# Standard HuggingFace speculative decoding
out = target.generate(
    **inputs,
    max_new_tokens=256,
    do_sample=False,
    assistant_model=draft,
)
transcript = processor.tokenizer.decode(
    out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True,
)

Known caveat: as of transformers==5.14.1, HF's built-in assistant_model= path exhibits inefficiencies with VoxtralForConditionalGeneration (multimodal target). The final output is byte-identical to greedy baseline, but per-iteration acceptance-rate reporting is incorrect and wall-clock speedup is often below the theoretical ceiling. The training-time analytical measurement (66.67% top-1 agreement) reflects the model's true quality; a correctly-implemented spec-decode wrapper should approach the K=2 -> 1.4x speedup shown above. See measurement/README.md in the companion training scripts for a working analytical measurement approach.

The intended production deployment for this draft is on AWS Neuron/Trainium via a custom NxD Inference spec-decode wrapper (planned separately). GPU is useful for validation but is not the target platform.

Standalone quality (illustrative)

The draft is not intended for standalone transcription. If run alone on a 9.67 s clip about nuclear reactor refueling, it emits:

"That's right, you're always refiling the revolution. So I have lots of people and lots of conditions that can go wrong. That thing where you'm opening things in and moving things in and move."

The target on the same clip emits:

"That's right. Today, you're always refueling the reactor. So you have lots of people and lots of controls that can go wrong. That thing where you're opening it up and moving things in and out, that's..."

The draft's output is topically related and grammatically English but has word substitutions ("refueling the reactor" -> "refiling the revolution", "controls" -> "conditions"). This is a fundamental property of a highly-shrunk distilled model. When used inside a spec-decode loop, the target's verification pass corrects these mispredictions on the fly, so the final output is byte-identical to the target's greedy output.

Training

Dataset

LibriSpeech train subsets (dev-clean + test-clean + test-other, public, CC BY 4.0). Sampled 2000 clips, mean duration 6.95 s, total ~230 minutes of English audio.
Pseudo-labeling: for each clip, ran Voxtral-Mini-3B-2507 end-to-end greedy on CPU to obtain the "teacher transcript" token IDs. These are the distillation targets (not the ground-truth LibriSpeech transcripts). Rationale: we want the draft to match the target's idiosyncrasies (its capitalization, punctuation, word choice), not to be a good ASR model on its own.

Distillation recipe

Framework: torch-neuronx custom training loop on AWS Trainium (trn2.3xlarge, LNC=2, TP=1, ~5-day capacity block; actual training used ~20 minutes).
Loss: cross-entropy on teacher's argmax tokens (teacher forcing). No KL on soft logits, no auxiliary losses.
Optimizer: AdamW, betas (0.9, 0.95), weight decay 0.01.
Learning rate: peak 1e-4, linear warmup 100 steps, cosine decay to 10% of peak over remaining 2900 steps.
Batch size: 2. Max sequence length: 512 (padded).
Precision: bfloat16 model, fp32 loss computation.
Gradient clipping: max norm 1.0.
Total steps: 3000. Approximately 3 epochs over 2000 clips at batch=2.
Frozen: all params outside the 4 kept decoder layers.

Convergence

Held-out validation top-1 agreement (measured every 200 steps on a validation set of unrelated short English audio clips):

Table with columns: Step, Val top-1 agreement
Step	Val top-1 agreement
200	12.82%
400	19.23%
600	25.05%
800	31.36%
1000	36.59%
1200	41.62%
1400	49.01%
1600	51.78%
1800	53.75%

Monotonically increasing until ~step 2400; plateau at 62-67% for the last 600 steps. The training data appears to be the bottleneck at this scale; scaling to ~10k+ LibriSpeech clips should push val top-1 into the 75-85% range and improve downstream spec-decode speedup accordingly.

Trainium-specific workaround (XLA path only)

What breaks: On the torch-neuronx XLA path (the default), Voxtral's masked_scatter call in VoxtralModel.forward -- which splices the audio-encoder output into inputs_embeds at the audio-token positions -- lowers to an HLO sort op on a s32[1176576] index tensor (for the standard [1, 383, 3072] inputs_embeds shape). Neuron's XLA backend rejects this with [NCC_EVRF029] Operation sort is not supported on trn2. The compiler's suggestion is to use TopK or an NKI kernel instead, but no such lowering exists in torch_xla for masked_scatter at this size as of neuronx-cc 2.26.6360.0+6f180f47.

What is fine: The audio encoder (audio_tower + multi_modal_projector) compiles cleanly on Neuron via both the XLA and Beta 4 eager paths. This is not a hardware or model-architecture limitation.

Workaround (used by the training scripts 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. Since the audio path is frozen (never trained), the audio-encoder forward pass only needs to run once per clip anyway, so moving it to CPU costs nothing in the amortized case.

What PyTorch Native Beta 4 does differently: In eager mode Beta 4 dispatches masked_scatter directly to ATen kernels on the Neuron device, no XLA compilation involved. It runs natively and Voxtral works end-to-end on Neuron without any CPU workaround. Verified 2026-07-31: full model(**inputs) including audio encoder + projector + splice + 30-layer decoder runs in one call on torch.device("neuron:0"), ~1150 ms steady state. The same failure mode does NOT reproduce on Beta 4.

For training on either stack, we still use the CPU-precompute workaround: it's more efficient than re-running the frozen audio encoder every training step, and it makes the training loop's static-shape graph simpler (no variable-length audio inside the compiled region). See training_scripts/README.md for details.

For inference on Beta 4, no workaround is needed -- move the whole model to neuron:0 and call .generate(**inputs) normally.

Reproducibility

Training scripts and full recipe: see the training_scripts/ directory alongside this model card. Also linked from the file browser above.

Quick reproduction on a fresh trn2.3xlarge (SDK 2.31 DLAMI 20260708 or later):

bash
# --- 0. Setup
source /opt/aws_neuronx_venv_pytorch_2_9/bin/activate
pip install -q -U "transformers>=4.54" "mistral-common[audio]>=1.8.1" librosa soundfile

# --- 1. Sanity check that the LLM decoder loads on trn2
python training_scripts/sanity_voxtral_xla_v2.py

# --- 2. Download LibriSpeech dev-clean (~340 MB) + test-clean + test-other (~700 MB)
cd /mnt/data
for split in dev-clean test-clean test-other; do
  curl -sSL -o ${split}.tar.gz https://www.openslr.org/resources/12/${split}.tar.gz
  tar xzf ${split}.tar.gz
done

# --- 3. Build a training manifest (500 dev-clean clips)
cd - && cd training_scripts
python build_librispeech_manifest.py \
  --librispeech-root /mnt/data/LibriSpeech/dev-clean \
  --output /mnt/data/manifest_train.csv --limit 500

# Optional: combine multiple splits for more data
python build_librispeech_combined_manifest.py \
  --librispeech-roots /mnt/data/LibriSpeech/test-clean /mnt/data/LibriSpeech/test-other \
  --output /mnt/data/manifest_extra.csv --limit 1500 \
  --exclude-existing-dir /mnt/data/pseudo_labels_train

# --- 4. Precompute pseudo-labels (~20 min per 500 clips on CPU)
python precompute_pseudo_labels.py \
  --manifest /mnt/data/manifest_train.csv --dataset-root / \
  --output-dir /mnt/data/pseudo_labels_train

# --- 5. Precompute a small validation set of your own audio for validation
#     (any short English audio clips work; ~10-20 clips is sufficient)
python precompute_pseudo_labels.py \
  --manifest /path/to/your/val_manifest.csv --dataset-root /path/to/val \
  --output-dir /mnt/data/pseudo_labels_val

# --- 6. Train (~16 minutes on trn2.3xlarge for 3000 steps at batch=2)
python train_distill_v2.py \
  --train-data-dir /mnt/data/pseudo_labels_train \
  --val-data-dir /mnt/data/pseudo_labels_val \
  --output-dir /mnt/drafts/my-voxtral-draft \
  --keep-layers 0,10,20,29 \
  --max-seq-len 512 --batch-size 2 --lr 1e-4 \
  --warmup-steps 100 --total-steps 3000 \
  --save-every 500 --validate-every 200

# --- 7. Validate on one clip
python validate_draft.py \
  --draft /mnt/drafts/my-voxtral-draft/best \
  --audio /mnt/data/LibriSpeech/dev-clean/2035/147960/2035-147960-0015.flac

Expected outcome: val top-1 climbs from ~4% (step 200) to ~65-68% (step 2800), and standalone draft output on any English clip is coherent (not gibberish).

Files in this repo

Table with columns: File, Purpose
File	Purpose
`model.safetensors`	4-layer draft weights (~3.6 GB in bfloat16)
`config.json`	Voxtral config with `text_config.num_hidden_layers=4` and pruned `layer_idx` values on the 4 kept layers
`generation_config.json`	Copied from target
`processor_config.json`, `tekken.json`	Copied from target (same tokenizer + processor)

Limitations

Not a standalone ASR model. See "Standalone quality" above.
Small training corpus. Only 2000 LibriSpeech clips (~230 minutes English audio). Domain shift for non-English or non-conversational audio (medical, financial, etc.) will lower acceptance rate. Retraining on domain-matched pseudo-labels is straightforward and recommended.
English only. Trained on LibriSpeech English. Untested on other languages. Voxtral supports multilingual audio; the draft's other-language behavior is undefined.
HuggingFace assistant_model= inefficiency on GPU. Documented above; use a custom spec-decode loop or wait for Neuron NxD Inference integration.
4-layer choice was not swept. Alternative layer selections (5-6 layers, different indices) may give better accepance/cost tradeoffs. Reproducibility recipe supports easy re-training with --keep-layers.

Citation

If you use this draft, please cite the base Voxtral model:

bibtex
@misc{voxtral2025,
  title={Voxtral},
  author={Mistral AI},
  year={2025},
  url={https://huggingface.co/mistralai/Voxtral-Mini-3B-2507}
}

And LibriSpeech for the training corpus:

bibtex
@inproceedings{panayotov2015librispeech,
  title={Librispeech: an ASR corpus based on public domain audio books},
  author={Panayotov, Vassil and Chen, Guoguo and Povey, Daniel and Khudanpur, Sanjeev},
  booktitle={2015 IEEE international conference on acoustics, speech and signal processing (ICASSP)},
  pages={5206--5210},
  year={2015},
  organization={IEEE}
}

License

Apache 2.0. This model is a distilled variant of mistralai/Voxtral-Mini-3B-2507 (Apache 2.0), retaining the base model's tokenizer, config schema, audio encoder weights, projector weights, embed_tokens, lm_head, and final norm. Only the 4 kept decoder layers were trained by us on LibriSpeech.

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Model Details

Model Provider