From f0c2758c78133c4d9ca4a87356df8c9a4d9c7afd Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Fri, 16 Oct 2026 03:53:11 -0700 Subject: [PATCH 1/8] docs(qwen3.6): add data-format ablations and move SciCode diagnosis to ABLATIONS.md Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .../tutorials/Qwen3.6-35B-A3B/ABLATIONS.md | 107 ++++++++++++++++++ .../tutorials/Qwen3.6-35B-A3B/README.md | 52 +-------- examples/megatron_bridge/tutorials/README.md | 2 +- 3 files changed, 112 insertions(+), 49 deletions(-) create mode 100644 examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/ABLATIONS.md diff --git a/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/ABLATIONS.md b/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/ABLATIONS.md new file mode 100644 index 00000000000..33789d9b34d --- /dev/null +++ b/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/ABLATIONS.md @@ -0,0 +1,107 @@ +# 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** | pending | +| MMMU-Pro (3 / 8)1 | 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.992 | +| SciCode capped sub-steps | 0.7% | 4.3% | 3.0% | 3.1% | 3.1% | 5.9% | **3.6%** | 5.7% | +| SciCode median tokens | 1,806 | 2,684 | 2,244 | 2,236 | 3,726 | 3,208 | **3,482** | 2,766 | + +>1 3 runs at iterations 50 and 300, 8 at 500. 2 Over the five finished benchmarks; the pre-tokenized run averages −0.28 on the same five. +> +> *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 is ahead on 5 of 6 benchmarks (mean −0.34 vs −1.30) and cuts the SciCode runaway rate from 4.3% to 3.0%. At 300 the two formats tie, and at 500 the pre-tokenized run is better (−0.28 vs −0.99 on the same five benchmarks, mostly tau2-bench and MMMU-Pro). 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. + +### 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 ``, `