example(launcher): Megatron-Bridge NVFP4 QAD launcher example for Nemotron-3.5-Lightning-30B-A3B - #2142
Merged
Merged
Conversation
Contributor
|
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:
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)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds a Megatron-Bridge launcher pipeline for NVIDIA Nemotron 3.5 Lightning 30B-A3B. The pipeline preprocesses data, applies NVFP4 PTQ, performs 200-step QAD distillation, and exports a unified Hugging Face checkpoint. ChangesNemotron QAD pipeline
Estimated code review effort: 2 (Simple) | ~15 minutes Sequence Diagram(s)sequenceDiagram
participant PostTrainingDataset
participant PreprocessingTask
participant NVFP4PTQTask
participant QADDistillationTask
participant HuggingFaceExportTask
PostTrainingDataset->>PreprocessingTask: chat messages
PreprocessingTask->>NVFP4PTQTask: tokenized Megatron data
NVFP4PTQTask->>QADDistillationTask: NVFP4 student checkpoint
QADDistillationTask->>HuggingFaceExportTask: distilled Megatron checkpoint
HuggingFaceExportTask->>HuggingFaceExportTask: write unified Hugging Face checkpoint
Possibly related PRs
Suggested reviewers: 🚥 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 #2142 +/- ##
==========================================
- Coverage 78.74% 77.73% -1.02%
==========================================
Files 522 522
Lines 60368 61469 +1101
==========================================
+ Hits 47538 47781 +243
- Misses 12830 13688 +858
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:
|
yueshen2016
force-pushed
the
feat/mbridge-qad-launcher-example
branch
from
August 11, 2026 12:34
dc3fbba to
cbb4453
Compare
yueshen2016
changed the base branch from
feat/mbridge-sft-distillation
to
main
August 11, 2026 12:34
yueshen2016
force-pushed
the
feat/mbridge-qad-launcher-example
branch
from
August 11, 2026 12:41
cbb4453 to
98fb149
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
| # CPU-bound and long-running; it needs no GPU beyond the allocation minimum. | ||
| task_0: | ||
| inline: >- | ||
| python -m modelopt.torch.utils.plugins.megatron_preprocess_data |
Collaborator
There was a problem hiding this comment.
Contributor
Author
There was a problem hiding this comment.
It's an example that will be posted on tech blog as the M-Bridge counterpart of #2146. The settings will stay the same as the megatron_lm_qad.
| <<: *sc | ||
| nodes: 8 | ||
|
|
||
| # 4) Export the distilled (still quantized) checkpoint to a deployable unified-HF checkpoint. |
Collaborator
There was a problem hiding this comment.
should we add a lm_eval_hf.py MMLU check as well?
Contributor
Author
There was a problem hiding this comment.
Same as prior response, will keep the same as megatron_lm_qad.
kevalmorabia97
approved these changes
Aug 11, 2026
yueshen2016
changed the title
Adds mbridge_qad.yaml next to the existing mbridge_prune / mbridge_quantize examples: tokenize the Nemotron-Post-Training-Dataset-v2 chat split, PTQ to NVFP4, distill against the BF16 teacher, then export to unified HF. Mirrors megatron_lm_qad.yaml's recipe and data through the Megatron-Bridge scripts instead of Megatron-LM. Megatron-LM's finetune path reads an HF parquet shard directly; Megatron-Bridge trains from pre-tokenized data, so the split is tokenized once with megatron_preprocess_data and passed via --data_paths. Signed-off-by: James Shen <yueshen@nvidia.com>
…ng-30B-A3B Adds mbridge_qad.yaml: tokenize the Nemotron-Post-Training-Dataset-v2 chat split, PTQ to NVFP4, distill against the BF16 teacher, then export to unified HF. The Megatron-LM counterpart is megatron_lm_qad.yaml. Megatron-LM's finetune path reads an HF parquet shard directly; Megatron-Bridge trains from pre-tokenized data, so the split is tokenized once with megatron_preprocess_data and passed via --data_paths. PTQ uses the same recipe as the Megatron-LM example, huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6. Signed-off-by: James Shen <yueshen@nvidia.com>
Contributor
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)
tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml (1)
87-89: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftConfigure the multi-node
torchrunrendezvous.
task_2runs the inline command on all 32 Slurm tasks. Each task starts fourtorchrunworkers, producing 128 workers instead of 32. Pass--nnodes, a Slurm-derived--node_rank, and a shared rendezvous endpoint. Setntasks_per_node: 1so onetorchrunprocess manages the four GPUs on each node.🤖 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.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml` around lines 87 - 89, Update the task_2 launcher configuration around LAUNCH_SCRIPT to set ntasks_per_node to 1 and configure torchrun for multi-node execution with --nnodes, a Slurm-derived --node_rank, and a shared rendezvous endpoint; ensure one torchrun process manages all four GPUs on each node.
🤖 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
`@tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml`:
- Around line 87-89: Update the task_2 launcher configuration around
LAUNCH_SCRIPT to set ntasks_per_node to 1 and configure torchrun for multi-node
execution with --nnodes, a Slurm-derived --node_rank, and a shared rendezvous
endpoint; ensure one torchrun process manages all four GPUs on each node.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 71ac454b-bb87-4d6f-b7d0-212cfe5e000f
📒 Files selected for processing (1)
tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml
yueshen2016
force-pushed
the
feat/mbridge-qad-launcher-example
branch
from
August 12, 2026 08:04
ba7f364 to
031682b
Compare
GlobalVariables is a fixed dataclass (hf_model, hf_data, hf_local, output_dir, draft_model), so data_dir / data_prefix / ptq_ckpt / qad_dir / export_dir were rejected at launch with 'No parameter named ...'. Keep hf_model and output_dir and inline the rest. Found by running the pipeline: the validate-launcher-YAML hook checks recipe and template paths, not global_vars names. Signed-off-by: James Shen <yueshen@nvidia.com>
The tokenize task requested time: 08:00:00, which sbatch rejects on clusters
whose batch partition caps at 4h ('Requested time limit is invalid'). Use
04:00:00, matching the launcher's slurm_factory default.
Signed-off-by: James Shen <yueshen@nvidia.com>
A 1-GPU request is rejected on clusters that enforce a minimum GPU count per job (QOSMinGRES -> 'Job violates accounting/QOS policy'). Ask for a full node while keeping ntasks_per_node: 1, since the tokenizer is a single process and multiple tasks would race writing the same output directory. Signed-off-by: James Shen <yueshen@nvidia.com>
task_2 is the only multi-node task (8 nodes, 32 ranks). The launcher defaults
TRITON_CACHE_DIR under the shared /{title} mount, which is on Lustre; ranks race
to write and mmap the same compiled kernels there, and a stalled rank leaves the
other 31 waiting at the next expert-parallel allgather until the 600s NCCL
watchdog aborts the job ~17 minutes in, before iteration 1. Observed as
OSError: [Errno 14] Bad address on /cicd/triton-cache. With a node-local cache
the same run completes all 200 iterations.
Signed-off-by: James Shen <yueshen@nvidia.com>
yueshen2016
force-pushed
the
feat/mbridge-qad-launcher-example
branch
from
August 13, 2026 00:05
c4cf4b6 to
cc3a8c0
Compare
yueshen2016
enabled auto-merge (squash)
August 13, 2026 00:06
Contributor
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Adds
mbridge_qad.yaml, a launcher example running NVFP4 quantization-awaredistillation for Nemotron-3.5-Lightning-30B-A3B through the Megatron-Bridge
scripts in
examples/megatron_bridge/, alongside the existingmbridge_prune.yaml/mbridge_quantize.yaml.megatron_lm_qad.yaml(#2146) runs the same recipe and the same data throughMegatron-LM. This is the Megatron-Bridge counterpart: the same
huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6recipe and the same
nvidia/Nemotron-Post-Training-Dataset-v2chat data, with thetraining hyperparameters from our public-data QAD run.
Four tasks: tokenize the training data, PTQ the student, distill it against the
frozen BF16 teacher, export to unified HF.
Why the extra tokenize task
Megatron-LM's finetune path reads an HF parquet shard directly. Megatron-Bridge
trains from pre-tokenized data, so
distill.pyconsumes Megatron.bin/.idxvia
--data_paths. The chat split is therefore tokenized once withmodelopt.torch.utils.plugins.megatron_preprocess_data.--hf_streamingavoidsthe Arrow cast errors this dataset's nested tool-call fields trigger in
non-streaming mode, and
--append_eodis omitted because chat rows alreadyterminate each conversation via the chat template.
Details
gbs64 atmbs1 is 8 gradient-accumulation microbatches. 200 iters x 64 x 32768 = 419Mtraining tokens.
shard.
--calib_dataset_nameis left unset, selecting the default publiccnn_nemotron_v2_mix(cnn_dailymail + Nemotron-Post-Training-Dataset-v2).52 layers 13/stage.
nvcr.io/nvidia/nemo:26.06like the othermbridge_*examples.Dependencies
Based on
main; the PTQ recipe ships in #2146 (merged). No other PR required.Nemotron-3.5-Lightning has
tie_word_embeddings: false, so a correct quantizedlm_headin the exported checkpoint also depends on #2112.Testing
The PTQ -> export -> QAD flow and these hyperparameters were run end to end on
Nemotron-3.5-Lightning (
main+ #2112 + #2113):mtp_num_layers: 1) withall 278
mtp.*quantizers disabled by the recipe.The YAML itself is validated against the launcher's conventions
(
ntasks_per_node == gpus_per_nodeon Slurm, single-lineinline, noargsalongside
inline, all<<global_vars.X>>resolve, output prefix matchesmegatron_preprocess_data's naming) and by the repo'svalidate launcher YAML referencespre-commit hook. Topology arithmetic checked: EP dividesworld/(TPPP),
gbsdivisible by DPmbs.Before your PR is "Ready for review"
🤖 Generated with Claude Code
Summary by CodeRabbit