diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index a6428d79022..db46db9e43c 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -15,12 +15,10 @@ This directory contains examples of using Model Optimizer with the [NeMo Megatro +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 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..b36b1d32727 --- /dev/null +++ b/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, ~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)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.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 | + +>1 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 ``, `