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 new file mode 100644 index 00000000000..3ac7c40a9a1 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml @@ -0,0 +1,236 @@ +# 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: +# - 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 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). +# - 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 +# 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 +# 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 + - --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" + 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 + + 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 + # 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 + - --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 + - --finetune + - --freeze-all-layers + - --logits-save-top-k 16 + - --logits-save-dir /scratchspace/logits + - --async-save + - --use-persistent-ckpt-worker + environment: + - CP: 2 + - TP: 1 + - PP: 1 + - 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 + - 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: 4 + gpus_per_node: 4 + + 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 + - --export-default-te-spec + - --dist-ckpt-strictness log_all + - --sequence-parallel + - --logits-load-dir /scratchspace/logits + environment: + - CP: 2 + - TP: 2 + - PP: 1 + - EP: 4 + - 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: 4 + gpus_per_node: 4 + + 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 + - 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: 4 + gpus_per_node: 4 + + 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