Offline-KD QAD example - #1998
Conversation
|
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 a five-task Slurm pipeline for Nemotron NVFP4 quantization, offline knowledge distillation, checkpoint export, and TensorRT-LLM evaluation. Also updates four local pre-commit hooks to run through the frozen uv development environment. ChangesOffline KD QAD workflow
uv-managed pre-commit hooks
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Slurm
participant TeacherQuantization
participant TeacherSFT
participant LogitCache
participant StudentSFT
participant HFExport
participant TensorRTLLM
Slurm->>TeacherQuantization: Create NVFP4 teacher checkpoint
TeacherQuantization->>TeacherSFT: Provide quantized checkpoint
TeacherSFT->>LogitCache: Save top-K logits
LogitCache->>StudentSFT: Load cached logits
StudentSFT->>HFExport: Provide trained NVFP4 checkpoint
HFExport->>TensorRTLLM: Evaluate exported HF checkpoint
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml (1)
44-50: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider deduplicating shared slurm_config values.
The
containerandmodelopt_install_pathvalues are identical across tasks 0-3. Per coding guidelines,<<global_vars.X>>interpolation is available for shared values. If global_vars supports partial slurm_config overrides, this would reduce four repetitions to a single definition.As per coding guidelines: "
<<global_vars.X>>interpolation for shared values".Also applies to: 100-106, 153-159, 169-175
🤖 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 `@tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml` around lines 44 - 50, Deduplicate the repeated container and modelopt_install_path values in the slurm_config blocks for tasks 0–3 by defining them once in global_vars and referencing them through <<global_vars.X>> interpolation. Update the corresponding blocks near the referenced task sections while preserving their existing task-specific Slurm settings.Source: Coding guidelines
🤖 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.
Nitpick comments:
In
`@tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml`:
- Around line 44-50: Deduplicate the repeated container and
modelopt_install_path values in the slurm_config blocks for tasks 0–3 by
defining them once in global_vars and referencing them through <<global_vars.X>>
interpolation. Update the corresponding blocks near the referenced task sections
while preserving their existing task-specific Slurm settings.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3413039f-4a7c-4124-88d0-d4897f9b1eaf
📒 Files selected for processing (1)
tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1998 +/- ##
==========================================
+ Coverage 60.23% 60.28% +0.05%
==========================================
Files 519 519
Lines 59399 59399
==========================================
+ Hits 35778 35808 +30
+ Misses 23621 23591 -30
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:
|
e6957fb to
97d2ccc
Compare
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| - EP: 4 | ||
| - ETP: 1 | ||
| - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 | ||
| - MLM_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16-MCore |
There was a problem hiding this comment.
[SUGGESTION] task_1's teacher loads MLM_MODEL_CKPT: .../NVIDIA-Nemotron-3-Nano-30B-A3B-BF16-MCore, but this MCore checkpoint is never produced by any task in this pipeline (task_0 produces the NVFP4 ckpt under /cicd/megatron-lm/, not this BF16 MCore). It's an implicit external prerequisite — the user must run the HF→MCore import (e.g. common/megatron_bridge/import/import.sh) beforehand, or task_1 fails at load with a missing-path error.
The header "Key constraints" block is otherwise thorough; consider adding a one-line prerequisite note there, e.g.:
# - PREREQUISITE: task_1's teacher loads a BF16 MCore checkpoint at
# MLM_MODEL_CKPT (NVIDIA-Nemotron-3-Nano-30B-A3B-BF16-MCore). Import it from
# HF first (common/megatron_bridge/import/import.sh); this pipeline does not
# create it.Not blocking — just saves the next user a confusing first-run failure.
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1
Scope: reviewed both changed files — the new offline_kd_qad.yaml (209 lines) and the .pre-commit-config.yaml hook updates. Traced the checkpoint dataflow across all five tasks against the launcher scripts (quantize.sh, sft.sh, export.sh, eval.sh) and slurm_config.py / core.py.
Assessment: example/config-only PR (no library code), so the risk surface is a broken example rather than a regression. The pipeline is internally consistent:
- Checkpoint chaining is coherent: task_0 saves NVFP4 to the default /cicd/megatron-lm/${MLM_MODEL_CFG}; task_2 (no explicit MLM_MODEL_CKPT) loads/overwrites that same default; task_3 exports it to /scratchspace/export; task_4 evals that directory.
- Teacher/student alignment holds: CP=8, --seq-length 32768, and --global-batch-size 8 match between task_1 and task_2, which the cached-logit per-(iteration, microbatch) indexing requires.
- TE-spec discipline is correct: --export-default-te-spec present on task_0/2/3 and correctly absent on the BF16 teacher (task_1); --sequence-parallel appears only on the TP=2 task.
- Parallelism x node counts check out (task_1 CP8xTP1=8 GPUs/2 nodes; task_2 CP8xTP2=16 GPUs/4 nodes).
- task_4 omitting container correctly falls back to slurm_factory's default TRT-LLM image for trtllm-eval — intentional.
- Pre-commit 'uv run --frozen --extra dev' change is applied consistently across the three local Python hooks.
Most impactful finding (non-blocking SUGGESTION): task_1's BF16 MCore teacher checkpoint (...-BF16-MCore) is an implicit external prerequisite — it is not produced by any task in this pipeline and must be imported from HF beforehand. A one-line note in the header would save the next user a confusing first-run load failure.
Risk level: low.
There was a problem hiding this comment.
what do these changes do? Is it supposed to make it faster?
There was a problem hiding this comment.
I got a new machine and decided during setup I will avoid system python and just use uv for project dependencies. That way I don't need to worry about installing the pre-commit deps like pyyaml into something outside of this project's uv.
Adds a 5-task offline knowledge distillation pipeline for NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4: task_0: PTQ quantize BF16 → NVFP4 task_1: teacher dump (freeze, save top-16 logits) task_2: student training with --logits-load-dir task_3: export task_4: TRT-LLM eval Also adds common/megatron_lm/train/sft.sh for SFT/QAD training jobs. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
8f8ccf7 to
2c1ff0e
Compare
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)
tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml (1)
48-61: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRepeated shared config values should use
<<global_vars.X>>interpolation.
container: nvcr.io/nvidia/nemo:26.06,modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt, andMLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16are duplicated verbatim across task_0 (48-61), task_1 (105-121), task_2 (161-176), and task_3 (183-199). Cross-task path strings like/scratchspace/logits(102, 160) and/scratchspace/export(187, 204) are similarly duplicated. A single typo in any repeated copy would silently break the pipeline dependency chain.As per coding guidelines,
tools/launcher/**/*.yamlfiles use`<<global_vars.X>>` interpolation for shared values; consider defining these once underglobal_varsand referencing them.🤖 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 `@tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml` around lines 48 - 61, Define the repeated container, modelopt installation path, MLM model configuration, and shared scratchspace paths once under global_vars in offline_kd_qad.yaml. Replace the duplicate values in task_0, task_1, task_2, and task_3, including slurm_config and dependency path fields, with the corresponding <<global_vars.X>> interpolations while preserving each task’s existing behavior.Source: Coding guidelines
🤖 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
`@tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml`:
- Line 75: Update the data-cache-path values in task_1 and task_2 to use the
consistent nemotron-3-super-sft directory name, matching the calibration and
blend paths for the same corpus; leave the surrounding configuration unchanged.
---
Nitpick comments:
In
`@tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml`:
- Around line 48-61: Define the repeated container, modelopt installation path,
MLM model configuration, and shared scratchspace paths once under global_vars in
offline_kd_qad.yaml. Replace the duplicate values in task_0, task_1, task_2, and
task_3, including slurm_config and dependency path fields, with the
corresponding <<global_vars.X>> interpolations while preserving each task’s
existing 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: 28132d90-7b6d-420c-9e35-87b3929a8aed
📒 Files selected for processing (2)
.pre-commit-config.yamltools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml
|
So 2 args to save and 1 args to enable offline kd? |
| gpus_per_node: 4 | ||
|
|
||
| task_1: | ||
| script: common/megatron_lm/train/sft.sh |
There was a problem hiding this comment.
Ideally we use the shell script in megatron-lm directly. I think it is possible but require refactorization of the script.
| - --log-interval 100 | ||
| - --finetune | ||
| - --freeze-all-layers | ||
| - --logits-save-top-k 16 |
There was a problem hiding this comment.
So these 2 args to save the logits?
There was a problem hiding this comment.
Yes. (Soon technically it will be only 1 arg since in next update I'll make default k=128)
- data-cache-path: test-nemotron-sft -> nemotron-3-super-sft (CodeRabbit), matching the dataset/blend/calib paths (task_1 + task_2) - header: note task_1's BF16 MCore teacher ckpt is an external prerequisite, not produced by this pipeline (Claude review) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
3 participants
What does this PR do?
Type of change: new feature
Adds a Nano-v3 launcher example which uses the new MLM offline-logits KD feature
Usage
# Add a code snippet demonstrating how to use thisTesting
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
Summary by CodeRabbit
uvfor consistent dev-time tooling.