Add ModelOpt QAD skill for Slurm workflows - #2010
Conversation
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
📝 WalkthroughWalkthroughAdds a QAD skill specification, a Nemotron-Cascade-2 blend configuration, a shared skill symlink, and configurable distillation checkpoint and early-exit controls with strengthened QAD test assertions. ChangesQAD workflow support
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2010 +/- ##
==========================================
+ Coverage 66.83% 75.51% +8.67%
==========================================
Files 519 519
Lines 58916 58916
==========================================
+ Hits 39376 44488 +5112
+ Misses 19540 14428 -5112
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com> # Conflicts: # examples/megatron_bridge/distill.py
Signed-off-by: Meng Xin <mxin@nvidia.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
🧹 Nitpick comments (1)
.agents/skills/qad/assets/nemotron-cascade-2-blend.yaml (1)
5-43: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winPin the dataset revision for reproducible data.
All entries use the floating
nvidia/Nemotron-Cascade-2-SFT-Dataidentifier without a revision, while.agents/skills/qad/SKILL.mdLines 206-209 require reporting an exact source revision. Pin the revision in the blend schema if supported, or capture and validate the resolved revision in the generated artifact.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In @.agents/skills/qad/assets/nemotron-cascade-2-blend.yaml around lines 5 - 43, Pin every nvidia/Nemotron-Cascade-2-SFT-Data entry in the blend configuration to an immutable dataset revision using the schema’s supported revision field. If the blend schema cannot express revisions, update the generation flow to record and validate the resolved revision in the output artifact, preserving the requirement from SKILL.md for an exact source revision.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.agents/skills/qad/SKILL.md:
- Around line 160-164: Update the initial-stage instructions around the training
configuration and the QAD-150 export/evaluation steps to record the actual exit
iteration, distinguish duration-triggered exits before iteration 150, and run
the 150-step gate only when checkpoint 150 exists; otherwise mark the stage
incomplete and direct the operator to resume or relaunch instead of labeling an
earlier checkpoint as QAD-150.
In `@examples/megatron_bridge/distill.py`:
- Around line 164-184: Validate eval_interval, eval_iters, save_interval,
exit_interval, and exit_duration_in_mins immediately after CLI parsing,
rejecting zero or negative values before they reach scheduling or validation
logic; if any zero value is intentionally supported as a disable sentinel,
handle it explicitly and consistently instead.
---
Nitpick comments:
In @.agents/skills/qad/assets/nemotron-cascade-2-blend.yaml:
- Around line 5-43: Pin every nvidia/Nemotron-Cascade-2-SFT-Data entry in the
blend configuration to an immutable dataset revision using the schema’s
supported revision field. If the blend schema cannot express revisions, update
the generation flow to record and validate the resolved revision in the output
artifact, preserving the requirement from SKILL.md for an exact source revision.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8c080adc-5bf5-4c53-95fb-5a0e5b8a1bad
📒 Files selected for processing (5)
.agents/skills/qad/SKILL.md.agents/skills/qad/assets/nemotron-cascade-2-blend.yaml.claude/skills/qadexamples/megatron_bridge/distill.pytests/examples/megatron_bridge/test_qad.py
| 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. |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Handle duration exits before iteration 150 explicitly.
exit_duration_in_mins can save and exit before exit_interval=150. The following instructions nevertheless require checkpoints 50, 100, and 150 and direct the operator to export/evaluate “QAD-150”, which can cause a missing-checkpoint failure or mislabel an earlier checkpoint. Record the actual exit iteration; run the 150-step gate only when checkpoint 150 exists, otherwise mark the stage incomplete and resume or relaunch.
Also applies to: 184-190
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In @.agents/skills/qad/SKILL.md around lines 160 - 164, Update the initial-stage
instructions around the training configuration and the QAD-150 export/evaluation
steps to record the actual exit iteration, distinguish duration-triggered exits
before iteration 150, and run the 150-step gate only when checkpoint 150 exists;
otherwise mark the stage incomplete and direct the operator to resume or
relaunch instead of labeling an earlier checkpoint as QAD-150.
| 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", |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the target file around the referenced lines.
sed -n '140,230p' examples/megatron_bridge/distill.py
# Find where these args are consumed.
rg -n "eval_interval|eval_iters|save_interval|exit_interval|exit_duration_in_mins" examples/megatron_bridge/distill.pyRepository: NVIDIA/Model-Optimizer
Length of output: 4465
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Search for validation patterns and relevant config wiring.
rg -n "max\(.*0|assert .*interval|if .*interval|interval .*>=|>= 0|> 0|validate.*interval|check.*interval" examples/megatron_bridge/distill.pyRepository: NVIDIA/Model-Optimizer
Length of output: 441
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- distill.py excerpt ---'
sed -n '1,260p' examples/megatron_bridge/distill.py
echo
echo '--- interval references ---'
rg -n "eval_interval|eval_iters|save_interval|exit_interval|exit_duration_in_mins" examples/megatron_bridge/distill.py
echo
echo '--- validation patterns ---'
rg -n "max\(.*0|assert .*interval|if .*interval|interval .*>=|>= 0|> 0|validate.*interval|check.*interval" examples/megatron_bridge/distill.pyRepository: NVIDIA/Model-Optimizer
Length of output: 11985
🏁 Script executed:
#!/bin/bash
set -euo pipefail
nl -ba examples/megatron_bridge/distill.py | sed -n '150,240p'Repository: NVIDIA/Model-Optimizer
Length of output: 200
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Look for other validation or semantic handling of the interval fields in the repo.
rg -n "save_interval|exit_interval|exit_duration_in_mins|eval_interval|eval_iters" .
# Inspect the distillation config wiring where these args are passed through.
sed -n '340,420p' examples/megatron_bridge/distill.pyRepository: NVIDIA/Model-Optimizer
Length of output: 8113
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Search for these names in the repository, but keep the output compact.
rg -n "save_interval|exit_interval|exit_duration_in_mins|eval_interval|eval_iters" examples megatron_bridge . -g '!**/.git/**' -g '!**/node_modules/**' || true
echo
echo '--- wiring excerpt ---'
sed -n '340,420p' examples/megatron_bridge/distill.pyRepository: NVIDIA/Model-Optimizer
Length of output: 11754
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- wiring excerpt ---'
sed -n '340,420p' examples/megatron_bridge/distill.py | cat -nRepository: NVIDIA/Model-Optimizer
Length of output: 3825
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Find any local handling or documentation of the interval flags.
rg -n "save_interval|exit_interval|exit_duration_in_mins|eval_interval|eval_iters" examples/megatron_bridgeRepository: NVIDIA/Model-Optimizer
Length of output: 2564
Validate the interval args at the CLI boundary. --eval_interval and --eval_iters are only checked under --validate_only, and --save_interval, --exit_interval, and --exit_duration_in_mins are passed through unchecked. Reject zero/negative values explicitly, or document a sentinel if 0 is meant to disable scheduling.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/megatron_bridge/distill.py` around lines 164 - 184, Validate
eval_interval, eval_iters, save_interval, exit_interval, and
exit_duration_in_mins immediately after CLI parsing, rejecting zero or negative
values before they reach scheduling or validation logic; if any zero value is
intentionally supported as a disable sentinel, handle it explicitly and
consistently instead.
Source: Coding guidelines
1 participant
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:
PTQ configuration or recipe;
using its master-rank quantizer summary as a scoped
amaxsanity check;nvidia/Nemotron-Cascade-2-SFT-Datatoken budget and uses Megatron sequencepacking;
1e-5with cosine decay, a 1000-step cap, andGBS 512;
batches every 25 steps, saves every 50 steps, and monitors a decreasing
smoothed loss trend;
benchmark recovery and the loss trend justify more training;
duplicating mutable commands from the Megatron Bridge README.
Also exposes Megatron Bridge
save_interval,exit_interval, andexit_duration_in_minsthroughexamples/megatron_bridge/distill.py, with example-test coverage for checkpointand ModelOpt-state preservation at an early exit.
Usage
Testing
PYTHONPATH=$PWD pre-commit run --all-filesYAML/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/qadSkill is valid!Qwen3-0.6B result-bearing validation:
Resources: one exclusive node, 8 H100 GPUs
Container:
nvcr.io/nvidia/nemo:26.06Quantization: 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, LR1e-5/ minimum LR1e-6, 50 warmup iterations,cosine decay,
eval_interval=150,exit_interval=150,exit_duration_in_mins=220This 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-Dataconfigs:Megatron built packed 32K GPT samples from the materialized prefixes; the full
dataset was not downloaded.
QAD loss was finite and decreased from
0.2640341at iteration 10 to0.1060580at iteration 150. Final gradient norm was0.747, with zeroskipped and zero NaN iterations. Validation distillation loss was
0.09715855.The iteration-150 checkpoint saved successfully with
modelopt_state, andboth PTQ and QAD-150 exported to unified Hugging Face format.
Identical full MMLU 0-shot comparison through the Megatron evaluator:
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"
you follow guidance in
CONTRIBUTING.md: N/Alifecycle flags.
Additional Information
All seven branch commits are cryptographically signed and include a
Signed-off-bytrailer.Summary by CodeRabbit
New Features
Enhancements
Tests