Add ModelOpt QAD skill for Slurm workflows

.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 <copied-blend.yaml>
+```
+
+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.

.agents/skills/qad/assets/nemotron-cascade-2-blend.yaml

@@ -0,0 +1,44 @@
+tokenizer: <target-model-tokenizer>
+output_dir: <session-model-workspace>/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

.claude/skills/qad

@@ -0,0 +1 @@
+../../.agents/skills/qad

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 <N> steps")
     parser.add_argument(
-        "--eval_interval", type=int, default=100, help="Validate + checkpoint every <N> 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 <N> 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)

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"
     )