Offline-KD QAD example

[AAnoosheh]

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 this

Testing

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.).

• Is this change backward compatible?: N/A
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: N/A
• Did you update Changelog?: ❌
• Did you get Claude approval on this PR?: ❌

Additional Information

Summary by CodeRabbit

• New Features
  • Added an end-to-end offline knowledge distillation workflow for Nemotron-3-Nano NVFP4 models.
  • Covers teacher quantization, teacher and student fine-tuning, asynchronous top‑K logits sharing, checkpoint export to Hugging Face format, and TensorRT‑LLM evaluation.
  • Provides predefined Slurm execution settings and hardware parallelism options for repeatable runs.
• Chores
  • Updated local pre-commit Python hooks to run via uv for consistent dev execution.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a873f3b1-6f8c-40b0-bb39-7a8e74749596

📥 Commits

Reviewing files that changed from the base of the PR and between e6957fb and 97d2ccc.

📒 Files selected for processing (2)

• .pre-commit-config.yaml
• tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml

🚧 Files skipped from review as they are similar to previous changes (1)

• tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml

---

📝 Walkthrough

Walkthrough

Adds 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.

Changes

Offline KD QAD workflow

Layer / File(s)	Summary
Define staged KD and evaluation workflow tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml	Adds teacher quantization, asynchronous top-K logit caching, NVFP4 student SFT, HF export, and TensorRT-LLM evaluation tasks with pinned NeMo runtime settings, Slurm resources, parallelism, and scratch paths.

uv-managed pre-commit hooks

Layer / File(s)	Summary
Run local hooks through uv .pre-commit-config.yaml	Updates recipe checks, launcher YAML validation, and argument documentation generation to execute with uv run --frozen --extra dev.

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

Loading

Suggested reviewers: kevalmorabia97

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title captures the main change: adding an offline knowledge-distillation QAD example.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	No touched Python/dependency files; the diff contains none of torch.load(weights_only=False), allow_pickle=True, trust_remote_code=True, eval/exec, or # nosec.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch aanoosheh/offline-qad-example

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 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 value

Consider deduplicating shared slurm_config values.

The container and modelopt_install_path values 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

📥 Commits

Reviewing files that changed from the base of the PR and between 8ae7407 and e6957fb.

📒 Files selected for processing (1)

• tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.87%. Comparing base (9392dfe) to head (97d2ccc).
⚠️ Report is 4 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1998      +/-   ##
==========================================
+ Coverage   77.79%   77.87%   +0.07%     
==========================================
  Files         518      518              
  Lines       58204    58611     +407     
==========================================
+ Hits        45279    45642     +363     
- Misses      12925    12969      +44     

Flag	Coverage Δ
regression	15.14% <ø> (+0.07%)	⬆️
unit	54.86% <ø> (+0.06%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.