diff --git a/.agents/skills/qad/SKILL.md b/.agents/skills/qad/SKILL.md new file mode 100644 index 00000000000..d016fdf65d0 --- /dev/null +++ b/.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. 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/README.md b/examples/megatron_bridge/README.md index 45606d691ee..5618234fddb 100644 --- a/examples/megatron_bridge/README.md +++ b/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): diff --git a/examples/megatron_bridge/data/nemotron-cascade-2-blend.yaml b/examples/megatron_bridge/data/nemotron-cascade-2-blend.yaml new file mode 100644 index 00000000000..ca4afdd4bdb --- /dev/null +++ b/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: +output_dir: /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 diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index b0f6f1af865..16dd37d5f8d 100644 --- a/examples/megatron_bridge/distill.py +++ b/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 steps" + "--eval_interval", type=_positive_int, default=100, help="Validate every 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 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) 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" )