From 4e7af9c05db1f8eec9c1fb72951547a2f4d5a2b2 Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Fri, 19 Jun 2026 17:25:50 +0200 Subject: [PATCH 1/2] feat(launcher): add Nemotron-3-Nano-30B NVFP4 offline KD QAD example MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) Signed-off-by: Asha Anoosheh --- .../offline_kd_qad.yaml | 159 ++++++++++++++++++ 1 file changed, 159 insertions(+) create mode 100644 tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml new file mode 100644 index 00000000000..b5979d72658 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml @@ -0,0 +1,159 @@ +# Nemotron-3-Nano-30B-A3B NVFP4 Offline KD QAD pipeline. +# +# Two-stage knowledge distillation using cached teacher logits: +# task_0 (quantize): PTQ of BF16 teacher → NVFP4 MCore checkpoint. +# task_1 (dump): Run BF16 teacher frozen, save top-K logits to /scratchspace/logits. +# task_2 (student): Train NVFP4 student with --logits-load-dir. +# task_3 (export): Export trained student checkpoint to HF format. +# task_4 (eval): TRT-LLM evaluation of the exported checkpoint. +# +# Requirements: +# - Container must include the Megatron-LM offline-logits-kd branch +# (github.com/AAnoosheh/Megatron-LM @ aanoosheh/offline-logits-kd). +# Set PYTHONPATH=/nemo_run/code/modules/Megatron-LM to override the +# container's installed Megatron and use the submodule's version. +# - ModelOpt must include the MambaMixer conv1d compat fix (PR #1730). +# +# Usage: +# source .env-slurm +# cd tools/launcher +# uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml --yes + +job_name: NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4-QAD-OfflineKD +pipeline: + task_0: + script: common/megatron_lm/quantize/quantize.sh + args: + - --calib-dataset-path-or-name /hf-local/abisee/cnn_dailymail + - --calib-size 32 + - --calib-max-sequence-length 512 + environment: + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - QUANT_CFG: MAMBA_MOE_NVFP4_CONSERVATIVE_CFG + - EP: 4 + - MMLU_LOWER_BOUND: 0.0 + - PYTHONPATH: /nemo_run/code/modules/Megatron-LM + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 8 + gpus_per_node: 8 + + task_1: + script: common/megatron_lm/train/sft.sh + args: + - --seq-length 8192 --max-position-embeddings 8192 + - --train-samples 1024 + - --lr-decay-samples 256 + - --lr-warmup-samples 128 + - --sft + - --sft-tokenizer-prompt-format identity + - --tokenizer-type SFTTokenizer + - --no-create-attention-mask-in-dataloader + - --per-split-data-args-path /hf-local/modelopt/nemotron-3-super-sft/blend.json + - --data-cache-path /hf-local/modelopt/test-nemotron-sft/cache + - --micro-batch-size 1 + - --attention-dropout 0.0 + - --hidden-dropout 0.0 + - --no-check-for-nan-in-loss-and-grad + - --recompute-granularity selective + - --recompute-modules layernorm moe + - --lr 5.0e-5 + - --min-lr 5.0e-6 + - --lr-decay-style cosine + - --clip-grad 1.0 + - --weight-decay 0.0 + - --adam-beta1 0.9 + - --adam-beta2 0.95 + - --init-method-std 0.010 + - --use-distributed-optimizer + - --eval-iters 4 + - --eval-interval 1000 + - --save-interval 1000 + - --log-interval 100 + - --freeze-all-layers + - --logits-save-top-k 16 + - --logits-save-dir /scratchspace/logits + - --async-save + - --use-persistent-ckpt-worker + environment: + - DP: 1 + - CP: 2 + - TP: 1 + - PP: 1 + - EP: 8 + - 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 + - MLM_DATA_ARGS: --sft + - PYTHONPATH: /nemo_run/code/modules/Megatron-LM + slurm_config: + _factory_: "slurm_factory" + nodes: 2 + + task_2: + script: common/megatron_lm/train/sft.sh + args: + - --seq-length 8192 --max-position-embeddings 8192 + - --train-samples 1024 + - --lr-decay-samples 256 + - --lr-warmup-samples 128 + - --sft + - --sft-tokenizer-prompt-format identity + - --tokenizer-type SFTTokenizer + - --no-create-attention-mask-in-dataloader + - --per-split-data-args-path /hf-local/modelopt/nemotron-3-super-sft/blend.json + - --data-cache-path /hf-local/modelopt/test-nemotron-sft/cache + - --micro-batch-size 1 + - --global-batch-size 8 + - --attention-dropout 0.0 + - --hidden-dropout 0.0 + - --no-check-for-nan-in-loss-and-grad + - --recompute-granularity selective + - --recompute-modules layernorm moe + - --lr 5.0e-5 + - --min-lr 5.0e-6 + - --lr-decay-style cosine + - --clip-grad 1.0 + - --weight-decay 0.0 + - --adam-beta1 0.9 + - --adam-beta2 0.95 + - --init-method-std 0.010 + - --use-distributed-optimizer + - --eval-iters 4 + - --eval-interval 1000 + - --save-interval 1000 + - --log-interval 100 + - --sequence-parallel + - --logits-load-dir /scratchspace/logits + environment: + - DP: 1 + - CP: 2 + - TP: 2 + - PP: 1 + - EP: 8 + - ETP: 1 + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - MLM_DATA_ARGS: --sft + - PYTHONPATH: /nemo_run/code/modules/Megatron-LM + slurm_config: + _factory_: "slurm_factory" + nodes: 2 + + task_3: + script: common/megatron_lm/export/export.sh + environment: + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - PP: 4 + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 8 + gpus_per_node: 8 + + task_4: + script: common/tensorrt_llm/eval.sh + slurm_config: + _factory_: "slurm_factory" + nodes: 1 From 97d2ccc5a73d93651e413ca1c0592f0b78f52501 Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Mon, 20 Jul 2026 19:27:34 +0200 Subject: [PATCH 2/2] Use Nemo 26.06 container Signed-off-by: Asha Anoosheh --- .pre-commit-config.yaml | 8 +- .../offline_kd_qad.yaml | 228 ++++++++++-------- 2 files changed, 131 insertions(+), 105 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 43b869800d7..2b9bf162e85 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -57,13 +57,13 @@ repos: hooks: - id: normalize-yaml-ext name: normalize .yml to .yaml in required places, right now only yaml files in modelopt_recipes - entry: python tools/precommit/normalize_yaml_ext.py + entry: uv run --frozen --extra dev python tools/precommit/normalize_yaml_ext.py language: system files: ^modelopt_recipes/.*\.yml$ - id: check-modelopt-recipes name: validate modelopt recipes - entry: python tools/precommit/check_modelopt_recipes.py + entry: uv run --frozen --extra dev python tools/precommit/check_modelopt_recipes.py language: system files: ^modelopt_recipes/ # configs/ contains reusable snippets (not full recipes) — skip recipe validation @@ -78,7 +78,7 @@ repos: - id: check-launcher-yaml name: validate launcher YAML references to recipes and templates - entry: python tools/precommit/check_launcher_yaml.py + entry: uv run --frozen --extra dev python tools/precommit/check_launcher_yaml.py language: system files: ^(tools/launcher/examples/.*\.yaml|tools/precommit/check_launcher_yaml\.py)$ @@ -158,7 +158,7 @@ repos: hooks: - id: generate-arguments-md name: Regenerate examples/llm_qat/ARGUMENTS.md - entry: bash -c 'python examples/llm_qat/arguments.py --generate_docs examples/llm_qat/ARGUMENTS.md' + entry: uv run --frozen --extra dev python examples/llm_qat/arguments.py --generate_docs examples/llm_qat/ARGUMENTS.md language: system files: >- (?x)^( diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml index b5979d72658..6e90364d262 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml @@ -8,11 +8,19 @@ # task_4 (eval): TRT-LLM evaluation of the exported checkpoint. # # Requirements: -# - Container must include the Megatron-LM offline-logits-kd branch -# (github.com/AAnoosheh/Megatron-LM @ aanoosheh/offline-logits-kd). -# Set PYTHONPATH=/nemo_run/code/modules/Megatron-LM to override the -# container's installed Megatron and use the submodule's version. -# - ModelOpt must include the MambaMixer conv1d compat fix (PR #1730). +# - nvcr.io/nvidia/nemo:26.06 (ships nvidia-resiliency-ext>=0.6.0, needed +# by task_1's --async-save + --use-persistent-ckpt-worker). Tasks 0-3 pin +# this container explicitly since the launcher's default slurm_factory +# container is a TensorRT-LLM image with no Megatron-LM/NeMo deps. +# - GOTCHA: nemo containers install ModelOpt into a venv at /opt/venv, whose +# site-packages precede /usr/local on sys.path, so the default +# modelopt_install_path is never consulted. Point it at the venv +# site-packages so the mounted (submodule-pinned) modelopt actually +# overrides the container's. +# +# MMLU_LOWER_BOUND=0.0 on task_0 is a temporary bypass — the MMLU gate needs +# revisiting for this container/branch pairing before this pipeline is used +# for production runs (see OfflineKD MMLU accuracy note in project history). # # Usage: # source .env-slurm @@ -24,17 +32,19 @@ pipeline: task_0: script: common/megatron_lm/quantize/quantize.sh args: - - --calib-dataset-path-or-name /hf-local/abisee/cnn_dailymail - - --calib-size 32 - - --calib-max-sequence-length 512 + - --calib-dataset-path-or-name /hf-local/abisee/cnn_dailymail + - --calib-size 32 + - --calib-max-sequence-length 512 environment: - - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - - QUANT_CFG: MAMBA_MOE_NVFP4_CONSERVATIVE_CFG - - EP: 4 - - MMLU_LOWER_BOUND: 0.0 - - PYTHONPATH: /nemo_run/code/modules/Megatron-LM + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - QUANT_CFG: MAMBA_MOE_NVFP4_CONSERVATIVE_CFG + - EP: 4 + - MMLU_LOWER_BOUND: 0.0 + - PYTHONPATH: /nemo_run/code/modules/Megatron-LM 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: 8 gpus_per_node: 8 @@ -42,118 +52,134 @@ pipeline: task_1: script: common/megatron_lm/train/sft.sh args: - - --seq-length 8192 --max-position-embeddings 8192 - - --train-samples 1024 - - --lr-decay-samples 256 - - --lr-warmup-samples 128 - - --sft - - --sft-tokenizer-prompt-format identity - - --tokenizer-type SFTTokenizer - - --no-create-attention-mask-in-dataloader - - --per-split-data-args-path /hf-local/modelopt/nemotron-3-super-sft/blend.json - - --data-cache-path /hf-local/modelopt/test-nemotron-sft/cache - - --micro-batch-size 1 - - --attention-dropout 0.0 - - --hidden-dropout 0.0 - - --no-check-for-nan-in-loss-and-grad - - --recompute-granularity selective - - --recompute-modules layernorm moe - - --lr 5.0e-5 - - --min-lr 5.0e-6 - - --lr-decay-style cosine - - --clip-grad 1.0 - - --weight-decay 0.0 - - --adam-beta1 0.9 - - --adam-beta2 0.95 - - --init-method-std 0.010 - - --use-distributed-optimizer - - --eval-iters 4 - - --eval-interval 1000 - - --save-interval 1000 - - --log-interval 100 - - --freeze-all-layers - - --logits-save-top-k 16 - - --logits-save-dir /scratchspace/logits - - --async-save - - --use-persistent-ckpt-worker + - --seq-length 8192 --max-position-embeddings 8192 + - --train-samples 1024 + - --lr-decay-samples 256 + - --lr-warmup-samples 128 + - --sft + - --sft-tokenizer-prompt-format identity + - --tokenizer-type SFTTokenizer + - --no-create-attention-mask-in-dataloader + - --per-split-data-args-path /hf-local/modelopt/nemotron-3-super-sft/blend.json + - --data-cache-path /hf-local/modelopt/test-nemotron-sft/cache + - --micro-batch-size 1 + - --attention-dropout 0.0 + - --hidden-dropout 0.0 + - --no-check-for-nan-in-loss-and-grad + - --recompute-granularity selective + - --recompute-modules layernorm moe + - --lr 5.0e-5 + - --min-lr 5.0e-6 + - --lr-decay-style cosine + - --clip-grad 1.0 + - --weight-decay 0.0 + - --adam-beta1 0.9 + - --adam-beta2 0.95 + - --init-method-std 0.010 + - --use-distributed-optimizer + - --eval-iters 4 + - --eval-interval 1000 + - --save-interval 1000 + - --log-interval 100 + - --freeze-all-layers + - --logits-save-top-k 16 + - --logits-save-dir /scratchspace/logits + - --async-save + - --use-persistent-ckpt-worker environment: - - DP: 1 - - CP: 2 - - TP: 1 - - PP: 1 - - EP: 8 - - 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 - - MLM_DATA_ARGS: --sft - - PYTHONPATH: /nemo_run/code/modules/Megatron-LM + - DP: 1 + - CP: 2 + - TP: 1 + - PP: 1 + - EP: 8 + - 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 + - MLM_DATA_ARGS: --sft + - PYTHONPATH: /nemo_run/code/modules/Megatron-LM slurm_config: _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.06 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt nodes: 2 + ntasks_per_node: 8 + gpus_per_node: 8 task_2: script: common/megatron_lm/train/sft.sh args: - - --seq-length 8192 --max-position-embeddings 8192 - - --train-samples 1024 - - --lr-decay-samples 256 - - --lr-warmup-samples 128 - - --sft - - --sft-tokenizer-prompt-format identity - - --tokenizer-type SFTTokenizer - - --no-create-attention-mask-in-dataloader - - --per-split-data-args-path /hf-local/modelopt/nemotron-3-super-sft/blend.json - - --data-cache-path /hf-local/modelopt/test-nemotron-sft/cache - - --micro-batch-size 1 - - --global-batch-size 8 - - --attention-dropout 0.0 - - --hidden-dropout 0.0 - - --no-check-for-nan-in-loss-and-grad - - --recompute-granularity selective - - --recompute-modules layernorm moe - - --lr 5.0e-5 - - --min-lr 5.0e-6 - - --lr-decay-style cosine - - --clip-grad 1.0 - - --weight-decay 0.0 - - --adam-beta1 0.9 - - --adam-beta2 0.95 - - --init-method-std 0.010 - - --use-distributed-optimizer - - --eval-iters 4 - - --eval-interval 1000 - - --save-interval 1000 - - --log-interval 100 - - --sequence-parallel - - --logits-load-dir /scratchspace/logits + - --seq-length 8192 --max-position-embeddings 8192 + - --train-samples 1024 + - --lr-decay-samples 256 + - --lr-warmup-samples 128 + - --sft + - --sft-tokenizer-prompt-format identity + - --tokenizer-type SFTTokenizer + - --no-create-attention-mask-in-dataloader + - --per-split-data-args-path /hf-local/modelopt/nemotron-3-super-sft/blend.json + - --data-cache-path /hf-local/modelopt/test-nemotron-sft/cache + - --micro-batch-size 1 + - --global-batch-size 8 + - --attention-dropout 0.0 + - --hidden-dropout 0.0 + - --no-check-for-nan-in-loss-and-grad + - --recompute-granularity selective + - --recompute-modules layernorm moe + - --lr 5.0e-5 + - --min-lr 5.0e-6 + - --lr-decay-style cosine + - --clip-grad 1.0 + - --weight-decay 0.0 + - --adam-beta1 0.9 + - --adam-beta2 0.95 + - --init-method-std 0.010 + - --use-distributed-optimizer + - --eval-iters 4 + - --eval-interval 1000 + - --save-interval 1000 + - --log-interval 100 + - --sequence-parallel + - --logits-load-dir /scratchspace/logits environment: - - DP: 1 - - CP: 2 - - TP: 2 - - PP: 1 - - EP: 8 - - ETP: 1 - - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - - MLM_DATA_ARGS: --sft - - PYTHONPATH: /nemo_run/code/modules/Megatron-LM + - DP: 1 + - CP: 2 + - TP: 2 + - PP: 1 + - EP: 8 + - ETP: 1 + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - MLM_DATA_ARGS: --sft + - PYTHONPATH: /nemo_run/code/modules/Megatron-LM slurm_config: _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.06 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt nodes: 2 + ntasks_per_node: 8 + gpus_per_node: 8 task_3: script: common/megatron_lm/export/export.sh environment: - - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - - PP: 4 + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - QUANT_CFG: MAMBA_MOE_NVFP4_CONSERVATIVE_CFG + - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - EXPORT_DIR: /scratchspace/export + - PP: 4 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: 8 gpus_per_node: 8 task_4: script: common/tensorrt_llm/eval.sh + environment: + - HF_MODEL_CKPT: /scratchspace/export slurm_config: _factory_: "slurm_factory" nodes: 1 + ntasks_per_node: 4 + gpus_per_node: 4