From 1eb8ce64db3a299c338aeaaf2270ba651c3d6445 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 23 Jul 2026 12:29:47 +0000 Subject: [PATCH 1/7] Add ModelOpt QAD skill Signed-off-by: Meng Xin --- .agents/skills/qad/SKILL.md | 160 ++++++++++++++++++ .../qad/assets/nemotron-cascade-2-blend.yaml | 44 +++++ .claude/skills/qad | 1 + examples/megatron_bridge/distill.py | 14 ++ tests/examples/megatron_bridge/test_qad.py | 10 +- 5 files changed, 226 insertions(+), 3 deletions(-) create mode 100644 .agents/skills/qad/SKILL.md create mode 100644 .agents/skills/qad/assets/nemotron-cascade-2-blend.yaml create mode 120000 .claude/skills/qad diff --git a/.agents/skills/qad/SKILL.md b/.agents/skills/qad/SKILL.md new file mode 100644 index 00000000000..be8e57a88de --- /dev/null +++ b/.agents/skills/qad/SKILL.md @@ -0,0 +1,160 @@ +--- +name: qad +description: >- + Recover a measured BF16-to-PTQ accuracy gap on Slurm with ModelOpt + Quantization-Aware Distillation (QAD) through Megatron Bridge. Use when PTQ + has already been benchmarked against the matching BF16 model and the + quantized checkpoint needs short, evidence-driven distillation; also use for + QAD topology, data preparation, Slurm launch, resume, checkpoint export, or + recovery decisions. +--- + +# ModelOpt Quantization-Aware Distillation + +Use QAD only to recover a material, apples-to-apples PTQ benchmark gap. Do not +train merely because a quantized checkpoint exists. + +## 1. Read the supported workflow + +Read these sources before constructing commands: + +- `examples/megatron_bridge/README.md`, especially Post-Training Quantization, + Data Preparation, Quantization Aware Distillation, export, and Slurm usage +- `examples/dataset/MEGATRON_DATA_PREP.md`, especially token-budgeted blends +- `examples/megatron_bridge/{quantize.py,distill.py}` via `--help` + +Treat them as the source of truth for mutable flags, containers, model support, +and checkpoint formats. Do not copy their full command lines into new wrappers. +The deprecated `examples/llm_qad` flow is not the default. + +## 2. Establish the gap + +Use the `evaluation` and `compare-results` skills to run and validate the same +benchmark configuration on the BF16 source and the export of the exact +quantized Megatron checkpoint that will initialize QAD. If the existing PTQ +result came from another checkpoint or its lineage is unclear, export and +benchmark the intended Megatron checkpoint first. Record its path or identifier, +metric, direction, uncertainty, and acceptable delta. Stop if the gap is not +material or the runs are not comparable. + +## 3. Set up Slurm execution + +Read `skills/common/environment-setup.md` and detect whether the Slurm target is +local or remote. This version of the skill supports Slurm only; stop if the +target is not Slurm. + +- On local Slurm, follow `skills/common/workspace-management.md` and + `skills/common/slurm-setup.md` directly. Do not add SSH or sync steps. +- On remote Slurm, also follow `skills/common/remote-execution.md`: establish + its persistent SSH session, create matching local and remote session/model + workspaces, and sync the source plus the two-script Slurm wrapper/inner runner. + +Use the common account, partition, registry-auth, submission, and monitoring +procedures in either case. + +Keep BF16, the exact benchmarked PTQ Megatron checkpoint and its export, QAD +Megatron, exported checkpoints, data, logs, and benchmark results in the same +session/model workspace. QAD must start from the checkpoint produced by +`examples/megatron_bridge/quantize.py` because it carries the ModelOpt state; an +exported HF PTQ checkpoint is not a substitute. + +## 4. Choose topology explicitly + +Inspect model parameter count and architecture, sequence length, GPU count/type +and memory, then choose the smallest topology that fits both student and teacher: + +1. Choose node count from model-state plus activation memory; do not start with + an oversized multi-node topology. +2. Use TP only when model-state/GEMM size warrants it, and preserve attention + head and hidden-size divisibility. +3. Use CP to distribute the 32K context when activation/attention memory is the + constraint. Prefer CP before TP for a small dense model at long context. +4. Keep EP=1 for dense models. For MoE, choose EP from expert count, memory, and + available ranks; verify expert and data-parallel divisibility. +5. Add PP only when layer/model state still does not fit. + +Verify `DP = world_size / (TP * PP * CP)` is integral and +`GBS % (MBS * DP) == 0`. Record nodes, GPUs/node, TP/PP/CP/EP/DP, MBS, GPU +memory, and why each non-one dimension is needed. + +Example only: for the requested one-node, eight-H100 Qwen3-0.6B validation at +32K, TP=1, PP=1, CP=4, EP=1, DP=2, MBS=1 is a reasonable starting point. This +is validation evidence, not a default topology; derive a fresh topology for +every target model and cluster. + +## 5. Prepare only the next data tranche + +Copy `assets/nemotron-cascade-2-blend.yaml`, then set its tokenizer, output +directory, and token budget. Run: + +```bash +python -m modelopt.torch.utils.plugins.prepare_megatron_data_blend \ + --config +``` + +The utility streams `nvidia/Nemotron-Cascade-2-SFT-Data`, applies a deterministic +approximate buffer shuffle, and stops at the token budget. Never snapshot or +materialize the full dataset. For the initial step-150 gate, the default +2.6-billion-token budget covers +`150 * 512 * 32768` tokens plus a small margin. Use a much smaller explicit +budget for workflow validation. At very small budgets, use a representative +subset of source configs rather than giving a rare config too few documents +for Megatron's train/validation split. + +Pass the generated `data_blend.txt` entries as `distill.py --data_paths`. +The supported real-data `GPTDatasetConfig` shuffles documents and concatenates +them into fixed-length 32K samples, so sequences are packed instead of padded. +Do not use mock data as training evidence. + +## 6. Retain PTQ lineage, then run staged QAD + +Use the exact PTQ Megatron checkpoint established in step 2, or produce and +benchmark it before QAD. Follow the current Megatron Bridge README for PTQ, QAD, +resume, and export commands. Apply these QAD defaults: + +| Setting | Default | +| --- | --- | +| Sequence length | 32768 | +| Peak / minimum LR | `1e-5` / `1e-6` | +| LR schedule | cosine | +| Training cap | 1000 iterations | +| Global batch size | 512 | +| Dataset | `nvidia/Nemotron-Cascade-2-SFT-Data` | +| Initial eval/save/exit | iteration 150 | +| Slurm duration exit | 220 minutes for a 4-hour allocation | + +Keep `train_iters=1000` from the first run so the cosine schedule has a stable +horizon. Set `eval_interval=150`, `exit_interval=150`, and +`exit_duration_in_mins=220` for the initial stage. The duration exit is a +checkpointing safety margin, not the training target. + +Use one result-bearing Slurm job per stage and put Pyxis container flags on the +final `srun`. Fold startup validation into that job; do not submit separate GPU +preflight or smoke jobs. + +## 7. Benchmark before continuing + +At the initial exit: + +1. Confirm finite QAD/KD loss, sensible learning rate, non-pathological gradient + norm, and an iteration-150 checkpoint. +2. Export that QAD checkpoint with the current quantized Megatron exporter. +3. Run the identical benchmark on BF16, the step-2 PTQ export, and QAD-150. +4. Calculate the original gap and recovery. For higher-is-better metrics: + `gap = BF16 - PTQ`, `recovered = QAD - PTQ`, and + `recovery_fraction = recovered / gap`. + +Continue only if benchmark recovery is positive beyond run noise and loss is +stable. Choose the next absolute checkpoint from the evidence (for example 300, +500, 750, or 1000), resume the same output directory with `train_iters=1000`, +and set both `eval_interval` and `exit_interval` to that target. Stop on +recovered gap, plateau, regression, divergence, or iteration 1000. Never raise +the 1000-step cap without explicit user direction. + +## 8. Report evidence + +Report exact source revision, commands, Slurm job/account/partition, container, +paths, topology derivation, dataset configs/seed/token budget/materialized +counts, PTQ format, checkpoint iterations, loss/LR/grad trend, scheduler state, +and BF16/PTQ/QAD benchmark scores. Support success with both scheduler and log +evidence; identify the first real error when a run fails. diff --git a/.agents/skills/qad/assets/nemotron-cascade-2-blend.yaml b/.agents/skills/qad/assets/nemotron-cascade-2-blend.yaml new file mode 100644 index 00000000000..cb28ce5b147 --- /dev/null +++ b/.agents/skills/qad/assets/nemotron-cascade-2-blend.yaml @@ -0,0 +1,44 @@ +tokenizer: +output_dir: /data/nemotron-cascade-2-through-150 +target_tokens: 2_600_000_000 +sources: + - hf_dataset: nvidia/Nemotron-Cascade-2-SFT-Data + config: math + split: train + content_field: messages + weight: 21.1 + - hf_dataset: nvidia/Nemotron-Cascade-2-SFT-Data + config: science + split: train + content_field: messages + weight: 10.9 + - hf_dataset: nvidia/Nemotron-Cascade-2-SFT-Data + config: chat + split: train + content_field: messages + weight: 56.2 + - hf_dataset: nvidia/Nemotron-Cascade-2-SFT-Data + config: instruction_following + split: train + content_field: messages + weight: 3.3 + - hf_dataset: nvidia/Nemotron-Cascade-2-SFT-Data + config: safety + split: train + content_field: messages + weight: 0.02 + - hf_dataset: nvidia/Nemotron-Cascade-2-SFT-Data + config: conversational_agent + split: train + content_field: messages + weight: 3.3 + - hf_dataset: nvidia/Nemotron-Cascade-2-SFT-Data + config: swe + split: train + content_field: messages + weight: 1.8 + - hf_dataset: nvidia/Nemotron-Cascade-2-SFT-Data + config: terminal_agent + split: train + content_field: messages + weight: 3.3 diff --git a/.claude/skills/qad b/.claude/skills/qad new file mode 120000 index 00000000000..9ac1b527220 --- /dev/null +++ b/.claude/skills/qad @@ -0,0 +1 @@ +../../.agents/skills/qad \ No newline at end of file diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 96880d2ae82..b1d2bca363a 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -167,6 +167,18 @@ def get_args(): parser.add_argument( "--eval_iters", type=int, default=32, help="Number of batches per validation stage" ) + parser.add_argument( + "--exit_interval", + type=int, + default=None, + help="Save a checkpoint and exit when the iteration is divisible by this value", + ) + parser.add_argument( + "--exit_duration_in_mins", + type=int, + default=None, + help="Save a checkpoint and exit after this many minutes", + ) # Logging arguments parser.add_argument("--log_interval", type=int, default=10, help="Write to log every steps") parser.add_argument( @@ -328,6 +340,8 @@ def _restore_student_hook(model_chunks): train_iters=args.train_iters, global_batch_size=args.gbs, micro_batch_size=args.mbs, + exit_interval=args.exit_interval, + exit_duration_in_mins=args.exit_duration_in_mins, manual_gc=True, manual_gc_interval=100, ), diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py index f36b727b9c4..1d6f6731b7f 100644 --- a/tests/examples/megatron_bridge/test_qad.py +++ b/tests/examples/megatron_bridge/test_qad.py @@ -68,7 +68,8 @@ def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe): hf_model_path = create_student(tmp_path) quantized_megatron_path = tmp_path / "quantized_megatron" distill_output_dir = tmp_path / "qad_output" - train_iters = 2 + train_iters = 3 + early_exit_iter = 2 # TODO: VLMs disable sequence parallelism, so tensor parallelism can't be used here. # Flip to tp_size=num_gpus in nemo:26.08 container @@ -109,13 +110,16 @@ def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe): gbs=4, train_iters=train_iters, lr_warmup_iters=2, - eval_interval=train_iters, + eval_interval=early_exit_iter, eval_iters=1, log_interval=1, + exit_interval=early_exit_iter, + exit_duration_in_mins=60, ) run_example_command(distill_cmd, example_path="megatron_bridge", setup_free_port=True) distilled_megatron_path = distill_output_dir / "checkpoints" - assert (distilled_megatron_path / "latest_checkpointed_iteration.txt").exists() + tracker = distilled_megatron_path / "latest_checkpointed_iteration.txt" + assert tracker.read_text(encoding="utf-8").strip() == str(early_exit_iter) assert list(distilled_megatron_path.rglob("modelopt_state")), ( "Expected modelopt_state to be preserved in the distilled (QAD) checkpoint" ) From afd59b43bc19fa8b1dd05e07e2389b0555c89407 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 23 Jul 2026 12:40:02 +0000 Subject: [PATCH 2/7] Preserve PTQ configuration for QAD Signed-off-by: Meng Xin --- .agents/skills/qad/SKILL.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.agents/skills/qad/SKILL.md b/.agents/skills/qad/SKILL.md index be8e57a88de..4b7d4734c82 100644 --- a/.agents/skills/qad/SKILL.md +++ b/.agents/skills/qad/SKILL.md @@ -29,6 +29,11 @@ The deprecated `examples/llm_qad` flow is not the default. ## 2. Establish the gap +When QAD follows `evaluation` or `quant-recipe-search`, reuse its exact PTQ +checkpoint and complete config or recipe. Preserve the quantization format, +layer selection, calibration dataset, sample count, sequence length, seed, and +topology; do not substitute a cheaper or otherwise different PTQ run. + Use the `evaluation` and `compare-results` skills to run and validate the same benchmark configuration on the BF16 source and the export of the exact quantized Megatron checkpoint that will initialize QAD. If the existing PTQ From 8abffb7fb3648bf8b9e98b3c50e57741af3048b9 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 23 Jul 2026 12:42:28 +0000 Subject: [PATCH 3/7] Clarify PTQ handoff to QAD Signed-off-by: Meng Xin --- .agents/skills/qad/SKILL.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/.agents/skills/qad/SKILL.md b/.agents/skills/qad/SKILL.md index 4b7d4734c82..9d23d28f9d0 100644 --- a/.agents/skills/qad/SKILL.md +++ b/.agents/skills/qad/SKILL.md @@ -29,10 +29,11 @@ The deprecated `examples/llm_qad` flow is not the default. ## 2. Establish the gap -When QAD follows `evaluation` or `quant-recipe-search`, reuse its exact PTQ -checkpoint and complete config or recipe. Preserve the quantization format, -layer selection, calibration dataset, sample count, sequence length, seed, and -topology; do not substitute a cheaper or otherwise different PTQ run. +When QAD follows `evaluation` or `quant-recipe-search`, reuse its complete PTQ +config or recipe to reproduce PTQ with `examples/megatron_bridge/quantize.py`; +do not use the usual HF checkpoint as the QAD student. Preserve the quantization +format, layer selection, calibration dataset, sample count, sequence length, +seed, and topology; do not substitute a cheaper or otherwise different PTQ run. Use the `evaluation` and `compare-results` skills to run and validate the same benchmark configuration on the BF16 source and the export of the exact From 494d46ceb52174cee9194d24bdc10159068471cf Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 23 Jul 2026 12:47:23 +0000 Subject: [PATCH 4/7] Reuse existing PTQ evaluation for QAD Signed-off-by: Meng Xin --- .agents/skills/qad/SKILL.md | 60 ++++++++++++++++++++----------------- 1 file changed, 32 insertions(+), 28 deletions(-) diff --git a/.agents/skills/qad/SKILL.md b/.agents/skills/qad/SKILL.md index 9d23d28f9d0..96a2980386b 100644 --- a/.agents/skills/qad/SKILL.md +++ b/.agents/skills/qad/SKILL.md @@ -27,21 +27,20 @@ Treat them as the source of truth for mutable flags, containers, model support, and checkpoint formats. Do not copy their full command lines into new wrappers. The deprecated `examples/llm_qad` flow is not the default. -## 2. Establish the gap - -When QAD follows `evaluation` or `quant-recipe-search`, reuse its complete PTQ -config or recipe to reproduce PTQ with `examples/megatron_bridge/quantize.py`; -do not use the usual HF checkpoint as the QAD student. Preserve the quantization -format, layer selection, calibration dataset, sample count, sequence length, -seed, and topology; do not substitute a cheaper or otherwise different PTQ run. - -Use the `evaluation` and `compare-results` skills to run and validate the same -benchmark configuration on the BF16 source and the export of the exact -quantized Megatron checkpoint that will initialize QAD. If the existing PTQ -result came from another checkpoint or its lineage is unclear, export and -benchmark the intended Megatron checkpoint first. Record its path or identifier, -metric, direction, uncertainty, and acceptable delta. Stop if the gap is not -material or the runs are not comparable. +## 2. Reuse or establish the gap + +When QAD follows `evaluation`, `compare-results`, or `quant-recipe-search`, +reuse the validated BF16/PTQ scores and exact benchmark configuration; do not +rerun valid baselines. Record their run IDs, metric, direction, uncertainty, and +acceptable delta. Use `evaluation` and `compare-results` only for missing, +unvalidated, or non-comparable results. Stop if the gap is not material. + +Reuse the preceding PTQ config or recipe to reproduce PTQ with +`examples/megatron_bridge/quantize.py`; do not use the usual HF checkpoint as +the QAD student. Preserve the quantization format, layer selection, calibration +dataset, sample count, sequence length, and seed. Choose Megatron execution +topology separately without changing quantization semantics. If any quantization +setting must change, treat it as a new PTQ candidate and evaluate it before QAD. ## 3. Set up Slurm execution @@ -58,11 +57,11 @@ target is not Slurm. Use the common account, partition, registry-auth, submission, and monitoring procedures in either case. -Keep BF16, the exact benchmarked PTQ Megatron checkpoint and its export, QAD -Megatron, exported checkpoints, data, logs, and benchmark results in the same -session/model workspace. QAD must start from the checkpoint produced by -`examples/megatron_bridge/quantize.py` because it carries the ModelOpt state; an -exported HF PTQ checkpoint is not a substitute. +Keep the reused BF16/PTQ run references and configs with the reproduced PTQ +Megatron checkpoint, QAD checkpoints, data, logs, exports, and new benchmark +results in the same session/model workspace. QAD must start from the checkpoint +produced by `examples/megatron_bridge/quantize.py` because it carries the +ModelOpt state; an exported HF PTQ checkpoint is not a substitute. ## 4. Choose topology explicitly @@ -112,11 +111,13 @@ The supported real-data `GPTDatasetConfig` shuffles documents and concatenates them into fixed-length 32K samples, so sequences are packed instead of padded. Do not use mock data as training evidence. -## 6. Retain PTQ lineage, then run staged QAD +## 6. Reproduce PTQ, then run staged QAD -Use the exact PTQ Megatron checkpoint established in step 2, or produce and -benchmark it before QAD. Follow the current Megatron Bridge README for PTQ, QAD, -resume, and export commands. Apply these QAD defaults: +Produce the required Megatron checkpoint with the PTQ config or recipe preserved +in step 2. Reuse a valid preceding PTQ evaluation; run a new PTQ evaluation only +when the baseline is missing, invalid, or non-comparable, or a quantization +setting changed. Follow the current Megatron Bridge README for PTQ, QAD, resume, +and export commands. Apply these QAD defaults: | Setting | Default | | --- | --- | @@ -145,7 +146,9 @@ At the initial exit: 1. Confirm finite QAD/KD loss, sensible learning rate, non-pathological gradient norm, and an iteration-150 checkpoint. 2. Export that QAD checkpoint with the current quantized Megatron exporter. -3. Run the identical benchmark on BF16, the step-2 PTQ export, and QAD-150. +3. Evaluate QAD-150 with the exact step-2 benchmark configuration. Reuse + validated, comparable BF16/PTQ results; run only missing, invalid, or + non-comparable baselines. 4. Calculate the original gap and recovery. For higher-is-better metrics: `gap = BF16 - PTQ`, `recovered = QAD - PTQ`, and `recovery_fraction = recovered / gap`. @@ -161,6 +164,7 @@ the 1000-step cap without explicit user direction. Report exact source revision, commands, Slurm job/account/partition, container, paths, topology derivation, dataset configs/seed/token budget/materialized -counts, PTQ format, checkpoint iterations, loss/LR/grad trend, scheduler state, -and BF16/PTQ/QAD benchmark scores. Support success with both scheduler and log -evidence; identify the first real error when a run fails. +counts, PTQ format, reused evaluation run IDs, checkpoint iterations, +loss/LR/grad trend, scheduler state, and BF16/PTQ/QAD benchmark scores. State +which scores were reused versus newly evaluated. Support success with both +scheduler and log evidence; identify the first real error when a run fails. From f8f87645f09617c80776b795fbb0c4541d159a17 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 23 Jul 2026 12:53:10 +0000 Subject: [PATCH 5/7] Clarify Megatron expert parallel topology Signed-off-by: Meng Xin --- .agents/skills/qad/SKILL.md | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/.agents/skills/qad/SKILL.md b/.agents/skills/qad/SKILL.md index 96a2980386b..ae862471274 100644 --- a/.agents/skills/qad/SKILL.md +++ b/.agents/skills/qad/SKILL.md @@ -75,12 +75,19 @@ and memory, then choose the smallest topology that fits both student and teacher 3. Use CP to distribute the 32K context when activation/attention memory is the constraint. Prefer CP before TP for a small dense model at long context. 4. Keep EP=1 for dense models. For MoE, choose EP from expert count, memory, and - available ranks; verify expert and data-parallel divisibility. + available ranks, and inspect the current `distill.py` value for expert tensor + parallelism (ETP); do not assume ETP equals TP. 5. Add PP only when layer/model state still does not fit. -Verify `DP = world_size / (TP * PP * CP)` is integral and -`GBS % (MBS * DP) == 0`. Record nodes, GPUs/node, TP/PP/CP/EP/DP, MBS, GPU -memory, and why each non-one dimension is needed. +For MoE, Megatron folds two overlapping meshes onto the same ranks within each +PP stage; EP/ETP do not multiply the dense TP/CP mesh: + +- Attention DP: `DP = world_size / (TP * PP * CP)` +- Expert DP: `EDP = world_size / (ETP * EP * PP)` + +Require both divisions to be integral, `num_experts % EP == 0`, and +`GBS % (MBS * DP) == 0`. Record nodes, GPUs/node, TP/PP/CP/EP/ETP/DP/EDP, MBS, +GPU memory, and why each non-one dimension is needed. Example only: for the requested one-node, eight-H100 Qwen3-0.6B validation at 32K, TP=1, PP=1, CP=4, EP=1, DP=2, MBS=1 is a reasonable starting point. This From ef9be330ce77cf615d3660f28f6f85c8f57b2bc9 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Thu, 23 Jul 2026 13:37:43 +0000 Subject: [PATCH 6/7] Add QAD training safety controls Signed-off-by: Meng Xin --- .agents/skills/qad/SKILL.md | 61 +++++++++++++++------- examples/megatron_bridge/distill.py | 12 +++-- tests/examples/megatron_bridge/test_qad.py | 2 + 3 files changed, 54 insertions(+), 21 deletions(-) diff --git a/.agents/skills/qad/SKILL.md b/.agents/skills/qad/SKILL.md index ae862471274..6ac55f6cfe5 100644 --- a/.agents/skills/qad/SKILL.md +++ b/.agents/skills/qad/SKILL.md @@ -1,18 +1,18 @@ --- name: qad description: >- - Recover a measured BF16-to-PTQ accuracy gap on Slurm with ModelOpt - Quantization-Aware Distillation (QAD) through Megatron Bridge. Use when PTQ - has already been benchmarked against the matching BF16 model and the - quantized checkpoint needs short, evidence-driven distillation; also use for - QAD topology, data preparation, Slurm launch, resume, checkpoint export, or - recovery decisions. + Run explicitly requested ModelOpt Quantization-Aware Distillation (QAD) on + Slurm through Megatron Bridge to recover a measured BF16-to-PTQ accuracy gap. + Use only when the user explicitly asks for QAD, including its topology, data + preparation, Slurm launch, resume, checkpoint export, or recovery decisions. --- # ModelOpt Quantization-Aware Distillation -Use QAD only to recover a material, apples-to-apples PTQ benchmark gap. Do not -train merely because a quantized checkpoint exists. +QAD is expensive. Use it only to recover a material, apples-to-apples PTQ +benchmark gap and only when the user explicitly requests QAD for the target +model or run. Do not infer permission from a quantized checkpoint, evaluation +gap, or recipe-search result. ## 1. Read the supported workflow @@ -116,7 +116,16 @@ for Megatron's train/validation split. Pass the generated `data_blend.txt` entries as `distill.py --data_paths`. The supported real-data `GPTDatasetConfig` shuffles documents and concatenates them into fixed-length 32K samples, so sequences are packed instead of padded. -Do not use mock data as training evidence. +It deterministically holds out 1% of those shuffled documents for validation +with `split="99,1,0"`; bounded materialization uses shuffle seed 42, while the +training seed controls ordering within the Megatron splits. Do not download a +duplicate validation copy. With the defaults, two GBS-512 32K validation +batches consume 33.6M tokens from the nominal 26M-token holdout, so about 1.3 +passes per validation event is accepted. Recalculate this ratio when changing +GBS, sequence length, or `eval_iters`, and increase the token budget rather than +allowing substantially more repetition. Treat this loss as a training-health +signal; the independent benchmark remains the recovery gate. Do not use mock +data as training evidence. ## 6. Reproduce PTQ, then run staged QAD @@ -134,24 +143,39 @@ and export commands. Apply these QAD defaults: | Training cap | 1000 iterations | | Global batch size | 512 | | Dataset | `nvidia/Nemotron-Cascade-2-SFT-Data` | -| Initial eval/save/exit | iteration 150 | +| Validation | every 25 iterations; deterministic 1% holdout; 2 batches | +| Checkpoint interval | 50 iterations | +| Initial benchmark/exit | iteration 150 | | Slurm duration exit | 220 minutes for a 4-hour allocation | Keep `train_iters=1000` from the first run so the cosine schedule has a stable -horizon. Set `eval_interval=150`, `exit_interval=150`, and -`exit_duration_in_mins=220` for the initial stage. The duration exit is a +horizon. Set `save_interval=50`, `eval_interval=25`, `eval_iters=2`, +`exit_interval=150`, and `exit_duration_in_mins=220` for the initial stage. This +preserves checkpoints at iterations 50, 100, and 150. The duration exit is a checkpointing safety margin, not the training target. Use one result-bearing Slurm job per stage and put Pyxis container flags on the final `srun`. Fold startup validation into that job; do not submit separate GPU preflight or smoke jobs. -## 7. Benchmark before continuing +## 7. Monitor loss, then benchmark + +Monitor the running log instead of waiting for the stage to finish. Record total +or logits-distillation loss every log interval. The expected signal is a +decreasing smoothed trend, not a decrease at every noisy step: + +1. From iteration 100 onward, compare the median loss in each 50-step window + with the preceding window. +2. Treat non-finite loss, NaN iterations, repeated skipped iterations, or a + sharp sustained increase as immediate failure. +3. Flag one flat or rising window. If a second consecutive window also fails to + decrease, stop at the checkpoint ending that second window. Preserve the + latest known-good checkpoint and diagnose before resuming. At the initial exit: -1. Confirm finite QAD/KD loss, sensible learning rate, non-pathological gradient - norm, and an iteration-150 checkpoint. +1. Confirm the smoothed QAD/KD loss decreased, learning rate is sensible, + gradient norm is non-pathological, and checkpoints 50, 100, and 150 exist. 2. Export that QAD checkpoint with the current quantized Megatron exporter. 3. Evaluate QAD-150 with the exact step-2 benchmark configuration. Reuse validated, comparable BF16/PTQ results; run only missing, invalid, or @@ -163,9 +187,10 @@ At the initial exit: Continue only if benchmark recovery is positive beyond run noise and loss is stable. Choose the next absolute checkpoint from the evidence (for example 300, 500, 750, or 1000), resume the same output directory with `train_iters=1000`, -and set both `eval_interval` and `exit_interval` to that target. Stop on -recovered gap, plateau, regression, divergence, or iteration 1000. Never raise -the 1000-step cap without explicit user direction. +retain `eval_interval=25`, `eval_iters=2`, and `save_interval=50`, and set only +`exit_interval` to that target. Stop on recovered gap, plateau, regression, +divergence, or iteration 1000. Never raise the 1000-step cap without explicit +user direction. ## 8. Report evidence diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index b1d2bca363a..332c09224db 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -161,11 +161,15 @@ def get_args(): help="Modules to recompute with --recompute_granularity=selective. Defaults to ['core_attn']. " "Allowed: core_attn, mlp, moe, moe_act, layernorm, mla_up_proj, shared_experts.", ) + parser.add_argument("--eval_interval", type=int, default=100, help="Validate every steps") parser.add_argument( - "--eval_interval", type=int, default=100, help="Validate + checkpoint every steps" + "--eval_iters", type=int, default=32, help="Number of batches per validation stage" ) parser.add_argument( - "--eval_iters", type=int, default=32, help="Number of batches per validation stage" + "--save_interval", + type=int, + default=None, + help="Checkpoint every steps; defaults to --eval_interval", ) parser.add_argument( "--exit_interval", @@ -370,7 +374,9 @@ def _restore_student_hook(model_chunks): tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size ), checkpoint=CheckpointConfig( - save_interval=args.eval_interval, + save_interval=( + args.save_interval if args.save_interval is not None else args.eval_interval + ), save=checkpoint_dir, load=checkpoint_dir, # Resume from this directory (if exists) most_recent_k=5, # Keeps 5 most recent checkpoints (not metric-based) diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py index 1d6f6731b7f..92c122bbd8f 100644 --- a/tests/examples/megatron_bridge/test_qad.py +++ b/tests/examples/megatron_bridge/test_qad.py @@ -112,6 +112,7 @@ def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe): lr_warmup_iters=2, eval_interval=early_exit_iter, eval_iters=1, + save_interval=1, log_interval=1, exit_interval=early_exit_iter, exit_duration_in_mins=60, @@ -120,6 +121,7 @@ def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe): distilled_megatron_path = distill_output_dir / "checkpoints" tracker = distilled_megatron_path / "latest_checkpointed_iteration.txt" assert tracker.read_text(encoding="utf-8").strip() == str(early_exit_iter) + assert (distilled_megatron_path / "iter_0000001").is_dir() assert list(distilled_megatron_path.rglob("modelopt_state")), ( "Expected modelopt_state to be preserved in the distilled (QAD) checkpoint" ) From 61aa5f07cd17f8800f221efe96505538620eafb7 Mon Sep 17 00:00:00 2001 From: Meng Xin Date: Fri, 24 Jul 2026 04:01:01 +0000 Subject: [PATCH 7/7] Add QAD PTQ compatibility gate Signed-off-by: Meng Xin --- .agents/skills/qad/SKILL.md | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/.agents/skills/qad/SKILL.md b/.agents/skills/qad/SKILL.md index 6ac55f6cfe5..b9c26e0d4a5 100644 --- a/.agents/skills/qad/SKILL.md +++ b/.agents/skills/qad/SKILL.md @@ -42,6 +42,15 @@ dataset, sample count, sequence length, and seed. Choose Megatron execution topology separately without changing quantization semantics. If any quantization setting must change, treat it as a new PTQ candidate and evaluate it before QAD. +In the exact runtime, confirm `AutoBridge.can_handle()` accepts the model, then +reproduce PTQ with `quantize.py`. Sanity-check the checkpoint's master-rank +`.quant_summary.txt`: every enabled statically calibrated quantizer it lists +must have a finite, positive `amax`. Accept `dynamic` or format-defined `None` +only when the preserved recipe intentionally uses dynamic or MX quantization; +with model parallelism, do not treat this rank-local summary as exhaustive. +Stop on unsupported conversion, failed PTQ, missing ModelOpt state, or +unexpected `amax`. + ## 3. Set up Slurm execution Read `skills/common/environment-setup.md` and detect whether the Slurm target is @@ -129,11 +138,11 @@ data as training evidence. ## 6. Reproduce PTQ, then run staged QAD -Produce the required Megatron checkpoint with the PTQ config or recipe preserved -in step 2. Reuse a valid preceding PTQ evaluation; run a new PTQ evaluation only -when the baseline is missing, invalid, or non-comparable, or a quantization -setting changed. Follow the current Megatron Bridge README for PTQ, QAD, resume, -and export commands. Apply these QAD defaults: +Use the verified Megatron PTQ checkpoint from step 2. Reuse a valid preceding +PTQ evaluation; run a new PTQ evaluation only when the baseline is missing, +invalid, or non-comparable, or a quantization setting changed. Follow the +current Megatron Bridge README for QAD, resume, and export commands. Apply these +QAD defaults: | Setting | Default | | --- | --- |