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
• the common skill's environment-setup.md, workspace-management.md, and slurm-setup.md; also its 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.