Offline-KD QAD example - #1998
Conversation
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>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds a five-task Slurm pipeline for Nemotron NVFP4 quantization, offline knowledge distillation, checkpoint export, and TensorRT-LLM evaluation. Also updates three 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 StudentSFT
participant HFExport
participant TensorRTLLM
Slurm->>TeacherQuantization: Create NVFP4 teacher checkpoint
TeacherQuantization->>TeacherSFT: Provide quantized checkpoint
TeacherSFT->>StudentSFT: Write and load top-K 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://cdn.github.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 77.79% 76.14% -1.66%
==========================================
Files 518 519 +1
Lines 58204 62054 +3850
==========================================
+ Hits 45279 47249 +1970
- Misses 12925 14805 +1880
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: Asha Anoosheh <aanoosheh@nvidia.com>
e6957fb to
97d2ccc
Compare
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>
|
/claude review |
![claude[bot]](https://cdn.github.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.
2 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 development execution.