Add ModelOpt QAD skill for Slurm workflows

.agents/skills/qad/SKILL.md

@@ -0,0 +1,104 @@
+---
+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. 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/ETP from student and teacher architecture, 32K activation memory,
+   and available GPU memory. Prefer CP before TP for small long-context models;
+   keep EP=1 for dense models. For MoE require:
+
+   - `DP = world_size / (TP * PP * CP)`
+   - `EDP = world_size / (ETP * EP * PP)`
+   - integral DP/EDP, `num_experts % EP == 0`, and
+     `GBS % (MBS * DP) == 0`
+
+4. **Prepare the full capped dataset once.** Copy
+   `examples/megatron_bridge/data/nemotron-cascade-2-blend.yaml`, set the target
+   tokenizer and workspace path, and materialize it before training. Use the
+   generated random sample for packed 32K sequences and its 1% validation
+   holdout; do not download separate validation data or use mock data as
+   evidence.
+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 |
+| 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` |
+| Materialized token budget | 17.3B, prepared once before training |
+| 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

@@ -130,6 +130,13 @@ The distillation script expects pre-tokenized data in Megatron's binary format (
 See the **[Dataset Preparation README](../dataset/README.md#tokenizing-for-megatron-frameworks)**
 for full instructions on tokenizing JSONL files and Hugging Face datasets and get the list of output prefixes that you can use for `--data_paths` argument.
 
+If you do not already have a suitable 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.
+
 ### Distillation with Real Data
 
 Example usage to distill a 4B student (HF) from an 8B teacher (HF) on 8 GPUs (TP=8, PP=1):

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_distill.py

@@ -14,8 +14,10 @@
 # limitations under the License.
 """Tests for prune_minitron.py and distill.py scripts."""
 
+import sys
 from pathlib import Path
 
+import pytest
 import torch
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
 from _test_utils.torch.puzzletron.utils import create_and_save_small_hf_model
@@ -28,6 +30,71 @@
 
 from modelopt.torch.puzzletron.anymodel import convert_model
 
+# distill.py imports sibling example modules by filename, so make that directory importable first.
+sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "examples" / "megatron_bridge"))
+
+import distill
+
+
+@pytest.mark.parametrize(
+    ("option", "value", "message"),
+    [
+        ("--eval_interval", "0", "must be a positive integer"),
+        ("--eval_iters", "-1", "must be a non-negative integer"),
+        ("--save_interval", "0", "must be a positive integer"),
+        ("--exit_interval", "-1", "must be a positive integer"),
+        ("--exit_duration_in_mins", "0", "must be a positive integer"),
+    ],
+)
+def test_distill_rejects_invalid_intervals(monkeypatch, capsys, option, value, message):
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "distill.py",
+            "--student_hf_path",
+            "unused",
+            "--teacher_hf_path",
+            "unused",
+            "--output_dir",
+            "unused",
+            "--train_iters",
+            "1",
+            "--use_mock_data",
+            option,
+            value,
+        ],
+    )
+
+    with pytest.raises(SystemExit) as exc_info:
+        distill.get_args()
+
+    assert exc_info.value.code == 2
+    assert message in capsys.readouterr().err
+
+
+def test_distill_allows_zero_eval_iters(monkeypatch):
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "distill.py",
+            "--student_hf_path",
+            "unused",
+            "--teacher_hf_path",
+            "unused",
+            "--output_dir",
+            "unused",
+            "--train_iters",
+            "1",
+            "--use_mock_data",
+            "--eval_iters",
+            "0",
+        ],
+    )
+
+    assert distill.get_args().eval_iters == 0
+
 
 def test_distill_llm(tmp_path, num_gpus):
     teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True)

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