From 49fc88b1cf1c140f821d362d10e48a6996536f5c Mon Sep 17 00:00:00 2001 From: James Shen Date: Tue, 11 Aug 2026 05:31:10 -0700 Subject: [PATCH 1/6] example(launcher): Megatron-Bridge NVFP4 QAD for Nemotron-3-Nano-30B-A3B 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 --- .../mbridge_qad.yaml | 129 ++++++++++++++++++ 1 file changed, 129 insertions(+) create mode 100644 tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_qad.yaml diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_qad.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_qad.yaml new file mode 100644 index 00000000000..83769a35b6f --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_qad.yaml @@ -0,0 +1,129 @@ +# NVIDIA Nemotron 3 Nano 30B-A3B NVFP4 quantization-aware distillation (QAD) via Megatron-Bridge. +# +# Four tasks: tokenize the training data, PTQ the student to NVFP4, distill it against the BF16 +# teacher, and export a deployable unified-HF checkpoint. +# +# Training topology: 8 nodes x 4 GPUs, TP=1, PP=1, CP=4, EP=16. That leaves DP=8, so a +# global-batch-size of 64 at micro-batch-size 1 is 8 gradient-accumulation microbatches per step. +# 200 iterations x 64 sequences x 32768 tokens = 419M training tokens. +# +# Requirements: +# - HF_TOKEN can access the gated nvidia/Nemotron-Post-Training-Dataset-v2 dataset. +# +# Usage from tools/launcher: +# source .env-slurm +# uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_qad.yaml --yes + +job_name: Nemotron-3-Nano-30B-A3B_mbridge_qad_32k_200iter +pipeline: + note: "NVFP4 QAD at 32K for 200 iterations on Nemotron-Post-Training-Dataset-v2 chat (Megatron-Bridge)" + + global_vars: + hf_model: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + data_dir: /cicd/tokenized/nemotron-post-training-v2 + data_prefix: /cicd/tokenized/nemotron-post-training-v2/nvidia--Nemotron-Post-Training-Dataset-v2_default_chat_messages + ptq_ckpt: /cicd/megatron-bridge/Nemotron-3-Nano-30B-A3B-NVFP4-ptq + qad_dir: /cicd/megatron-bridge/Nemotron-3-Nano-30B-A3B-NVFP4-qad + export_dir: /cicd/export/Nemotron-3-Nano-30B-A3B-NVFP4-qad-hf + + # 1) Tokenize the QAD training data into Megatron .bin/.idx, which distill.py reads via + # --data_paths. megatron_lm_qad.yaml points Megatron-LM's finetune path at a single parquet + # shard instead; Megatron-Bridge trains from pre-tokenized data, so the split is tokenized + # once here. --hf_streaming avoids the Arrow cast errors that this dataset's nested tool-call + # fields trigger in non-streaming mode. No --append_eod: these are chat rows ("messages"), + # whose chat template already terminates each conversation. + # 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 + --hf_dataset nvidia/Nemotron-Post-Training-Dataset-v2 + --hf_name default + --hf_split chat + --hf_streaming + --json_keys messages + --tokenizer <> + --output_dir <> + --workers 32 + --max_sequence_length 256_000 + slurm_config: + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.06 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + time: "08:00:00" + + # 2) NVFP4 PTQ. Produces the quantized Megatron checkpoint that seeds the QAD student. + # TP=EP=PP=1 leaves pure DP=4, so each rank calibrates on its own shard of the samples. + # --calib_dataset_name is left unset, which selects the default public text mix. + task_1: + environment: + - LAUNCH_SCRIPT: torchrun --nproc_per_node 4 + inline: >- + $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/quantize.py + --hf_model_name_or_path <> + --trust_remote_code + --tp_size 1 + --pp_size 1 + --ep_size 1 + --quant_cfg MAMBA_MOE_NVFP4_CONSERVATIVE_CFG + --calib_batch_size 1 + --calib_num_samples 1000 + --seq_length 32768 + --skip_generate + --export_megatron_path <> + slurm_config: &sc + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.06 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + nodes: 1 + ntasks_per_node: 4 + gpus_per_node: 4 + + # 3) Distill the NVFP4 student from the BF16 teacher on the tokenized chat data. + task_2: + environment: + - LAUNCH_SCRIPT: torchrun --nproc_per_node 4 + inline: >- + $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/distill.py + --teacher_hf_path <> + --student_hf_path <> + --student_megatron_path <> + --trust_remote_code + --tp_size 1 + --pp_size 1 + --cp_size 4 + --ep_size 16 + --data_paths <> + --data_path_to_cache <>/cache + --seq_length 32768 + --mbs 1 + --gbs 64 + --lr 2e-5 + --min_lr 5e-6 + --lr_warmup_iters 30 + --train_iters 200 + --eval_interval 50 + --eval_iters 8 + --log_interval 10 + --checkpoint_keep_last 2 + --output_dir <> + slurm_config: + <<: *sc + nodes: 8 + + # 4) Export the distilled (still quantized) checkpoint to a deployable unified-HF checkpoint. + # TP must be 1 -- the HF writer does not gather TP shards -- and PP=4 splits 52 layers 13/stage. + task_3: + environment: + - LAUNCH_SCRIPT: torchrun --nproc_per_node 4 + inline: >- + $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/export_quantized_megatron_to_hf.py + --hf_model_name_or_path <> + --megatron_path <>/checkpoints + --trust_remote_code + --pp_size 4 + --export_unified_hf_path <> + slurm_config: + <<: *sc From 031682b83718fec4d2d29b5b51f0301c8802f2fd Mon Sep 17 00:00:00 2001 From: James Shen Date: Tue, 11 Aug 2026 08:28:26 -0700 Subject: [PATCH 2/6] example(launcher): Megatron-Bridge NVFP4 QAD for Nemotron-3.5-Lightning-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 --- .../mbridge_qad.yaml | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) rename tools/launcher/examples/nvidia/{NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 => NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16}/mbridge_qad.yaml (87%) diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_qad.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml similarity index 87% rename from tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_qad.yaml rename to tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml index 83769a35b6f..47a4b4affe4 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_qad.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml @@ -1,4 +1,4 @@ -# NVIDIA Nemotron 3 Nano 30B-A3B NVFP4 quantization-aware distillation (QAD) via Megatron-Bridge. +# NVIDIA Nemotron 3.5 Lightning 30B-A3B NVFP4 quantization-aware distillation (QAD) via Megatron-Bridge. # # Four tasks: tokenize the training data, PTQ the student to NVFP4, distill it against the BF16 # teacher, and export a deployable unified-HF checkpoint. @@ -12,19 +12,19 @@ # # Usage from tools/launcher: # source .env-slurm -# uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_qad.yaml --yes +# uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml --yes -job_name: Nemotron-3-Nano-30B-A3B_mbridge_qad_32k_200iter +job_name: Nemotron-3.5-Lightning-30B-A3B_mbridge_qad_32k_200iter pipeline: note: "NVFP4 QAD at 32K for 200 iterations on Nemotron-Post-Training-Dataset-v2 chat (Megatron-Bridge)" global_vars: - hf_model: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + hf_model: nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 data_dir: /cicd/tokenized/nemotron-post-training-v2 data_prefix: /cicd/tokenized/nemotron-post-training-v2/nvidia--Nemotron-Post-Training-Dataset-v2_default_chat_messages - ptq_ckpt: /cicd/megatron-bridge/Nemotron-3-Nano-30B-A3B-NVFP4-ptq - qad_dir: /cicd/megatron-bridge/Nemotron-3-Nano-30B-A3B-NVFP4-qad - export_dir: /cicd/export/Nemotron-3-Nano-30B-A3B-NVFP4-qad-hf + ptq_ckpt: /cicd/megatron-bridge/Nemotron-3.5-Lightning-30B-A3B-NVFP4-ptq + qad_dir: /cicd/megatron-bridge/Nemotron-3.5-Lightning-30B-A3B-NVFP4-qad + export_dir: /cicd/export/Nemotron-3.5-Lightning-30B-A3B-NVFP4-qad-hf # 1) Tokenize the QAD training data into Megatron .bin/.idx, which distill.py reads via # --data_paths. megatron_lm_qad.yaml points Megatron-LM's finetune path at a single parquet @@ -67,7 +67,7 @@ pipeline: --tp_size 1 --pp_size 1 --ep_size 1 - --quant_cfg MAMBA_MOE_NVFP4_CONSERVATIVE_CFG + --recipe huggingface/models/nvidia/Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/w4a16_nvfp4_4o6 --calib_batch_size 1 --calib_num_samples 1000 --seq_length 32768 From 73aedfcc7eb29d14101699bb349cb3fe99e2fb1a Mon Sep 17 00:00:00 2001 From: James Shen Date: Wed, 12 Aug 2026 01:18:12 -0700 Subject: [PATCH 3/6] fix(launcher): use only declared global_vars keys in mbridge_qad.yaml 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 --- .../mbridge_qad.yaml | 22 ++++++++----------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml index 47a4b4affe4..980bd7953fb 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml @@ -20,11 +20,7 @@ pipeline: global_vars: hf_model: nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 - data_dir: /cicd/tokenized/nemotron-post-training-v2 - data_prefix: /cicd/tokenized/nemotron-post-training-v2/nvidia--Nemotron-Post-Training-Dataset-v2_default_chat_messages - ptq_ckpt: /cicd/megatron-bridge/Nemotron-3.5-Lightning-30B-A3B-NVFP4-ptq - qad_dir: /cicd/megatron-bridge/Nemotron-3.5-Lightning-30B-A3B-NVFP4-qad - export_dir: /cicd/export/Nemotron-3.5-Lightning-30B-A3B-NVFP4-qad-hf + output_dir: /cicd/megatron-bridge/Nemotron-3.5-Lightning-30B-A3B-NVFP4 # 1) Tokenize the QAD training data into Megatron .bin/.idx, which distill.py reads via # --data_paths. megatron_lm_qad.yaml points Megatron-LM's finetune path at a single parquet @@ -42,7 +38,7 @@ pipeline: --hf_streaming --json_keys messages --tokenizer <> - --output_dir <> + --output_dir /cicd/tokenized/nemotron-post-training-v2 --workers 32 --max_sequence_length 256_000 slurm_config: @@ -72,7 +68,7 @@ pipeline: --calib_num_samples 1000 --seq_length 32768 --skip_generate - --export_megatron_path <> + --export_megatron_path <>-ptq slurm_config: &sc _factory_: "slurm_factory" container: nvcr.io/nvidia/nemo:26.06 @@ -89,14 +85,14 @@ pipeline: $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/distill.py --teacher_hf_path <> --student_hf_path <> - --student_megatron_path <> + --student_megatron_path <>-ptq --trust_remote_code --tp_size 1 --pp_size 1 --cp_size 4 --ep_size 16 - --data_paths <> - --data_path_to_cache <>/cache + --data_paths /cicd/tokenized/nemotron-post-training-v2/nvidia--Nemotron-Post-Training-Dataset-v2_default_chat_messages + --data_path_to_cache /cicd/tokenized/nemotron-post-training-v2/cache --seq_length 32768 --mbs 1 --gbs 64 @@ -108,7 +104,7 @@ pipeline: --eval_iters 8 --log_interval 10 --checkpoint_keep_last 2 - --output_dir <> + --output_dir <>-qad slurm_config: <<: *sc nodes: 8 @@ -121,9 +117,9 @@ pipeline: inline: >- $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/export_quantized_megatron_to_hf.py --hf_model_name_or_path <> - --megatron_path <>/checkpoints + --megatron_path <>-qad/checkpoints --trust_remote_code --pp_size 4 - --export_unified_hf_path <> + --export_unified_hf_path <>-qad-hf slurm_config: <<: *sc From 3e8d0e524eaba8b4ff2cf4fdd584055104e645fe Mon Sep 17 00:00:00 2001 From: James Shen Date: Wed, 12 Aug 2026 04:41:54 -0700 Subject: [PATCH 4/6] fix(launcher): cap the tokenize task at the 4h partition limit 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 --- .../NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml index 980bd7953fb..035a7855849 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml @@ -48,7 +48,7 @@ pipeline: nodes: 1 ntasks_per_node: 1 gpus_per_node: 1 - time: "08:00:00" + time: "04:00:00" # 2) NVFP4 PTQ. Produces the quantized Megatron checkpoint that seeds the QAD student. # TP=EP=PP=1 leaves pure DP=4, so each rank calibrates on its own shard of the samples. From 7f838c42a2e4a828c455cd6e31f705432bb8b4a1 Mon Sep 17 00:00:00 2001 From: James Shen Date: Wed, 12 Aug 2026 04:44:07 -0700 Subject: [PATCH 5/6] fix(launcher): request a full node for the tokenize task 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 --- .../mbridge_qad.yaml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml index 035a7855849..a87c23cdb84 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml @@ -46,8 +46,10 @@ pipeline: container: nvcr.io/nvidia/nemo:26.06 modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt nodes: 1 + # One process (the tokenizer is single-process), but a full node: clusters + # commonly enforce a minimum GPU count per job (QOSMinGRES). ntasks_per_node: 1 - gpus_per_node: 1 + gpus_per_node: 4 time: "04:00:00" # 2) NVFP4 PTQ. Produces the quantized Megatron checkpoint that seeds the QAD student. From cc3a8c0740e0e526c55191752cedaa0d55718883 Mon Sep 17 00:00:00 2001 From: James Shen Date: Wed, 12 Aug 2026 17:05:24 -0700 Subject: [PATCH 6/6] fix(launcher): keep the QAD Triton cache node-local 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 --- .../NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml index a87c23cdb84..e14e39a185f 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_qad.yaml @@ -83,6 +83,7 @@ pipeline: task_2: environment: - LAUNCH_SCRIPT: torchrun --nproc_per_node 4 + - TRITON_CACHE_DIR: /tmp/triton_cache inline: >- $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/distill.py --teacher_hf_path <>