From b2375d4ed245694f95ff32c2e3169f80529d7644 Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Fri, 19 Jun 2026 17:25:50 +0200 Subject: [PATCH 1/8] 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 d3ac9bc4246409b7b38e016d034e2d2abbf7dfff Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Mon, 20 Jul 2026 19:27:34 +0200 Subject: [PATCH 2/8] 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 From d988d7fc40fa6b38c89b382db8d284b6a34f0e2a Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Tue, 21 Jul 2026 19:17:04 +0200 Subject: [PATCH 3/8] Tasks 1-3 working Signed-off-by: Asha Anoosheh --- .../offline_kd_qad.yaml | 51 +++++++++++++++++-- 1 file changed, 48 insertions(+), 3 deletions(-) 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 6e90364d262..817852da041 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 @@ -17,10 +17,46 @@ # modelopt_install_path is never consulted. Point it at the venv # site-packages so the mounted (submodule-pinned) modelopt actually # overrides the container's. +# - task_1's --finetune is required, not optional: MLM_MODEL_CKPT +# (/hf-local/.../BF16-MCore) is a shared, multi-tenant path that may hold +# another run's completed checkpoint (e.g. iter_0000128 with +# consumed_train_samples already at --train-samples). Without --finetune, +# Megatron resumes that state instead of starting the frozen dump forward +# pass fresh, and MegatronPretrainingSampler asserts "no samples left to +# consume" immediately. +# - task_0 stays on ModelOpt's default "local spec" (no --export-default- +# te-spec): common/megatron_lm/quantize/quantize.sh's inline MMLU/export +# sub-steps each overwrite MLM_EXTRA_ARGS with their own hardcoded args, +# dropping any custom spec flag, so a save-time-only --export-default- +# te-spec would make those sub-steps reload the checkpoint under a +# different (mismatched) spec default and hit an _extra_state KeyError. +# RUN_MMLU=false / RUN_EXPORT=false skip those sub-steps regardless +# (see below), so this isn't strictly load-bearing, but keeps task_0 +# simple and matches the sibling BF16 PTQ examples for this model family. +# - task_2 needs --export-default-te-spec (switches MoE experts to +# TEGroupedMLP), for TP=2 + EP=8 together: ModelOpt's default +# SequentialMLP quant wrapper hard-disallows TP>1 and EP>1 at the same +# time, which is required to avoid an OOM in the KD loss's full-vocab +# softmax norm. task_1 doesn't need it: it loads the separate, unquantized +# BF16-MCore teacher checkpoint (not task_0's output) at TP=1/EP=8, which +# never hits that restriction regardless of spec, and it has completed +# successfully without this flag. task_0 stays on the local spec (see +# RUN_MMLU/RUN_EXPORT note below), so any task loading a checkpoint +# task_2 saved (itself on the TE spec) can hit the same _extra_state key +# mismatch the sibling Llama-3.2-1B QAD example hits -- task_2 and task_3 +# (export, which loads task_2's trained output) both carry +# --export-default-te-spec / --dist-ckpt-strictness log_all to build +# under the matching spec and tolerate mismatched keys (logs them instead +# of hard-raising). +# - task_0 sets RUN_MMLU=false and RUN_EXPORT=false: neither inline +# sub-step is needed here (the MMLU gate is already bypassed below, and +# task_3 does its own dedicated export), and skipping them avoids the +# spec-mismatch risk noted above. # -# 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). +# MMLU_LOWER_BOUND=0.0 was the temporary bypass on task_0's inline MMLU gate; +# RUN_MMLU=false now skips that gate entirely (see above), so this value is +# currently a no-op. Revisit both before this pipeline is used for production +# runs (see OfflineKD MMLU accuracy note in project history). # # Usage: # source .env-slurm @@ -35,11 +71,14 @@ pipeline: - --calib-dataset-path-or-name /hf-local/abisee/cnn_dailymail - --calib-size 32 - --calib-max-sequence-length 512 + - --skip-generate 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 + - RUN_MMLU: "false" + - RUN_EXPORT: "false" - PYTHONPATH: /nemo_run/code/modules/Megatron-LM slurm_config: _factory_: "slurm_factory" @@ -81,6 +120,7 @@ pipeline: - --eval-interval 1000 - --save-interval 1000 - --log-interval 100 + - --finetune - --freeze-all-layers - --logits-save-top-k 16 - --logits-save-dir /scratchspace/logits @@ -138,6 +178,8 @@ pipeline: - --eval-interval 1000 - --save-interval 1000 - --log-interval 100 + - --export-default-te-spec + - --dist-ckpt-strictness log_all - --sequence-parallel - --logits-load-dir /scratchspace/logits environment: @@ -160,6 +202,9 @@ pipeline: task_3: script: common/megatron_lm/export/export.sh + args: + - --export-default-te-spec + - --dist-ckpt-strictness log_all environment: - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - QUANT_CFG: MAMBA_MOE_NVFP4_CONSERVATIVE_CFG From ee293de9dd506ad7535525ae0ecae42aa8a84e05 Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Wed, 22 Jul 2026 20:24:00 +0200 Subject: [PATCH 4/8] Switch to B100s Signed-off-by: Asha Anoosheh --- .../offline_kd_qad.yaml | 84 ++++++++++--------- 1 file changed, 45 insertions(+), 39 deletions(-) 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 817852da041..3ac7c40a9a1 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 @@ -25,33 +25,33 @@ # pass fresh, and MegatronPretrainingSampler asserts "no samples left to # consume" immediately. # - task_0 stays on ModelOpt's default "local spec" (no --export-default- -# te-spec): common/megatron_lm/quantize/quantize.sh's inline MMLU/export -# sub-steps each overwrite MLM_EXTRA_ARGS with their own hardcoded args, -# dropping any custom spec flag, so a save-time-only --export-default- -# te-spec would make those sub-steps reload the checkpoint under a -# different (mismatched) spec default and hit an _extra_state KeyError. -# RUN_MMLU=false / RUN_EXPORT=false skip those sub-steps regardless -# (see below), so this isn't strictly load-bearing, but keeps task_0 -# simple and matches the sibling BF16 PTQ examples for this model family. -# - task_2 needs --export-default-te-spec (switches MoE experts to -# TEGroupedMLP), for TP=2 + EP=8 together: ModelOpt's default -# SequentialMLP quant wrapper hard-disallows TP>1 and EP>1 at the same -# time, which is required to avoid an OOM in the KD loss's full-vocab -# softmax norm. task_1 doesn't need it: it loads the separate, unquantized -# BF16-MCore teacher checkpoint (not task_0's output) at TP=1/EP=8, which -# never hits that restriction regardless of spec, and it has completed -# successfully without this flag. task_0 stays on the local spec (see -# RUN_MMLU/RUN_EXPORT note below), so any task loading a checkpoint -# task_2 saved (itself on the TE spec) can hit the same _extra_state key -# mismatch the sibling Llama-3.2-1B QAD example hits -- task_2 and task_3 -# (export, which loads task_2's trained output) both carry -# --export-default-te-spec / --dist-ckpt-strictness log_all to build -# under the matching spec and tolerate mismatched keys (logs them instead -# of hard-raising). -# - task_0 sets RUN_MMLU=false and RUN_EXPORT=false: neither inline -# sub-step is needed here (the MMLU gate is already bypassed below, and -# task_3 does its own dedicated export), and skipping them avoids the -# spec-mismatch risk noted above. +# te-spec) and sets RUN_MMLU=false / RUN_EXPORT=false: quantize.sh's +# inline MMLU/export sub-steps overwrite MLM_EXTRA_ARGS with their own +# args, dropping any custom spec flag, so a save-time-only +# --export-default-te-spec would make those sub-steps reload under a +# mismatched spec and hit an _extra_state KeyError. Neither sub-step is +# needed anyway (task_3 does its own dedicated export). +# - task_2 (and task_3, which loads task_2's output) need +# --export-default-te-spec + --dist-ckpt-strictness log_all: ModelOpt's +# default SequentialMLP quant wrapper hard-disallows TP>1 and EP>1 +# together, needed here to avoid an OOM in the KD loss's full-vocab +# softmax norm, and switching specs risks the same _extra_state mismatch +# the sibling Llama-3.2-1B QAD example works around the same way. task_1 +# doesn't need either (TP=1, and it loads an unrelated, unquantized +# checkpoint). +# - All node/GPU counts here target 4-GPU nodes (e.g. OCI-HSG's B100 nodes): +# task_4's TRT-LLM eval needs native NVFP4 tensor-core support (Blackwell +# only) -- on Hopper (H100), TRT-LLM's fused MoE GEMM kernel has no valid +# tactic for this model's shapes and segfaults (`FusedMoeRunner::runMoe`). +# Since the whole pipeline must share a cluster with task_4, everything +# targets 4-GPU nodes. +# - task_0/task_1/task_2 use EP=4 (not EP=8) to fit task_2's TP=2 within 2 +# nodes x 4 GPUs: its expert-eligible pool is world_size/TP = 8/2 = 4. +# - No DP env var on task_1/task_2: NeMo Run's SlurmExecutor sets +# LAUNCH_SCRIPT=python, bypassing the torchrun/DP-based nproc_per_node +# calculation in conf/arguments.sh entirely -- world size comes from +# Slurm's nodes x ntasks_per_node, and Megatron derives its own DP from +# world_size/(TP*PP*CP), so a DP env var here would be inert. # # MMLU_LOWER_BOUND=0.0 was the temporary bypass on task_0's inline MMLU gate; # RUN_MMLU=false now skips that gate entirely (see above), so this value is @@ -85,8 +85,8 @@ pipeline: 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 + ntasks_per_node: 4 + gpus_per_node: 4 task_1: script: common/megatron_lm/train/sft.sh @@ -102,6 +102,14 @@ pipeline: - --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 + # Must match task_2's --global-batch-size: the cached teacher logits are + # indexed per (iteration, microbatch), so teacher and student must see the + # same samples per iteration. Without this, task_1's global batch defaults + # to micro_batch*DP (=4 at DP=4 here) while task_2's is 8, and the student + # runs past the end of the teacher's dump (IndexError: microbatch index + # out of range). It only aligned on CW-DFW by coincidence (teacher DP=8 + # there made the auto-default equal task_2's 8). + - --global-batch-size 8 - --attention-dropout 0.0 - --hidden-dropout 0.0 - --no-check-for-nan-in-loss-and-grad @@ -127,11 +135,10 @@ pipeline: - --async-save - --use-persistent-ckpt-worker environment: - - DP: 1 - CP: 2 - TP: 1 - PP: 1 - - EP: 8 + - 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 @@ -142,8 +149,8 @@ pipeline: 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 + ntasks_per_node: 4 + gpus_per_node: 4 task_2: script: common/megatron_lm/train/sft.sh @@ -183,11 +190,10 @@ pipeline: - --sequence-parallel - --logits-load-dir /scratchspace/logits environment: - - DP: 1 - CP: 2 - TP: 2 - PP: 1 - - EP: 8 + - EP: 4 - ETP: 1 - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - MLM_DATA_ARGS: --sft @@ -197,8 +203,8 @@ pipeline: 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 + ntasks_per_node: 4 + gpus_per_node: 4 task_3: script: common/megatron_lm/export/export.sh @@ -216,8 +222,8 @@ pipeline: 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 + ntasks_per_node: 4 + gpus_per_node: 4 task_4: script: common/tensorrt_llm/eval.sh From a24a803c7e9d57f6e19d0aad465eb2efc644bc1c Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Mon, 27 Jul 2026 21:16:37 +0200 Subject: [PATCH 5/8] Fix keyerror Signed-off-by: Asha Anoosheh --- .../offline_kd_qad.yaml | 53 ++++++++++++------- 1 file changed, 33 insertions(+), 20 deletions(-) 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 3ac7c40a9a1..1fe0b25840e 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 @@ -24,21 +24,29 @@ # Megatron resumes that state instead of starting the frozen dump forward # pass fresh, and MegatronPretrainingSampler asserts "no samples left to # consume" immediately. -# - task_0 stays on ModelOpt's default "local spec" (no --export-default- -# te-spec) and sets RUN_MMLU=false / RUN_EXPORT=false: quantize.sh's -# inline MMLU/export sub-steps overwrite MLM_EXTRA_ARGS with their own -# args, dropping any custom spec flag, so a save-time-only -# --export-default-te-spec would make those sub-steps reload under a -# mismatched spec and hit an _extra_state KeyError. Neither sub-step is -# needed anyway (task_3 does its own dedicated export). -# - task_2 (and task_3, which loads task_2's output) need -# --export-default-te-spec + --dist-ckpt-strictness log_all: ModelOpt's -# default SequentialMLP quant wrapper hard-disallows TP>1 and EP>1 -# together, needed here to avoid an OOM in the KD loss's full-vocab -# softmax norm, and switching specs risks the same _extra_state mismatch -# the sibling Llama-3.2-1B QAD example works around the same way. task_1 -# doesn't need either (TP=1, and it loads an unrelated, unquantized -# checkpoint). +# - task_0, task_2, and task_3 ALL use --export-default-te-spec (TE +# TEGroupedMLP experts). This MUST be consistent across every task that +# touches the *quantized* checkpoint: the NVFP4 quantizer calibration +# scales (linear_fc*.{weight,input}_quantizer._amax) are keyed by the MoE +# expert layout, so if task_0 saves under the default local spec +# (per-expert SequentialMLP) but task_2 loads under TE spec (grouped), +# ~93k expert _amax keys silently mismatch and --dist-ckpt-strictness +# log_all drops them -> the expert quantizers run uncalibrated -> the model +# is effectively random (MMLU ~25%). task_2 needs TE spec anyway (its +# default SequentialMLP quant wrapper hard-disallows TP>1 and EP>1 together, +# which task_2 requires to avoid an OOM in the KD loss's full-vocab softmax +# norm), so task_0 and task_3 match it. With the spec aligned, the only +# remaining load mismatch is the benign pre_mlp_layernorm._extra_state +# keys, which --dist-ckpt-strictness log_all correctly tolerates on +# task_2/task_3. +# - task_0 also needs RUN_MMLU=false / RUN_EXPORT=false: quantize.sh's inline +# MMLU/export sub-steps overwrite MLM_EXTRA_ARGS with their own args, +# dropping --export-default-te-spec, so they would reload the just-saved +# TE-spec checkpoint under the mismatched default spec and hit a KeyError. +# Neither sub-step is needed here (task_3 does its own dedicated export). +# - task_1 needs neither spec flag nor strictness override: it's the +# unquantized BF16 teacher (no _amax keys) at TP=1, loading an unrelated +# checkpoint. # - All node/GPU counts here target 4-GPU nodes (e.g. OCI-HSG's B100 nodes): # task_4's TRT-LLM eval needs native NVFP4 tensor-core support (Blackwell # only) -- on Hopper (H100), TRT-LLM's fused MoE GEMM kernel has no valid @@ -53,10 +61,14 @@ # Slurm's nodes x ntasks_per_node, and Megatron derives its own DP from # world_size/(TP*PP*CP), so a DP env var here would be inert. # -# MMLU_LOWER_BOUND=0.0 was the temporary bypass on task_0's inline MMLU gate; -# RUN_MMLU=false now skips that gate entirely (see above), so this value is -# currently a no-op. Revisit both before this pipeline is used for production -# runs (see OfflineKD MMLU accuracy note in project history). +# MMLU_LOWER_BOUND=0.6 is a real quality gate now that the spec-mismatch bug is +# fixed (the trained model scores ~0.70 MMLU via task_4's TRT-LLM eval, up from +# the ~0.25 random-chance the old --export-default-te-spec-less task_0 produced). +# It is currently INERT because task_0's inline MMLU gate is disabled +# (RUN_MMLU=false) -- re-enabling it would reload the TE-spec checkpoint under a +# mismatched spec (see the RUN_MMLU note above). The effective quality signal is +# task_4's MMLU. The value is kept meaningful (not 0.0) to document the expected +# floor for whenever the inline gate is wired back up. # # Usage: # source .env-slurm @@ -72,11 +84,12 @@ pipeline: - --calib-size 32 - --calib-max-sequence-length 512 - --skip-generate + - --export-default-te-spec 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 + - MMLU_LOWER_BOUND: "0.6" - RUN_MMLU: "false" - RUN_EXPORT: "false" - PYTHONPATH: /nemo_run/code/modules/Megatron-LM From 67dc669d6a2ad15b3b3e015d42dd23ffce81b94b Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Tue, 28 Jul 2026 19:21:02 +0200 Subject: [PATCH 6/8] Increase performance Signed-off-by: Asha Anoosheh --- .../offline_kd_qad.yaml | 63 +++++++++++++------ 1 file changed, 45 insertions(+), 18 deletions(-) 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 1fe0b25840e..1c71ea47d78 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 @@ -53,13 +53,27 @@ # tactic for this model's shapes and segfaults (`FusedMoeRunner::runMoe`). # Since the whole pipeline must share a cluster with task_4, everything # targets 4-GPU nodes. -# - task_0/task_1/task_2 use EP=4 (not EP=8) to fit task_2's TP=2 within 2 -# nodes x 4 GPUs: its expert-eligible pool is world_size/TP = 8/2 = 4. +# - task_0/task_1/task_2 use EP=4 (128 experts / 4 = 32 experts per rank). +# - Long context (seq 32768): the SFT corpus (train.jsonl) is raw, not +# pre-tokenized, and every sample is huge (>=170k chars, ~30k-360k tokens), +# so seq 32768 is fully consumed (truncation-bound, no padding waste). The +# model has position-embedding-type=none (Mamba + NoPE attention), so +# extending past the checkpoint's 8192 max_position_embeddings is safe. +# teacher (task_1) and student (task_2) MUST share seq-length AND CP: the +# cached top-K logits are stored per (iteration, microbatch, token) and +# sharded by CP, so both dump and load must partition the sequence +# identically. +# - CP=8 (up from the original CP=2 at seq 8192): scaling CP with seq-length +# keeps the per-GPU sequence chunk at seq/CP = 32768/8 = 4096 -- identical +# to the validated 69.98%-MMLU run -- so per-device activation/logit memory +# is unchanged. task_1 (TP=1) fits CP=8 on 2 nodes (8 GPUs, DP=1); task_2 +# (TP=2) needs 4 nodes (16 GPUs, DP=1) for TP*CP = 2*8 = 16 = world. # - No DP env var on task_1/task_2: NeMo Run's SlurmExecutor sets # LAUNCH_SCRIPT=python, bypassing the torchrun/DP-based nproc_per_node # calculation in conf/arguments.sh entirely -- world size comes from # Slurm's nodes x ntasks_per_node, and Megatron derives its own DP from -# world_size/(TP*PP*CP), so a DP env var here would be inert. +# world_size/(TP*PP*CP), so a DP env var here would be inert. Here both +# resolve to DP=1 (task_1: 8/(1*1*8); task_2: 16/(2*1*8)). # # MMLU_LOWER_BOUND=0.6 is a real quality gate now that the spec-mismatch bug is # fixed (the trained model scores ~0.70 MMLU via task_4's TRT-LLM eval, up from @@ -80,9 +94,13 @@ 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 + # In-domain calibration on the same SFT corpus task_1/task_2 train on + # (raw .jsonl the calib loader reads via its "messages" branch), rather + # than out-of-domain cnn_dailymail -- better-matched NVFP4 amax scales. + # 512 samples @ 8192 tokens (up from 32 @ 512) for a fuller calibration. + - --calib-dataset-path-or-name /hf-local/modelopt/nemotron-3-super-sft/train.jsonl + - --calib-size 512 + - --calib-max-sequence-length 8192 - --skip-generate - --export-default-te-spec environment: @@ -104,10 +122,10 @@ 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 + - --seq-length 32768 --max-position-embeddings 32768 + - --train-samples 2048 + - --lr-decay-samples 512 + - --lr-warmup-samples 256 - --sft - --sft-tokenizer-prompt-format identity - --tokenizer-type SFTTokenizer @@ -148,7 +166,7 @@ pipeline: - --async-save - --use-persistent-ckpt-worker environment: - - CP: 2 + - CP: 8 - TP: 1 - PP: 1 - EP: 4 @@ -168,10 +186,10 @@ pipeline: 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 + - --seq-length 32768 --max-position-embeddings 32768 + - --train-samples 2048 + - --lr-decay-samples 512 + - --lr-warmup-samples 256 - --sft - --sft-tokenizer-prompt-format identity - --tokenizer-type SFTTokenizer @@ -196,14 +214,16 @@ pipeline: - --use-distributed-optimizer - --eval-iters 4 - --eval-interval 1000 - - --save-interval 1000 + # = train-samples / global-batch-size (2048/8) so the trained student is + # checkpointed exactly once, at the final iteration (task_3 exports it). + - --save-interval 256 - --log-interval 100 - --export-default-te-spec - --dist-ckpt-strictness log_all - --sequence-parallel - --logits-load-dir /scratchspace/logits environment: - - CP: 2 + - CP: 8 - TP: 2 - PP: 1 - EP: 4 @@ -215,7 +235,7 @@ pipeline: _factory_: "slurm_factory" container: nvcr.io/nvidia/nemo:26.06 modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt - nodes: 2 + nodes: 4 ntasks_per_node: 4 gpus_per_node: 4 @@ -230,6 +250,13 @@ pipeline: - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - EXPORT_DIR: /scratchspace/export - PP: 4 + # Put the packaged Megatron-LM AND Model-Optimizer first on PYTHONPATH. + # nemo:26.06 uses a uv venv, so the modelopt_install_path bind-mount at + # /opt/venv may be stale and the container's own modelopt shadows it -- + # export then fails importing modelopt.torch.sparsity.weight_sparsity.config + # (absent in the container's older modelopt). PYTHONPATH forces the + # packaged modelopt, which has it. See nemo-container-uv-venv-path. + - PYTHONPATH: /nemo_run/code/modules/Megatron-LM:/nemo_run/code/modules/Model-Optimizer slurm_config: _factory_: "slurm_factory" container: nvcr.io/nvidia/nemo:26.06 From 2c1ff0e2191492baf5bee0aa1ff6fae8ec0a5b65 Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Tue, 28 Jul 2026 19:34:32 +0200 Subject: [PATCH 7/8] Compact comments Signed-off-by: Asha Anoosheh --- .../offline_kd_qad.yaml | 145 +++++------------- 1 file changed, 39 insertions(+), 106 deletions(-) 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 1c71ea47d78..6bb21bff53a 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 @@ -1,92 +1,36 @@ -# Nemotron-3-Nano-30B-A3B NVFP4 Offline KD QAD pipeline. +# Nemotron-3-Nano-30B-A3B NVFP4 Offline-KD QAD pipeline (cached-teacher-logits KD). +# task_0 quantize : PTQ BF16 -> NVFP4 MCore ckpt (in-domain calibration). +# task_1 dump : frozen BF16 teacher, save top-K logits to /scratchspace/logits. +# task_2 student : train NVFP4 student with --logits-load-dir. +# task_3 export : trained student -> HF format. +# task_4 eval : TRT-LLM MMLU. # -# 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: -# - 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. -# - task_1's --finetune is required, not optional: MLM_MODEL_CKPT -# (/hf-local/.../BF16-MCore) is a shared, multi-tenant path that may hold -# another run's completed checkpoint (e.g. iter_0000128 with -# consumed_train_samples already at --train-samples). Without --finetune, -# Megatron resumes that state instead of starting the frozen dump forward -# pass fresh, and MegatronPretrainingSampler asserts "no samples left to -# consume" immediately. -# - task_0, task_2, and task_3 ALL use --export-default-te-spec (TE -# TEGroupedMLP experts). This MUST be consistent across every task that -# touches the *quantized* checkpoint: the NVFP4 quantizer calibration -# scales (linear_fc*.{weight,input}_quantizer._amax) are keyed by the MoE -# expert layout, so if task_0 saves under the default local spec -# (per-expert SequentialMLP) but task_2 loads under TE spec (grouped), -# ~93k expert _amax keys silently mismatch and --dist-ckpt-strictness -# log_all drops them -> the expert quantizers run uncalibrated -> the model -# is effectively random (MMLU ~25%). task_2 needs TE spec anyway (its -# default SequentialMLP quant wrapper hard-disallows TP>1 and EP>1 together, -# which task_2 requires to avoid an OOM in the KD loss's full-vocab softmax -# norm), so task_0 and task_3 match it. With the spec aligned, the only -# remaining load mismatch is the benign pre_mlp_layernorm._extra_state -# keys, which --dist-ckpt-strictness log_all correctly tolerates on -# task_2/task_3. -# - task_0 also needs RUN_MMLU=false / RUN_EXPORT=false: quantize.sh's inline -# MMLU/export sub-steps overwrite MLM_EXTRA_ARGS with their own args, -# dropping --export-default-te-spec, so they would reload the just-saved -# TE-spec checkpoint under the mismatched default spec and hit a KeyError. -# Neither sub-step is needed here (task_3 does its own dedicated export). -# - task_1 needs neither spec flag nor strictness override: it's the -# unquantized BF16 teacher (no _amax keys) at TP=1, loading an unrelated -# checkpoint. -# - All node/GPU counts here target 4-GPU nodes (e.g. OCI-HSG's B100 nodes): -# task_4's TRT-LLM eval needs native NVFP4 tensor-core support (Blackwell -# only) -- on Hopper (H100), TRT-LLM's fused MoE GEMM kernel has no valid -# tactic for this model's shapes and segfaults (`FusedMoeRunner::runMoe`). -# Since the whole pipeline must share a cluster with task_4, everything -# targets 4-GPU nodes. -# - task_0/task_1/task_2 use EP=4 (128 experts / 4 = 32 experts per rank). -# - Long context (seq 32768): the SFT corpus (train.jsonl) is raw, not -# pre-tokenized, and every sample is huge (>=170k chars, ~30k-360k tokens), -# so seq 32768 is fully consumed (truncation-bound, no padding waste). The -# model has position-embedding-type=none (Mamba + NoPE attention), so -# extending past the checkpoint's 8192 max_position_embeddings is safe. -# teacher (task_1) and student (task_2) MUST share seq-length AND CP: the -# cached top-K logits are stored per (iteration, microbatch, token) and -# sharded by CP, so both dump and load must partition the sequence -# identically. -# - CP=8 (up from the original CP=2 at seq 8192): scaling CP with seq-length -# keeps the per-GPU sequence chunk at seq/CP = 32768/8 = 4096 -- identical -# to the validated 69.98%-MMLU run -- so per-device activation/logit memory -# is unchanged. task_1 (TP=1) fits CP=8 on 2 nodes (8 GPUs, DP=1); task_2 -# (TP=2) needs 4 nodes (16 GPUs, DP=1) for TP*CP = 2*8 = 16 = world. -# - No DP env var on task_1/task_2: NeMo Run's SlurmExecutor sets -# LAUNCH_SCRIPT=python, bypassing the torchrun/DP-based nproc_per_node -# calculation in conf/arguments.sh entirely -- world size comes from -# Slurm's nodes x ntasks_per_node, and Megatron derives its own DP from -# world_size/(TP*PP*CP), so a DP env var here would be inert. Here both -# resolve to DP=1 (task_1: 8/(1*1*8); task_2: 16/(2*1*8)). -# -# MMLU_LOWER_BOUND=0.6 is a real quality gate now that the spec-mismatch bug is -# fixed (the trained model scores ~0.70 MMLU via task_4's TRT-LLM eval, up from -# the ~0.25 random-chance the old --export-default-te-spec-less task_0 produced). -# It is currently INERT because task_0's inline MMLU gate is disabled -# (RUN_MMLU=false) -- re-enabling it would reload the TE-spec checkpoint under a -# mismatched spec (see the RUN_MMLU note above). The effective quality signal is -# task_4's MMLU. The value is kept meaningful (not 0.0) to document the expected -# floor for whenever the inline gate is wired back up. +# Key constraints: +# - Container nvcr.io/nvidia/nemo:26.06 (needs nvidia-resiliency-ext>=0.6.0 for +# task_1's --async-save). It uses a uv venv, so point modelopt_install_path at +# the venv site-packages and put the packaged Megatron-LM/Model-Optimizer on +# PYTHONPATH so they override the container's copies (see task_3). +# - task_0/task_2/task_3 must all pass --export-default-te-spec (TEGroupedMLP): +# NVFP4 expert _amax scales are keyed by MoE layout, so a save/load spec +# mismatch silently drops ~93k _amax keys (log_all) -> uncalibrated model +# (MMLU ~25%). task_2 needs TE spec regardless (its default SequentialMLP +# wrapper forbids TP>1 & EP>1, needed to avoid a KD-loss OOM); +# --dist-ckpt-strictness log_all then only tolerates benign _extra_state. +# - task_0 RUN_MMLU/RUN_EXPORT=false: quantize.sh's inline sub-steps drop +# --export-default-te-spec and would KeyError on the TE-spec ckpt (task_3 exports). +# - task_1 --finetune: MLM_MODEL_CKPT is a shared path that may hold another run's +# completed ckpt; without it the sampler asserts "no samples left to consume". +# - 4-GPU nodes (OCI-HSG B100): task_4's NVFP4 TRT-LLM eval needs Blackwell (Hopper +# segfaults in the fused MoE FP4 GEMM); the whole pipeline shares that cluster. +# - seq 32768: the SFT corpus is raw jsonl with very long samples (~30k-360k tok), +# so it's truncation-bound (no padding waste); position-embedding-type=none makes +# >8192 context safe. teacher & student MUST share seq-length AND CP (logits are +# CP-sharded). CP=8 keeps per-GPU seq/CP=4096 (the validated footprint): task_1 +# TP1 on 2 nodes, task_2 TP2 on 4 nodes (TP*CP=16=world). DP resolves to 1; a DP +# env var is inert under NeMo Run's python launcher. # # Usage: -# source .env-slurm -# cd tools/launcher +# 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 @@ -94,10 +38,8 @@ pipeline: task_0: script: common/megatron_lm/quantize/quantize.sh args: - # In-domain calibration on the same SFT corpus task_1/task_2 train on - # (raw .jsonl the calib loader reads via its "messages" branch), rather - # than out-of-domain cnn_dailymail -- better-matched NVFP4 amax scales. - # 512 samples @ 8192 tokens (up from 32 @ 512) for a fuller calibration. + # In-domain calibration on the SFT corpus (raw .jsonl, read via the calib + # loader's "messages" branch) instead of cnn_dailymail -- better-matched amax. - --calib-dataset-path-or-name /hf-local/modelopt/nemotron-3-super-sft/train.jsonl - --calib-size 512 - --calib-max-sequence-length 8192 @@ -107,7 +49,6 @@ pipeline: - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - QUANT_CFG: MAMBA_MOE_NVFP4_CONSERVATIVE_CFG - EP: 4 - - MMLU_LOWER_BOUND: "0.6" - RUN_MMLU: "false" - RUN_EXPORT: "false" - PYTHONPATH: /nemo_run/code/modules/Megatron-LM @@ -133,13 +74,9 @@ pipeline: - --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 - # Must match task_2's --global-batch-size: the cached teacher logits are - # indexed per (iteration, microbatch), so teacher and student must see the - # same samples per iteration. Without this, task_1's global batch defaults - # to micro_batch*DP (=4 at DP=4 here) while task_2's is 8, and the student - # runs past the end of the teacher's dump (IndexError: microbatch index - # out of range). It only aligned on CW-DFW by coincidence (teacher DP=8 - # there made the auto-default equal task_2's 8). + # Must match task_2's --global-batch-size: cached logits are indexed per + # (iteration, microbatch), so teacher and student must see the same samples + # per iteration (else "microbatch index out of range"). - --global-batch-size 8 - --attention-dropout 0.0 - --hidden-dropout 0.0 @@ -214,8 +151,7 @@ pipeline: - --use-distributed-optimizer - --eval-iters 4 - --eval-interval 1000 - # = train-samples / global-batch-size (2048/8) so the trained student is - # checkpointed exactly once, at the final iteration (task_3 exports it). + # = train-samples / global-batch-size (2048/8): one save, at the final iter. - --save-interval 256 - --log-interval 100 - --export-default-te-spec @@ -250,12 +186,9 @@ pipeline: - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - EXPORT_DIR: /scratchspace/export - PP: 4 - # Put the packaged Megatron-LM AND Model-Optimizer first on PYTHONPATH. - # nemo:26.06 uses a uv venv, so the modelopt_install_path bind-mount at - # /opt/venv may be stale and the container's own modelopt shadows it -- - # export then fails importing modelopt.torch.sparsity.weight_sparsity.config - # (absent in the container's older modelopt). PYTHONPATH forces the - # packaged modelopt, which has it. See nemo-container-uv-venv-path. + # Force the packaged modelopt onto PYTHONPATH: nemo:26.06's uv venv can + # shadow the /opt/venv bind-mount with its own older modelopt, which lacks + # modelopt.torch.sparsity.weight_sparsity.config that export imports. - PYTHONPATH: /nemo_run/code/modules/Megatron-LM:/nemo_run/code/modules/Model-Optimizer slurm_config: _factory_: "slurm_factory" From 427e9820929300a00229de627438a8add93249e2 Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Thu, 30 Jul 2026 19:29:29 +0200 Subject: [PATCH 8/8] Address PR review: fix data-cache-path naming, note MCore prerequisite - data-cache-path: test-nemotron-sft -> nemotron-3-super-sft (CodeRabbit), matching the dataset/blend/calib paths (task_1 + task_2) - header: note task_1's BF16 MCore teacher ckpt is an external prerequisite, not produced by this pipeline (Claude review) Co-Authored-By: Claude Opus 4.8 Signed-off-by: Asha Anoosheh --- .../offline_kd_qad.yaml | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) 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 6bb21bff53a..29cb5280ab8 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 @@ -18,6 +18,9 @@ # --dist-ckpt-strictness log_all then only tolerates benign _extra_state. # - task_0 RUN_MMLU/RUN_EXPORT=false: quantize.sh's inline sub-steps drop # --export-default-te-spec and would KeyError on the TE-spec ckpt (task_3 exports). +# - PREREQUISITE: task_1's teacher loads a BF16 MCore ckpt at MLM_MODEL_CKPT +# (NVIDIA-Nemotron-3-Nano-30B-A3B-BF16-MCore); this pipeline does not create +# it -- import it from HF first (e.g. common/megatron_bridge/import/import.sh). # - task_1 --finetune: MLM_MODEL_CKPT is a shared path that may hold another run's # completed ckpt; without it the sampler asserts "no samples left to consume". # - 4-GPU nodes (OCI-HSG B100): task_4's NVFP4 TRT-LLM eval needs Blackwell (Hopper @@ -72,7 +75,7 @@ pipeline: - --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 + - --data-cache-path /hf-local/modelopt/nemotron-3-super-sft/cache - --micro-batch-size 1 # Must match task_2's --global-batch-size: cached logits are indexed per # (iteration, microbatch), so teacher and student must see the same samples @@ -132,7 +135,7 @@ pipeline: - --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 + - --data-cache-path /hf-local/modelopt/nemotron-3-super-sft/cache - --micro-batch-size 1 - --global-batch-size 8 - --attention-dropout 0.0