Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 3 additions & 5 deletions examples/megatron_bridge/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,10 @@ This directory contains examples of using Model Optimizer with the [NeMo Megatro

</div>

This README is the getting-started guide, with each workflow shown on small models (mostly Qwen3-8B).

> [!TIP]
> Checkout the [Nemotron-3-Nano-30B-A3B pruning + distillation (with data blend prep) + quantization tutorial](tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md) for a complete end-to-end workflow using Megatron-Bridge!
>
> Or the [Qwen3.6-35B-A3B W4A4 NVFP4 + QAD tutorial](tutorials/Qwen3.6-35B-A3B/README.md) for an end-to-end quantization-aware distillation workflow.
>
> Or the [NemotronH omni VL W4A4 + QAD tutorial](tutorials/NemotronH-Omni-W4A4-QAD/README.md) for a Mamba-Transformer MoE vision-language model.
> For complete end-to-end workflows on larger models, see the [tutorials](tutorials/README.md): [Nemotron-3-Nano-30B-A3B](tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md) pruning + distillation + quantization, [Qwen3.6-35B-A3B](tutorials/Qwen3.6-35B-A3B/README.md) W4A4 NVFP4 + QAD, and [NemotronH omni VL](tutorials/NemotronH-Omni-W4A4-QAD/README.md) W4A4 + QAD.

## Pre-Requisites

Expand Down
115 changes: 115 additions & 0 deletions examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/ABLATIONS.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,115 @@
# Ablations: Qwen3.6-35B-A3B W4A4 QAD

Experiments behind the [main tutorial](README.md) that did not become the recommended recipe. All builds use the same W4A4 NVFP4 PTQ student, the same evaluation setup and the same `mean ± sem` pooling as the [main results](README.md#results); deltas are against BF16.

- [Data format: pre-tokenized blend vs chat SFT data](#data-format-pre-tokenized-blend-vs-chat-sft-data)
- [SciCode non-termination: diagnosis and decode-time fix](#scicode-non-termination-diagnosis-and-decode-time-fix)

## Data format: pre-tokenized blend vs chat SFT data

The main run trains on the **pre-tokenized blend**: conversations rendered with the chat template, packed into 32K windows, with the KD loss on **every** token, prompts included. Two alternatives feed the same Cascade-2 conversations as a chat `.jsonl` through `distill.py --sft --sft_hf_dataset`, with the loss on **assistant tokens only**:

- **Chat SFT:** the Cascade-2 responses as they are.
- **Self-distilled chat SFT:** each final assistant turn regenerated by the BF16 teacher (thinking on, `T=0.6`, `top_p=0.95`, responses that hit 32K dropped), so the student trains on the teacher's own response style and length.

Everything else matches the main run: 32 nodes, EP=8, ~16.8M tokens per iteration (one optimizer step), LR 1e-5 → 1e-6 over a 500-iteration schedule, top-4096 KD. Iterations 50, 300 and 500 are therefore ~0.8B, ~5.0B and ~8.4B training tokens. Both alternatives were first compared at iteration 50, where the main run's SciCode runaway rate peaks; chat SFT, the stronger of the two, was then continued to 500.

### Results

| Benchmark (runs) | BF16 | Pre-tok it50 | Chat SFT it50 | Self-distilled it50 | Pre-tok it300 | Chat SFT it300 | **Pre-tok it500** | Chat SFT it500 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| GPQA Diamond (1) | 84.7 | −0.8 | +0.5 | −0.3 | −0.4 | −0.4 | **−0.4** | 0.0 |
| MMMU-Pro (3 / 8)<sup>1</sup> | 74.6 | −0.8 | −0.5 | −0.9 | −1.0 | −0.7 | **−0.7** | −1.2 |
| SciCode (8) | 39.9 | +0.2 | +0.2 | −1.7 | −1.2 | −1.0 | **+0.3** | −0.4 |
| IFBench (8) | 60.0 | −1.3 | −1.9 | −0.0 | −1.9 | −2.6 | **−0.4** | −0.7 |
| AA-LCR (8) | 69.1 | −0.9 | +1.0 | −1.4 | −0.6 | +1.1 | **+0.3** | +0.1 |
| tau2-bench Telecom (3) | 94.2 | −4.2 | −1.4 | −2.9 | −1.2 | −2.2 | **−0.9** | −2.8 |
| **Mean delta** | | −1.30 | **−0.34** | −1.19 | −1.05 | −0.98 | **−0.30** | −0.83 |
| SciCode capped sub-steps | 0.7% | 4.3% | 3.0% | 3.1% | 3.1% | 5.9% | **3.6%** | 5.7% |
| SciCode median tokens | 1,817 | 2,684 | 2,244 | 2,236 | 3,726 | 3,208 | **3,482** | 2,766 |

><sup>1</sup> 3 runs at iterations 50 and 300, 8 at 500.
>
> *Capped* = hit the 131,072-token limit without a stop token, out of 2,704 SciCode sub-steps (338 × 8 runs). See [SciCode non-termination](#scicode-non-termination-diagnosis-and-decode-time-fix).

**Chat SFT's early lead does not last.** It leads at iteration 50, is nearly tied at 300 and trails at 500, mostly on tau2-bench and SciCode. Its SciCode capped rate *rises* with training, to ~5.8% against 3.1-3.6% for the pre-tokenized run, so assistant-only loss did not remove the runaway here (chat SFT also changed packing; see the caveat below).

**Self-distillation does not fix the length drift either.** At iteration 50 its SciCode capped rate and median length match chat SFT's even though it trains on the teacher's own responses, and it is no better on average, so it was not continued.

**Caveat: the chat-SFT teachers were noisier.** Both chat-SFT variants pack two conversations into one THD micro-batch (`cu_seqlens`), and on that path the Gated DeltaNet kernels are not reproducible: the BF16 teacher's logits for the same packed input differ run to run by 0.001–0.007 nats/token of KL on average, against a KD loss of ~0.03. The same conversations unpacked are bit-exact, and the pre-tokenized run feeds unpacked 32K windows, so only the chat-SFT arms trained on these noisier targets. Treat their gap to the pre-tokenized run as provisional until the kernel is fixed.

### Cost

32 nodes × 4 GB200 = 128 GPUs for every run.

| | Pre-tokenized | Chat SFT | Self-distilled chat SFT |
| --- | --- | --- | --- |
| Loss on | all tokens, prompts included | assistant tokens | assistant tokens |
| Data prep | tokenize: 5.8 h on one 96-CPU node | jsonl build, Bridge scan, row pairing (CPU only) | + teacher generation: ~43 GPU-h per 50 iterations of data |
| Tokens per micro-batch | 32,768 (packed window) | ~9K (two conversations) | ~10K (two conversations) |
| Micro-batches per rank per step | 4 | 14 | 12 |
| Steady step time | ~37 s | ~95-105 s | ~90 s |
| QAD GPU-h, 50 iterations | ~100 | ~290 (Triton autotune cache still filling) | ~194 |
| QAD GPU-h, 500 iterations | **~735** (5.7 h) | ~2,150 (16.8 h) | — |

Chat SFT costs **~2.9× the GPU-hours** for the same tokens. Each micro-batch holds two unpadded conversations, ~9K tokens on average, instead of a full 32K window, so a step runs 3.5× as many micro-batches, and the fixed per-micro-batch work (teacher and student forward, full recompute, MoE dispatch) dominates. Pairs cannot grow past ~32K tokens: the 248,320-token vocabulary's fp32 logits are ~30 GiB per 32K tokens.

### How the chat-SFT runs were configured

Launch details for the two chat-SFT runs above, not a separate ablation. They need `distill.py --sft_in_batch_packing` and `--moe_aux_loss_coeff`, which [#2873](https://github.com/NVIDIA/Model-Optimizer/pull/2873) adds. To repeat them:

- **Pack rows in the batch.** With EP > 1, Megatron-Bridge pads every SFT row to `--seq_length`, which made a step ~8× the pre-tokenized cost here. In-batch packing (micro-batch size ≥ 2, unpadded THD) avoids that. Pair long rows with short ones so each pack stays under ~32K tokens.
- **Disable the MoE auxiliary loss** (coefficient 0). It runs over padding tokens, and the linear-attention layers amplify the resulting gradient into the pad embedding row; after clipping, everything else gets near-zero updates.
- **Drop rows that Bridge's assistant-mask builder rejects** (~1.2% of Cascade-2 rows) and rows containing a literal `<image>`, `<video>` or `<audio>` string (HTML in code answers): a single such row switches the whole dataset to the VLM collator, which then fails for a text tokenizer.
- **Persist Triton autotuning.** The linear-attention kernels re-autotune for every new packed length. Set `TRITON_CACHE_AUTOTUNING=1` with a node-local `TRITON_CACHE_DIR` seeded from shared storage; 128 ranks writing into one Lustre cache directory stalled a run.

## SciCode non-termination: diagnosis and decode-time fix

QAD 500 raises SciCode's mean completion length by +90.9% over BF16 at an unchanged score ([Output length](README.md#output-length)). It is not verbosity. It is a **failure to terminate on a small fraction of sub-steps** — SciCode splits its 80 problems into 338 sub-steps, and `Subtask` scores the fraction of those whose code passes:

- Sub-steps that hit the 131,072-token cap inside `<think>` go from **0.7% (20/2704, BF16) to 3.6% (96/2704, QAD 500)**. Almost all return **zero answer tokens** (93 of those 96) — the model writes a complete solution, says *"I think I've been going in circles"*, and writes it again. In the case we inspected, a 20-word window repeats **352 times** and 97.7% of the trace's 20-word windows are duplicates.
- Those 3.6% of sub-steps burn **45.7% of all completion tokens**, so they dominate the mean: excluding them it is **+30.4%** rather than +90.9%.
- The **median** also roughly doubles (+91.6%), so the whole distribution shifted right — this is not *only* a tail effect.
- Capped rate peaks at **iteration 50** (4.3%) and settles at 3.1% / 3.6% by 300 / 500; it is not gradual drift.

The obvious suspect — that `--logit_kl_top_k 4096` leaves the stop tokens outside the loss — **did not hold up**. We scored 16 BF16 sub-steps that stop normally and 16 QAD-500 runaway sub-steps, with the eval's exact prompts, under the BF16 teacher and the PTQ, QAD-50 and QAD-500 students, over the full vocabulary:

- **The top-4096 tokens hold essentially all of the teacher's probability**: on average 1e-5 falls outside. `</think>` is often ranked outside the top-4096 mid-reasoning, but only with negligible probability, so the top-k loss sees nearly the whole distribution.
- **QAD does not suppress ending the reasoning.** Fed the teacher's sampled text, every student puts nearly all probability on `</think>` exactly where the teacher closes its reasoning. That is the token that matters here, since the runaways hit the cap inside `<think>`; the final `<|im_end|>` was not tested.
- **Inside the loops the teacher keeps going too.** Its stop probability there is about e<sup>-22</sup>, and it predicts the next repeated token with probability ~0.98.

So once a loop has started, a teacher-forced loss has nothing to correct: teacher and student agree. The difference is that the student enters loops in its own free-running generation, which KD on fixed text never exercises. Neither assistant-only loss nor training on the teacher's own responses removes it ([data format](#data-format-pre-tokenized-blend-vs-chat-sft-data)). The students were scored with their quantized weights dequantized to BF16 and BF16 activations, so the second finding does not measure W4A4 serving; the first and last depend only on the teacher.

### `presence_penalty: 1.5` removes it at decode time

Qwen's own thinking-mode recommendation for this model removes nearly all of it, with no retraining. Every cell is SciCode **without → with** the penalty, 8 runs per side, capped counts out of 2,704 sub-steps:

| Model | Capped sub-steps | SciCode (Subtask) | Mean tokens | Median tokens |
| --- | --- | --- | --- | --- |
| **BF16** (teacher) | 20 → **0** | 39.9 → 39.9 | 5,348 → 3,590 | 1,817 → 1,927 |
| W4A16 NVFP4 PTQ | 37 → **0** | 38.5 → 39.4 | 6,572 → 3,838 | 1,917 → 1,978 |
| **W4A4 NVFP4 PTQ** | 32 → **1** | 39.1 → 39.1 | 6,713 → 4,052 | 2,116 → 2,184 |
| **↳ + QAD 500 iters** | 96 → **2** | 40.2 → 39.6 | 10,209 → 4,641 | 3,482 → 3,302 |

Accuracy is neutral — the deltas span −0.59 to +0.89, run-to-run noise in both directions. The QAD row recovers 94 capped sub-steps, up to +3.5 pp of mechanical headroom, and the score does not move. Whether those sub-steps simply fail their tests anyway or the penalty costs a little elsewhere is not separable at this noise level — either way the win is cost, not accuracy: SciCode wall-clock generation drops 12-44%.

Means are `avg_completion_tokens`, so the left-hand column matches the main README's output-length table; medians come from the per-response `output.jsonl`.

**The penalty removes the tail; it does not make the model concise.** The median is flat or slightly *up* on three of four builds, so a typical response is unchanged; the mean falls 33-55% purely because the runaway generations are gone. QAD's median is still **+71%** over BF16 with the penalty on, against +91.6% without.

**Scope it to this failure mode.** On GPQA Diamond, which never caps, the same setting leaves accuracy unchanged and costs **4.8× throughput**. The penalty has a *rate* cost that scales with concurrency, and SciCode pays it back in tokens it no longer emits:

| Benchmark | `parallelism` | tok/s without → with | rate cost | tokens removed | net |
| --- | --- | --- | --- | --- | --- |
| SciCode | 8 | 1,475 → 1,192 | 1.24× | −55% | **faster** |
| GPQA Diamond | 32 | 5,109 → 1,067 | 4.8× | −7% | **slower** |

Measured on single-Slurm-segment runs on both sides, so neither figure includes requeue overhead. The likely cause is that the penalty is applied per sequence across the 248,320-token vocabulary on every decode step, so its cost grows with the number of concurrent sequences. We measured the correlation with `parallelism`, not the implementation, so treat the attribution as inference. The practical rule is the same either way: **the penalty pays for itself only where it removes enough tokens to cover its rate cost**, and that means workloads that actually cap.

[Lotfi et al., *Quantized Reasoning Models Think They Need to Think Longer, but They Do Not*](https://arxiv.org/abs/2606.00206) reports a related effect — PTQ amplifies "overthinking", where the model reaches a correct answer and then talks itself out of it — and proposes a logit penalty on hesitation markers ("Wait", "But", "Alternatively"). We tested their penalty here: it helps, but less than `presence_penalty` and at a small accuracy cost, and their marker-density signature does not reproduce on our W4A4 PTQ. Their failure opens *new* reasoning branches; ours repeats a finished solution.

### For the next QAD run

- Track the **length-capped rate** as a first-class metric alongside accuracy: a benchmark score can stay flat while 3.6% of responses return nothing.
- Changing the KD loss's vocabulary coverage (top-p, full-vocab KL) is unlikely to help: the top-4096 tokens already hold nearly all of the teacher's probability. We confirmed it with a 50-iteration rerun that adds a "ghost" token carrying the teacher's mass outside the top-4096: SciCode 39.35 ± 0.56 vs 40.13 ± 0.60, length-capped 5.0% vs 4.3%, all within noise.
- **On-policy distillation did not fix it either**, in one offline round. From the iteration-50 weights, 50 more KD iterations on 189K responses the student generated itself (looping ones kept) for the first 190K prompts of the chat-SFT data give SciCode 39.24 ± 0.47 with 5.2% length-capped, against 38.50 ± 0.74 and 5.5% for the same 50 iterations on the teacher's responses to the same prompts, and 5.0% at the start. The test is weak: on these training prompts the student loops barely more than the teacher (4.1% vs 3.4% capped), so its rollouts hold few loop entries; a retest would need SciCode-like prompts. `presence_penalty` remains the fix.
Loading
Loading