diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/megatron_lm_ptq.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/megatron_lm_ptq.yaml new file mode 100644 index 00000000000..9f7b1ec126b --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/megatron_lm_ptq.yaml @@ -0,0 +1,110 @@ +# NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 PTQ quantization + MMLU gate + export + vLLM smoke. +# NemotronH-class hybrid Mamba-Transformer MoE. Smallest NemotronH MoE we ship +# an example for, so it doubles as the fast MoE expert-parallel exerciser. +# +# Pipeline (1 node x 4 GPUs throughout): +# task_0 (quantize + MMLU): TP=1 PP=1 EP=4 ETP=1 — the MoE experts shard across +# the 4 ranks (expert-parallel collective). Quantize the HF +# weights from /hf-local (PTQ ckpt to /cicd), then run MMLU +# (5-shot) on that same EP=4 collective as a regression gate. +# MMLU runs at --fraction 0.01 (quantize.sh default) — a light +# sample, so MMLU_LOWER_BOUND is set conservatively vs a full +# run. Both stages exercise the MoE expert-parallel collective. +# task_1 (export): TP=1 PP=4 EP=1 ETP=1 — pipeline-parallel for the hybrid +# layer stack; writes the HF NVFP4 ckpt to /cicd/export. +# task_2 (smoke): serve the exported NVFP4 ckpt with vLLM (TP=4) and answer +# 8 GPQA-style questions. Inspect under /cicd/vllm//. +# +# No dedicated Nano-30B recipe exists (only Nano-4B / Super-120B / Ultra-550B), +# so this uses the named NVFP4_DEFAULT_CFG rather than a recipe path. +# +# GOTCHA — nemo containers install ModelOpt into a VENV at /opt/venv, whose +# site-packages precede /usr/local on sys.path. The default modelopt_install_path +# (/usr/local/lib/.../dist-packages/modelopt) is therefore never consulted and +# the container's modelopt loads instead of the submodule pin. Point it at the +# venv site-packages so the mounted modelopt actually overrides. modelopt_recipes +# follows automatically (derived from this path's parent in core.py). +# +# Usage: +# source .env-slurm +# cd tools/launcher +# uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/megatron_lm_ptq.yaml --yes + +job_name: Nemotron-3-Nano-30B-A3B_PTQ +pipeline: + skip: false + allow_to_fail: false + note: "PTQ on Nemotron-3-Nano-30B-A3B (NVFP4): quantize + MMLU gate + export + vLLM smoke, 1 node x 4 GPUs" + + task_0: + script: common/megatron_lm/quantize/quantize.sh + args: + - --seq-length 4096 --max-position-embeddings 4096 + - --skip-generate + # Fast calibration. Bump (e.g. --calib-size 512) for production. + - --calib-size 32 + environment: + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - QUANT_CFG: NVFP4_DEFAULT_CFG + - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + # MMLU regression gate on the quantized ckpt (same EP=4 expert-parallel + # collective). Export runs as its own task, so RUN_EXPORT stays false. + - RUN_MMLU: "true" + - MMLU_DATASET: /hf-local/cais/mmlu + - MMLU_LOWER_BOUND: "0.60" + - RUN_EXPORT: "false" + - TP: "1" + - PP: "1" + - EP: "4" + - ETP: "1" + slurm_config: + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.04 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + partition: batch + nodes: 1 + ntasks_per_node: 4 + gpus_per_node: 4 + time: "04:00:00" + + task_1: + script: common/megatron_lm/export/export.sh + environment: + - MLM_MODEL_CFG: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - QUANT_CFG: NVFP4_DEFAULT_CFG + - HF_MODEL_CKPT: /hf-local/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + - TP: "1" + - PP: "4" + - EP: "1" + - ETP: "1" + slurm_config: + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.04 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + partition: batch + nodes: 1 + ntasks_per_node: 4 + gpus_per_node: 4 + time: "02:00:00" + + # vLLM generation test: serve the exported HF NVFP4 ckpt and answer 8 + # GPQA-style questions. Inspect responses under /cicd/vllm//. + task_2: + script: common/vllm/query.sh + args: + - --model /cicd/export/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16_NVFP4_DEFAULT_CFG + - --tensor-parallel-size 4 + - --trust-remote-code + - -- + - --data common/vllm/gpqa_sample.jsonl + - --max-tokens 256 + - --num-shards 1 + - --save /cicd/vllm/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16_NVFP4_DEFAULT_CFG + slurm_config: + _factory_: "slurm_factory" + container: vllm/vllm-openai:v0.21.0 + partition: batch + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 4 + time: "01:00:00"