diff --git a/.agents/skills/qad/SKILL.md b/.agents/skills/qad/SKILL.md new file mode 100644 index 00000000000..b9c26e0d4a5 --- /dev/null +++ b/.agents/skills/qad/SKILL.md @@ -0,0 +1,211 @@ +--- +name: qad +description: >- + 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 + +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 + +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. 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. + +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 +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 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 + +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, 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. + +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 +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. +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 + +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 | +| --- | --- | +| 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` | +| 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 `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. 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 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 + 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`. + +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`, +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 + +Report exact source revision, commands, Slurm job/account/partition, container, +paths, topology derivation, dataset configs/seed/token budget/materialized +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. 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 b0f6f1af865..867ac2c9edd 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -161,11 +161,27 @@ 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", + 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", ) parser.add_argument( "--validate_only", @@ -347,6 +363,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, ), @@ -379,7 +397,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=args.checkpoint_keep_last, # Keeps most recent checkpoints (-1 keeps all) diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py index f36b727b9c4..92c122bbd8f 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,18 @@ 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, + save_interval=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 (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" )