docs(qwen3.6): add data-format ablations and move SciCode diagnosis to ABLATIONS.md

examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/ABLATIONS.md

@@ -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, ~16M tokens per step, LR 1e-5 → 1e-6 over a 500-iteration schedule, top-4096 KD. 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.** At iteration 50 it leads on 4 of 6 benchmarks, ties SciCode and trails only on IFBench (mean −0.34 vs −1.30); it also cuts the SciCode runaway rate from 4.3% to 3.0%. At 300 the two formats are nearly tied (mean −1.05 vs −0.98), and at 500 the pre-tokenized run is better (mean −0.30 vs −0.83, mostly tau2-bench and SciCode). Chat SFT keeps the median SciCode response shorter, but its capped rate *rises* with training, to ~5.8% against the pre-tokenized run's 3.1-3.6%. Masking the loss to assistant tokens is therefore not what drives the runaway.
+
+**Self-distillation does not fix the length drift either.** At iteration 50 its SciCode capped rate and median length match chat SFT's, though its training targets are the teacher's own responses; SciCode scores 1.7 points below BF16. It recovers IFBench fully (the only iteration-50 build to do so) but is no better on average and costs extra generation, 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.
+
+### Running QAD on chat SFT data
+
+These runs needed two `distill.py` options that it does not expose, so they are not reproducible from this repository as is; they were not upstreamed because the format did not win. For anyone repeating the comparison:
+
+- **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 stopping.** Fed the teacher's text, every student puts nearly all probability on `</think>` exactly where the teacher stops.
+- **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; the first and last findings 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 generated by the student itself (looping ones kept) 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.