Add ModelOpt QAD skill for Slurm workflows (#2010)

### What does this PR do?

Type of change: new feature

Adds a general Slurm-only QAD skill based on the supported Megatron
Bridge
workflow. The skill:

- starts from a measured BF16-to-PTQ benchmark gap and preserves the
preceding
  PTQ configuration or recipe;
- gates QAD on exact Megatron Bridge model support and successful
Megatron PTQ,
using its master-rank quantizer summary as a scoped `amax` sanity check;
- requires model- and hardware-derived TP/PP/CP/EP/ETP topology
selection;
- streams and randomly samples only the required
`nvidia/Nemotron-Cascade-2-SFT-Data` token budget and uses Megatron
sequence
  packing;
- defaults to 32K sequences, LR `1e-5` with cosine decay, a 1000-step
cap, and
  GBS 512;
- requires explicit user authorization because QAD is costly, validates
two
batches every 25 steps, saves every 50 steps, and monitors a decreasing
  smoothed loss trend;
- evaluates an early checkpoint around step 150 and continues only when
  benchmark recovery and the loss trend justify more training;
- follows the established common Slurm and remote-execution guidance
instead of
  duplicating mutable commands from the Megatron Bridge README.

Also exposes Megatron Bridge `save_interval`, `exit_interval`, and
`exit_duration_in_mins` through `examples/megatron_bridge/distill.py`,
with example-test coverage for checkpoint
and ModelOpt-state preservation at an early exit.

### Usage

```text
Use the QAD skill to recover the measured BF16-to-PTQ benchmark gap for
<model> on <Slurm cluster>, preserving the validated PTQ recipe.
```

### Testing

- `PYTHONPATH=$PWD pre-commit run --all-files`
- Passed every hook on the rebased branch, including Ruff, Ruff format,
mypy,
YAML/recipe validation, launcher reference validation, Bandit, generated
    arguments, symlink synchronization, and Markdown lint.
- `python
~/.codex/skills/.system/skill-creator/scripts/quick_validate.py
.agents/skills/qad`
  - `Skill is valid!`

Qwen3-0.6B result-bearing validation:

- Resources: one exclusive node, 8 H100 GPUs
- Container: `nvcr.io/nvidia/nemo:26.06`
- Quantization: NVFP4, group size 16, embedding excluded
- QAD topology: TP=1, PP=1, CP=4, EP=1, DP=2
- Training validation configuration: sequence length 32768, MBS=1,
GBS=8,
`train_iters=1000`, LR `1e-5` / minimum LR `1e-6`, 50 warmup iterations,
  cosine decay, `eval_interval=150`, `exit_interval=150`,
  `exit_duration_in_mins=220`
- This result-bearing run used the then-current coupled eval/save
cadence. The
final skill now validates two batches every 25 steps and saves every 50;
the
  example test covers the independent checkpoint cadence.
- The reduced GBS 8 is intentionally validation-only; the skill retains
GBS 512
  as the production default.
- Data: exactly 10,000,000 sampled tokens from four
  `nvidia/Nemotron-Cascade-2-SFT-Data` configs:
  - math: 2,306,011 tokens / 364 documents
  - science: 1,191,257 tokens / 285 documents
  - chat: 6,142,077 tokens / 1,800 documents
  - instruction following: 360,655 tokens / 411 documents
- Megatron built packed 32K GPT samples from the materialized prefixes;
the full
  dataset was not downloaded.
- QAD loss was finite and decreased from `0.2640341` at iteration 10 to
`0.1060580` at iteration 150. Final gradient norm was `0.747`, with zero
  skipped and zero NaN iterations. Validation distillation loss was
  `0.09715855`.
- The iteration-150 checkpoint saved successfully with `modelopt_state`,
and
  both PTQ and QAD-150 exported to unified Hugging Face format.
- Identical full MMLU 0-shot comparison through the Megatron evaluator:

  | Model | Accuracy |
  | --- | ---: |
  | BF16 | 0.39517164 |
  | PTQ | 0.32851446 |
  | QAD-150 | 0.38740921 |

QAD-150 recovered `0.05889475 / 0.06665718 = 88.35%` of the measured PTQ
gap,
so validation stopped at the early evidence gate rather than continuing
  blindly toward 1000 iterations.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did
  you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update Changelog?: N/A — this adds an agent skill and
example-only
  lifecycle flags.
- Did you get Claude approval on this PR?: N/A

### Additional Information

All seven branch commits are cryptographically signed and include a
`Signed-off-by` trailer.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Updated the QAD skill documentation with a clear “Execute in this
order” workflow, including a revised default recovery training policy.
* Added a new `nemotron-cascade-2` dataset blend configuration with an
increased token budget.
* Enhanced the MeGatron Bridge distillation CLI with stricter interval
argument validation and support for configurable save-and-exit controls.
* **Documentation**
* Expanded Megatron Bridge README guidance for dataset preparation,
token-budget recalculation, and resume expectations.
* **Tests**
* Improved distillation and QAD tests to validate early-exit behavior
and checkpoint expectations.
  * Added unit tests covering distillation CLI interval validation.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Meng Xin <mxin@nvidia.com>

Author: mxinO

.agents/skills/qad/SKILL.md

@@ -0,0 +1,106 @@
+---
+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. Run it only when the user explicitly authorizes QAD for the
+target model or run. A Day-0, PTQ, evaluation, comparison, or recipe-search
+request alone is not authorization to start QAD.
+
+## Follow the supported workflow
+
+Before constructing commands, read:
+
+- `examples/megatron_bridge/README.md`, especially PTQ, data preparation, QAD,
+  export, and Slurm usage
+- `examples/megatron_bridge/{quantize.py,distill.py}` via `--help`
+- `skills/common/{environment-setup,workspace-management,slurm-setup}.md`; also
+  `skills/common/remote-execution.md` for remote Slurm
+
+Treat the example README and `--help` output as authoritative for mutable flags,
+commands, containers, and checkpoint formats. This skill supports Slurm only.
+
+## Execute in this order
+
+1. **Confirm the gap.** Reuse only validated, comparable BF16/PTQ results and
+   the exact benchmark configuration from preceding evaluation or recipe
+   search; run missing, invalid, or non-comparable baselines. Confirm the target
+   benchmarks and their context-length needs. Stop if the PTQ gap to BF16 is
+   already below 1%.
+2. **Reproduce PTQ and verify compatibility.** In the target runtime, require
+   `AutoBridge.can_handle()` for the target model and PTQ through `quantize.py`
+   to succeed while preserving the exact preceding PTQ config or recipe:
+   format, layer selection, calibration data/count, sequence length, and seed.
+   A changed quantization setting is a new PTQ candidate and must be evaluated
+   before QAD. In the master-rank `.quant_summary.txt`, require finite positive
+   `amax` for enabled static quantizers; accept `dynamic`/format-defined `None`
+   only when the recipe intends it. Treat the summary as rank-local under model
+   parallelism.
+3. **Choose topology explicitly.** Derive the smallest fitting node count and
+   TP/PP/CP/EP from student and teacher architecture, the chosen sequence length,
+   and available GPU memory. Prefer CP before TP for small long-context models;
+   keep EP=1 for dense models and ETP=1 because the current `distill.py`
+   workflow does not support expert tensor parallelism. For MoE require:
+
+   - `DP = world_size / (TP * PP * CP)`
+   - `EDP = world_size / (EP * PP)`
+   - integral DP/EDP, `num_experts % EP == 0`, and
+     `GBS % (MBS * DP) == 0`
+
+4. **Prepare the full capped dataset once.** Use suitable user-provided data, or
+   copy `examples/megatron_bridge/data/nemotron-cascade-2-blend.yaml` as the
+   default. Set the target tokenizer and workspace path, then materialize the
+   randomly sampled subset before training. Pack the chosen sequence length;
+   Megatron's `99,1,0` split creates the 1% validation holdout from the same
+   data.
+5. **Run and monitor QAD.** Run one QAD training job at a time and fold startup
+   validation into it; do not submit separate GPU preflight jobs or split at
+   recovery iterations. Let training continue while evaluating saved
+   checkpoints, and cancel it when a stop condition below is met.
+
+## Default training policy
+
+| Setting | Default |
+| --- | --- |
+| Sequence length | 32768; adjust for target benchmarks |
+| 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` by default |
+| Materialized token budget | 17.3B at 32K; cover the full cap at the chosen length |
+| Training validation | every 25 iterations; deterministic 1% holdout; 2 batches |
+| Checkpoint interval | 50 iterations |
+| Loss logging | every 10 iterations |
+| Recovery benchmark | 150, then every 100 iterations while training runs |
+| Slurm duration exit | 220 minutes for a 4-hour allocation |
+
+## Run policy
+
+- Keep `train_iters=1000` and leave `exit_interval` unset.
+- From initial step timing, submit only enough sequential jobs to reach
+  checkpoint 150; never submit through iteration 1000 upfront. At each recovery
+  checkpoint, submit to the next only after its targeted evaluation and any
+  triggered full suite, and only if the BF16 gap remains at least 1% and
+  recovery has neither plateaued nor regressed.
+- Give all training jobs the same run-specific job name and
+  `--dependency=singleton`; record job IDs and, on any stop, cancel pending jobs
+  before the active job.
+- Cancel on non-finite loss, repeated skipped iterations, or a sustained spike.
+  At iteration 50, require the loss aggregate to be lower than at iteration 10.
+- At each recovery checkpoint, first evaluate the one to three benchmarks with
+  the largest PTQ drops. Run the remaining original PTQ suite at that checkpoint
+  only after recovery beyond run noise.
+- Cancel when the full-suite gap to BF16 is below 1%, benchmark recovery
+  regresses beyond run noise, or benchmark recovery and loss both plateau.
+- After a duration exit, resume the latest QAD checkpoint in the same output
+  directory with unchanged prepared data paths/cache, seed, topology, optimizer,
+  scheduler, iteration, and consumed-sample state; do not restart from PTQ.
+- Report the PTQ recipe, data sample, Slurm topology, loss/state, checkpoints,
+  and comparable BF16/PTQ/QAD results.

.claude/skills/qad

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

examples/megatron_bridge/README.md

@@ -250,7 +250,16 @@ To export selected iterations instead, use `--export_iterations 200 400 600`.
 
 ### Quantization Aware Distillation (QAD)
 
-To recover the accuracy lost during [Post-Training Quantization](#post-training-quantization), distill the quantized model (student) from the original, unquantized model (teacher). Pass the quantized **Megatron checkpoint** produced by `quantize.py` via `--student_megatron_path` (the ModelOpt quantizers are restored automatically, so distillation trains the fake-quantized student), while `--student_hf_path` provides the student architecture and `--teacher_hf_path` points to the original unquantized model. We also use a smaller learning rate for QAD:
+To recover the accuracy lost during [Post-Training Quantization](#post-training-quantization), distill the quantized model (student) from the original, unquantized model (teacher). Pass the quantized **Megatron checkpoint** produced by `quantize.py` via `--student_megatron_path` (the ModelOpt quantizers are restored automatically, so distillation trains the fake-quantized student), while `--student_hf_path` provides the student architecture and `--teacher_hf_path` points to the original unquantized model.
+
+If you do not already have a suitable QAD dataset, start with
+[data/nemotron-cascade-2-blend.yaml](data/nemotron-cascade-2-blend.yaml). It defines a general-purpose
+mixture of SFT data for QAD. Copy it, set the tokenizer for the target model, and adjust the output directory,
+sources, and weights as needed before preparing data. Its default 17.3-billion-token budget covers 1000
+iterations at global batch size 512 and sequence length 32768, including a 1% validation holdout and margin.
+Recalculate the budget when changing those settings, and keep the prepared data unchanged when resuming.
+
+We also use a smaller learning rate for QAD:
 
 ```bash
 torchrun --nproc_per_node 8 distill.py \

examples/megatron_bridge/data/nemotron-cascade-2-blend.yaml

@@ -0,0 +1,45 @@
+# Set to the target model's Hugging Face ID or local tokenizer path.
+tokenizer: <target-model-tokenizer>
+output_dir: <session-model-workspace>/data/nemotron-cascade-2-through-1000
+target_tokens: 17_300_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

examples/megatron_bridge/distill.py

@@ -59,6 +59,20 @@
     import modelopt.torch.puzzletron.plugins.mbridge  # noqa: F401
 
 
+def _positive_int(value: str) -> int:
+    parsed = int(value)
+    if parsed <= 0:
+        raise argparse.ArgumentTypeError("must be a positive integer")
+    return parsed
+
+
+def _nonnegative_int(value: str) -> int:
+    parsed = int(value)
+    if parsed < 0:
+        raise argparse.ArgumentTypeError("must be a non-negative integer")
+    return parsed
+
+
 def get_args():
     """Parse command-line arguments."""
     parser = argparse.ArgumentParser(description="Distillation for Megatron-Bridge.")
@@ -162,10 +176,31 @@ def get_args():
         "Allowed: core_attn, mlp, moe, moe_act, layernorm, mla_up_proj, shared_experts.",
     )
     parser.add_argument(
-        "--eval_interval", type=int, default=100, help="Validate + checkpoint every <N> steps"
+        "--eval_interval", type=_positive_int, default=100, help="Validate every <N> steps"
     )
     parser.add_argument(
-        "--eval_iters", type=int, default=32, help="Number of batches per validation stage"
+        "--eval_iters",
+        type=_nonnegative_int,
+        default=32,
+        help="Number of batches per validation stage; set to 0 to disable validation",
+    )
+    parser.add_argument(
+        "--save_interval",
+        type=_positive_int,
+        default=None,
+        help="Checkpoint every <N> steps; defaults to --eval_interval",
+    )
+    parser.add_argument(
+        "--exit_interval",
+        type=_positive_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=_positive_int,
+        default=None,
+        help="Save a checkpoint and exit after this many minutes",
     )
     parser.add_argument(
         "--validate_only",
@@ -218,8 +253,8 @@ def get_args():
         args.student_hf_model = args.student_hf_path
     if args.checkpoint_keep_last < -1:
         raise ValueError("--checkpoint_keep_last must be >= -1.")
-    if args.validate_only and (args.eval_interval <= 0 or args.eval_iters <= 0):
-        raise ValueError("--validate_only requires --eval_interval > 0 and --eval_iters > 0.")
+    if args.validate_only and args.eval_iters == 0:
+        raise ValueError("--validate_only requires --eval_iters > 0.")
 
     print_args(args)
 
@@ -347,6 +382,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 +416,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=10,
     )
     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"
     )