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. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds an authorized Slurm-based QAD workflow, a materialized Nemotron Cascade dataset blend, and configurable distillation save/exit controls. Tests cover CLI validation, zero evaluation iterations, early termination, and checkpoint state. ChangesQAD workflow support
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant DistillCLI
participant TrainingConfig
participant CheckpointConfig
participant CheckpointFilesystem
DistillCLI->>TrainingConfig: pass exit interval and duration
DistillCLI->>CheckpointConfig: select save interval
TrainingConfig->>CheckpointFilesystem: save checkpoint and exit
CheckpointFilesystem-->>DistillCLI: report latest checkpointed iteration
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 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 60.21% 61.82% +1.60%
==========================================
Files 519 519
Lines 59377 59377
==========================================
+ Hits 35755 36709 +954
+ Misses 23622 22668 -954
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
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
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: 1
🧹 Nitpick comments (1)
tests/examples/megatron_bridge/test_distill.py (1)
68-69: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert the exit code directly.
match="2"performs a regex match against the exception text and is not an exact code assertion. Capture the exception and checkexc_info.value.code == 2.🤖 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 `@tests/examples/megatron_bridge/test_distill.py` around lines 68 - 69, Update the get_args test to capture the SystemExit exception from distill.get_args() and assert that exc_info.value.code equals 2, replacing the regex-based match="2" check while preserving the expected SystemExit behavior.
🤖 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 `@tests/examples/megatron_bridge/test_distill.py`:
- Around line 45-48: Move the dynamic distill import out of both test bodies and
load the module once at file scope, after making the example directory
importable. Remove the duplicated importlib.import_module calls from
test_distill_rejects_invalid_intervals and the other test, preserving their
existing test behavior; only retain dynamic loading if a documented optional or
circular dependency requires it.
---
Nitpick comments:
In `@tests/examples/megatron_bridge/test_distill.py`:
- Around line 68-69: Update the get_args test to capture the SystemExit
exception from distill.get_args() and assert that exc_info.value.code equals 2,
replacing the regex-based match="2" check while preserving the expected
SystemExit behavior.
🪄 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: 5638b3cc-36d5-46f8-a5bf-961199ff07aa
📒 Files selected for processing (3)
.agents/skills/qad/SKILL.mdexamples/megatron_bridge/distill.pytests/examples/megatron_bridge/test_distill.py
🚧 Files skipped from review as they are similar to previous changes (2)
- examples/megatron_bridge/distill.py
- .agents/skills/qad/SKILL.md
Signed-off-by: Meng Xin <mxin@nvidia.com>
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: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/megatron_bridge/data/nemotron-cascade-2-blend.yaml (1)
5-45: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winAdd a dataset revision to each Hugging Face source
prepare_megatron_data_blend.pyforwardshf_dataset/hf_name/hf_splitwith streaming and no revision field, so these sources can drift and make BF16/PTQ/QAD runs non-reproducible. Add arevisionkey here and thread it through, or persist the resolved commit 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 `@examples/megatron_bridge/data/nemotron-cascade-2-blend.yaml` around lines 5 - 45, Add an explicit Hugging Face revision to every source entry in the YAML blend, using the same pinned revision for the Nemotron-Cascade-2-SFT-Data dataset. Ensure prepare_megatron_data_blend.py propagates each source’s revision through the streaming load path, or records the resolved commit in the generated artifact so runs remain reproducible.
🤖 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 96-101: Update the Megatron resume policy in the recovery
instructions to allow extending or removing the duration-based exit limit when a
later recovery target requires more time. Continue preserving prepared data,
checkpoint lineage, optimizer/scheduler state, iteration, consumed samples,
topology, and all other training arguments; retain the absolute exit_interval
only when sufficient.
---
Outside diff comments:
In `@examples/megatron_bridge/data/nemotron-cascade-2-blend.yaml`:
- Around line 5-45: Add an explicit Hugging Face revision to every source entry
in the YAML blend, using the same pinned revision for the
Nemotron-Cascade-2-SFT-Data dataset. Ensure prepare_megatron_data_blend.py
propagates each source’s revision through the streaming load path, or records
the resolved commit in the generated artifact so runs remain reproducible.
🪄 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: d29c7d8f-3e17-4f3d-b6c0-fd303c6a1c75
📒 Files selected for processing (4)
.agents/skills/qad/SKILL.mdexamples/megatron_bridge/README.mdexamples/megatron_bridge/data/nemotron-cascade-2-blend.yamltests/examples/megatron_bridge/test_distill.py
🚧 Files skipped from review as they are similar to previous changes (1)
- examples/megatron_bridge/README.md
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
.agents/skills/qad/SKILL.md (1)
45-53: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winConstrain ETP to the supported value of 1.
examples/megatron_bridge/distill.pyhard-setsprovider.expert_tensor_parallel_size = 1because expert tensor parallelism is unsupported. This skill currently lets operators derive/select ETP and uses it in the EDP calculation; choosingETP > 1will therefore be silently overridden, making the documented topology and sizing checks inconsistent with the actual launch.Proposed clarification
- derive the smallest fitting node count and TP/PP/CP/EP/ETP from + derive the smallest fitting node count and TP/PP/CP/EP from student and teacher architecture, 32K activation memory, and available GPU memory. ... - EDP = world_size / (ETP * EP * PP) + ETP = 1; EDP = world_size / (EP * PP)🤖 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 45 - 53, Update the topology guidance in “Choose topology explicitly” to require ETP = 1, remove ETP as a selectable/derived value, and simplify the EDP calculation to use the fixed supported value. Keep the existing integrality and expert-count checks consistent with this constraint, matching the hard-coded setting in distill.py.
🤖 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.
Outside diff comments:
In @.agents/skills/qad/SKILL.md:
- Around line 45-53: Update the topology guidance in “Choose topology
explicitly” to require ETP = 1, remove ETP as a selectable/derived value, and
simplify the EDP calculation to use the fixed supported value. Keep the existing
integrality and expert-count checks consistent with this constraint, matching
the hard-coded setting in distill.py.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: bfc1d9f2-5710-4e18-911d-568d0da8ecec
📒 Files selected for processing (1)
.agents/skills/qad/SKILL.md
Signed-off-by: Meng Xin <mxin@nvidia.com>
Signed-off-by: Meng Xin <mxin@nvidia.com>
4 participants
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
nemotron-cascade-2dataset blend configuration with an increased token budget.