diff --git a/examples/llm_eval/lm_eval_hf.py b/examples/llm_eval/lm_eval_hf.py index 51c0930e8f2..dbddb5ccafb 100755 --- a/examples/llm_eval/lm_eval_hf.py +++ b/examples/llm_eval/lm_eval_hf.py @@ -37,6 +37,10 @@ # See the License for the specific language governing permissions and # limitations under the License. import contextlib +import glob +import json +import os +import sys import warnings from importlib.metadata import version @@ -236,6 +240,17 @@ def _add_modelopt_args(parser): type=str, help="Sparse attention configuration (e.g., SKIP_SOFTMAX_DEFAULT, SKIP_SOFTMAX_CALIB)", ) + # Not a model arg (kept out of _MODELOPT_ARG_KEYS): popped separately and + # applied after the eval as an accuracy gate. + parser.add_argument( + "--accuracy_lower_bound", + type=float, + default=None, + help=( + "Exit non-zero if the requested task's acc is below this. Requires exactly one " + "--tasks and --output_path." + ), + ) def _inject_modelopt_args_into_model_args(args): @@ -243,9 +258,13 @@ def _inject_modelopt_args_into_model_args(args): args.model_args is a dict (parsed by lm-eval's MergeDictAction). The ModelOpt keys must be removed from the namespace so EvaluatorConfig.from_cli doesn't - reject them as unknown kwargs. + reject them as unknown kwargs. They only apply to the HF backend, so for other + backends (e.g. `--model vllm` on an already-quantized ckpt) they are dropped, + not passed into model_args (which would break the backend's constructor). """ model_args = dict(args.model_args) if args.model_args else {} + # HFLM and its subclasses (registered as hf / hf-auto / huggingface / hf-multimodal). + is_hf = getattr(args, "model", "hf") in ("hf", "hf-auto", "huggingface", "hf-multimodal") if getattr(args, "trust_remote_code", False): # Propagate the user-provided --trust_remote_code flag (not hardcoded). @@ -253,14 +272,68 @@ def _inject_modelopt_args_into_model_args(args): model_args["trust_remote_code"] = True args.trust_remote_code = None + # Args that actually request quantization/sparsity. The other _MODELOPT_ARG_KEYS are inert + # tuning knobs (calib_*, auto_quantize_* params) with non-None argparse defaults, so keying + # off "truthy" would false-positive on them — only the request args imply the wrong backend. + quant_request_keys = {"quant_cfg", "auto_quantize_bits", "compress", "sparse_cfg"} for key in _MODELOPT_ARG_KEYS: if hasattr(args, key): - model_args[key] = getattr(args, key) + value = getattr(args, key) + if is_hf: + model_args[key] = value + elif value and key in quant_request_keys: + # Don't silently drop a quantization/sparsity request on a non-HF backend, which + # would evaluate the wrong (unquantized) checkpoint and still report PASS. + raise ValueError( + f"--{key} requests quantization/sparsity, which only applies to the HF " + f"backend; with --model {getattr(args, 'model', '?')!r} the checkpoint must " + "already be quantized/sparsified (drop the ModelOpt args)." + ) delattr(args, key) args.model_args = model_args +def _pop_gate_bound(args): + """Pop --accuracy_lower_bound off the namespace; return it (None disables the gate).""" + bound = getattr(args, "accuracy_lower_bound", None) + if hasattr(args, "accuracy_lower_bound"): + delattr(args, "accuracy_lower_bound") + return bound + + +def _requested_tasks(args): + """Task names from lm-eval's --tasks (string or list).""" + tasks = getattr(args, "tasks", None) or [] + if isinstance(tasks, str): + tasks = [tasks] + return [p.strip() for item in tasks for p in str(item).split(",") if p.strip()] + + +def _enforce_accuracy_gate(output_path, task, lower_bound): + """Read lm-eval results at output_path; exit non-zero if /acc < lower_bound.""" + if not output_path: + raise ValueError("--accuracy_lower_bound requires --output_path.") + files = glob.glob(os.path.join(output_path, "**", "results*.json"), recursive=True) + if not files: + raise FileNotFoundError(f"No results*.json under {output_path}") + # Sort by mtime, not path: a reused output_path nests results under a + # / dir, and lexical order would pick the wrong run's file. + with open(max(files, key=os.path.getmtime)) as f: + scores = json.load(f)["results"].get(task, {}) + # lm-eval keys metrics by filter, e.g. "acc,none"; take acc (never acc_stderr). + acc = next((float(v) for k, v in scores.items() if k == "acc" or k.startswith("acc,")), None) + if acc is None: + raise KeyError(f"acc not found for '{task}' (have: {list(scores)})") + passed = acc >= lower_bound + print( + f"[accuracy_gate] {task}/acc = {acc:.4f} (lower_bound {lower_bound}) -> " + f"{'PASS' if passed else 'FAIL'}" + ) + if not passed: + sys.exit(1) + + if __name__ == "__main__": setup_logging() cli = HarnessCLI() @@ -277,4 +350,18 @@ def _inject_modelopt_args_into_model_args(args): _add_modelopt_args(run_parser) args = cli.parse_args() _inject_modelopt_args_into_model_args(args) + lower_bound = _pop_gate_bound(args) + output_path = getattr(args, "output_path", None) + gate_task = None + if lower_bound is not None: # fail fast before the (expensive) eval + if not output_path: + raise ValueError("--accuracy_lower_bound requires --output_path.") + tasks = _requested_tasks(args) + if len(tasks) != 1: + raise ValueError( + f"--accuracy_lower_bound needs exactly one --tasks (got {tasks or 'none'})." + ) + gate_task = tasks[0] cli.execute(args) + if lower_bound is not None: + _enforce_accuracy_gate(output_path, gate_task, lower_bound) diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 9f938495d4c..cd76d3c4696 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -42,6 +42,7 @@ import json import os import re +import sys import torch from megatron.bridge import AutoBridge @@ -237,6 +238,12 @@ def get_args() -> argparse.Namespace: "batch size for fast evaluation (default is mmlu_10pct_bs1)." ), ) + parser.add_argument( + "--score_lower_bound", + type=float, + default=None, + help="If set, fail the job when the NAS-based pruned model's score is below this bound.", + ) parser.add_argument( "--ss_channel_divisor", type=int, @@ -297,6 +304,11 @@ def get_args() -> argparse.Namespace: "At least one of --prune_export_config, --prune_target_params," " --prune_target_active_params, or --prune_target_memory_mb is required." ) + if args.score_lower_bound is not None and args.prune_export_config: + parser.error( + "--score_lower_bound requires NAS-based scoring (--prune_score_func), " + "not --prune_export_config." + ) # Post-process arguments if args.prune_intermediate_ckpt is None: @@ -668,6 +680,18 @@ def score_func(m): copy_hf_ckpt_remote_code(args.hf_model_name_or_path, args.output_hf_path) print_rank_0(f"Saved pruned model to {args.output_hf_path} in HF checkpoint format") + # Accuracy gate: exit non-zero if pruned model's score is below the bound + if args.score_lower_bound is not None: + best_score = pruning_scores["best"].get("score") + assert best_score is not None, "No scored best candidate in pruning_scores" + passed = best_score >= args.score_lower_bound + print_rank_0( + f"[score_gate] final pruned model {args.prune_score_func} score = {best_score:.4f} " + f"(lower_bound {args.score_lower_bound}) -> {'PASS' if passed else 'FAIL'}" + ) + if not passed: + sys.exit(1) + print_rank_0("Done!") diff --git a/modelopt/torch/prune/plugins/mcore_minitron.py b/modelopt/torch/prune/plugins/mcore_minitron.py index 876432fee01..28684fa3a4a 100644 --- a/modelopt/torch/prune/plugins/mcore_minitron.py +++ b/modelopt/torch/prune/plugins/mcore_minitron.py @@ -28,7 +28,7 @@ import sys from collections.abc import Callable from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import asdict, dataclass from functools import partial from itertools import product from typing import Any @@ -294,6 +294,7 @@ def default_state_dict(self) -> SearchStateDict: "layer_scores": {}, "sorted_layers": None, "all_candidates_per_constraint": {}, + "best": {}, } def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig: @@ -641,6 +642,7 @@ def search_best_arch_by_metrics(self) -> dict: dist.barrier() best = max(top_k_candidates, key=lambda x: x.score) # type: ignore[arg-type, return-value] + self.best = asdict(best) best_grid = Table.grid(padding=(0, 2)) best_grid.add_column(style="bold green", no_wrap=True) best_grid.add_column() diff --git a/tests/examples/llm_eval/test_llm_eval.py b/tests/examples/llm_eval/test_llm_eval.py index 1004a57e1b9..61934f2b141 100644 --- a/tests/examples/llm_eval/test_llm_eval.py +++ b/tests/examples/llm_eval/test_llm_eval.py @@ -36,6 +36,9 @@ def test_lm_eval_hf(tmp_path): num_fewshot=5, limit=0.1, batch_size=8, + # Exercise the accuracy gate: reads the results file and enforces the bound + output_path=str(tmp_path / "results"), + accuracy_lower_bound=0.1, ) run_example_command(cmd_parts, "llm_eval") diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index f468e8abcba..d5c6e8c823b 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -70,6 +70,7 @@ def test_prune_minitron(tmp_path, num_gpus, create_teacher, megatron_format): seq_length=16, prune_target_params=prune_target_params, prune_score_func="mmlu_1pct_bs32", + score_lower_bound=0.01, # exercise the score gate ss_channel_divisor=4, hparams_to_skip="num_attention_heads", top_k=1, @@ -138,6 +139,7 @@ def test_prune_minitron_vlm(tmp_path, num_gpus, create_teacher): seq_length=1024, prune_target_params=prune_target_params, prune_score_func="mmlu_1pct_bs32", + score_lower_bound=0.01, # exercise the score gate ss_channel_divisor=4, # Allow depth pruning (the primary param lever once hidden_size is fixed for VLMs). max_depth_pruning=0.6, diff --git a/tools/launcher/core.py b/tools/launcher/core.py index a35495e5795..47c93b48d45 100644 --- a/tools/launcher/core.py +++ b/tools/launcher/core.py @@ -51,17 +51,20 @@ def get_default_env(experiment_title=None): "SPECDEC_BENCH_S3_KEY_ID": os.getenv("SPECDEC_BENCH_S3_KEY_ID", ""), "SPECDEC_BENCH_S3_SECRET": os.getenv("SPECDEC_BENCH_S3_SECRET", ""), } + # HF_HOME / TRITON_CACHE_DIR default under the shared /{title} mount, but honor + # an env override so a user can point them at a personally-writable path (the + # shared cache is owned by the CI account and blocks other users' cache locks). slurm_env = { - "TRITON_CACHE_DIR": f"/{title}/triton-cache", - "HF_HOME": f"/{title}/hf-cache", + "TRITON_CACHE_DIR": os.getenv("TRITON_CACHE_DIR", f"/{title}/triton-cache"), + "HF_HOME": os.getenv("HF_HOME", f"/{title}/hf-cache"), "HF_TOKEN": os.getenv("HF_TOKEN", ""), "MLM_SKIP_INSTALL": "1", "LAUNCH_SCRIPT": "python", **specdec_s3, } local_env = { - "TRITON_CACHE_DIR": f"/{title}/triton-cache", - "HF_HOME": f"/{title}/hf-cache", + "TRITON_CACHE_DIR": os.getenv("TRITON_CACHE_DIR", f"/{title}/triton-cache"), + "HF_HOME": os.getenv("HF_HOME", f"/{title}/hf-cache"), "HF_TOKEN": os.getenv("HF_TOKEN", ""), "MLM_SKIP_INSTALL": "1", **specdec_s3, @@ -106,9 +109,21 @@ def register_factory(name, fn): @dataclass class SandboxTask: - """A single task with a script, slurm config, args, and environment.""" + """A single task with a script (or inline command), slurm config, args, and environment.""" script: str = None + # Inline shell command run instead of `script` (mutually exclusive; setting + # `args` too is rejected — put everything in the command). Lets one-liner jobs + # live in the YAML without a wrapper .sh. Must be a SINGLE line: the --yaml CLI + # layer rejects multi-line values, so YAMLs use a folded scalar (>-) and `&&`. + inline: str = None + # pip requirements installed in the container before the command runs + # (`pip install [-r reqs_file] [reqs] && `). `reqs` is a raw + # pip-install arg string (e.g. "transformers<5 fire"); `reqs_file` is a + # requirements.txt path relative to the run dir (e.g. + # modules/Model-Optimizer/examples/llm_eval/requirements.txt). + reqs: str = None + reqs_file: str = None slurm_config: object = None # Patched at runtime by set_slurm_config_type() args: list[str] = None environment: list[dict[str, str]] = None @@ -320,6 +335,12 @@ def _resolve(s): task.environment = {k: _resolve(v) for k, v in task.environment.items()} if task.args: task.args = [_resolve(a) for a in task.args] + if task.inline: + task.inline = _resolve(task.inline) + if task.reqs: + task.reqs = _resolve(task.reqs) + if task.reqs_file: + task.reqs_file = _resolve(task.reqs_file) # --------------------------------------------------------------------------- @@ -603,13 +624,16 @@ def build_docker_executor( f"{exp_title_src}:/{experiment_title}", ] + # Default to host uid:gid so artifacts aren't root-owned; docker_user="root" + # lets a job read root-only image paths (e.g. /opt/Megatron-Bridge in NeMo). + docker_user = getattr(slurm_config, "docker_user", None) or f"{os.getuid()}:{os.getgid()}" executor = run.DockerExecutor( num_gpus=-1, runtime="nvidia", ipc_mode="host", container_image=slurm_config.container, volumes=container_mounts, - additional_kwargs={"user": f"{os.getuid()}:{os.getgid()}", "entrypoint": ""}, + additional_kwargs={"user": docker_user, "entrypoint": ""}, packager=packager, ) return executor @@ -761,6 +785,12 @@ def run_jobs( continue task_name = f"{job_name}_{task_id}" task_args = [] if task.args is None else task.args + if bool(task.script) == bool(task.inline): + raise ValueError(f"{task_name}: set exactly one of `script` or `inline`.") + if task.inline and task_args: + raise ValueError( + f"{task_name}: `args` is only for `script`; put them in the `inline` command." + ) task_env = {} if task.environment is not None: @@ -803,7 +833,37 @@ def run_jobs( if job.allow_to_fail and hasattr(executor, "dependency_type"): executor.dependency_type = "afterany" - task_instance = run.Script(task.script, args=task_args, env=task_env) + # Optional reqs: pip-install before the command. reqs_file is a + # requirements.txt path; reqs is a raw arg string (shlex-quoted so + # < > = are literal, letting YAMLs write it unquoted). + reqs_prefix = "" + if task.reqs or task.reqs_file: + pkgs = ["-r", shlex.quote(task.reqs_file)] if task.reqs_file else [] + pkgs += [shlex.quote(tok) for tok in shlex.split(task.reqs or "")] + install = "python -m pip install " + " ".join(pkgs) + # On Slurm, srun runs this inline on every rank (ntasks_per_node), so install + # once on local rank 0 behind a filesystem barrier — concurrent pip on one node + # corrupts the env. The marker is relative to the per-task working dir + # (/nemo_run/code, fresh per task and shared across the node's ranks), so it is + # run-unique (no stale-marker reuse); rank 0 clears it first, and other ranks + # fail if it never appears (rank-0 install error). Local single-process runs + # just install as rank 0 and never wait. + marker = ".modelopt_launcher_reqs_done" + reqs_prefix = ( + f'if [ "${{SLURM_LOCALID:-0}}" -eq 0 ]; then rm -f {marker}; ' + f"{install} && touch {marker}; " + f"else for _ in $(seq 600); do [ -f {marker} ] && break; sleep 1; done; " + f"[ -f {marker} ]; fi && " + ) + if task.inline: + task_instance = run.Script(inline=reqs_prefix + task.inline, env=task_env) + elif reqs_prefix: # reqs + script path: wrap the bash call inline + # Quote the script path; task_args keep the launcher's shell-word-split + # convention (a "--flag value" item expands to two args), as run.Script does. + script_cmd = " ".join(["bash", shlex.quote(task.script), *task_args]) + task_instance = run.Script(inline=reqs_prefix + script_cmd, env=task_env) + else: + task_instance = run.Script(task.script, args=task_args, env=task_env) print(f"job {job_name} task {task_id} slurm_config: {task.slurm_config}") if dependency is None: diff --git a/tools/launcher/docs/configuration.md b/tools/launcher/docs/configuration.md index 5e2bc82108b..ad4703490f0 100644 --- a/tools/launcher/docs/configuration.md +++ b/tools/launcher/docs/configuration.md @@ -58,6 +58,70 @@ pipeline: gpus_per_node: 4 ``` +### Inline Command (no wrapper script) + +A task can carry an `inline` shell command instead of a `script:` path (no +wrapper `.sh` needed) — good for one-liner jobs like Megatron-Bridge's +`torchrun` scripts. `<>` is resolved; `args` must be omitted +(combining `inline` with non-empty `args` raises a `ValueError`). The +packaged repo is at `modules/Model-Optimizer/...` with the run dir as CWD. + +For the launcher (`torchrun` locally, `python`+`srun` on Slurm), reference +`$LAUNCH_SCRIPT` (plain `$VAR`, **not** `${...}` which collides with the config +loader) and set its local value in `environment`; on Slurm the launcher +overrides it to `python`, so set `ntasks_per_node = gpus_per_node`. + +> **`inline` must be a single line** — the CLI layer rejects multi-line values. +> Use a YAML **folded** scalar (`>-`), not a literal block (`|`) or `\` +> continuations; chain commands with `&&`. + +```yaml +job_name: Qwen3-8B_mbridge_prune +pipeline: + global_vars: + output_dir: /cicd/megatron-bridge + task_0: + environment: + - LAUNCH_SCRIPT: torchrun --nproc_per_node 2 + inline: >- + $LAUNCH_SCRIPT + modules/Model-Optimizer/examples/megatron_bridge/prune_minitron.py + --hf_model_name_or_path Qwen/Qwen3-8B + --pp_size 2 + --prune_target_params 6e9 + --output_hf_path <>/Qwen3-8B-Pruned-6B + 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: 2 + gpus_per_node: 2 +``` + +### Installing extra pip reqs (`reqs` / `reqs_file`) + +A task may pip-install **in the container before the command** +(`pip install [-r reqs_file] [reqs] && `): + +- `reqs` — a raw `pip install` argument string. Write specifiers unquoted; the + launcher shell-quotes each token, so `<` `>` `=` are safe. +- `reqs_file` — a `requirements.txt` path (relative to the run dir, i.e. under + `modules/Model-Optimizer/...`). + +```yaml + task_0: + reqs: "transformers<5" # or: "transformers<5 fire" for several packages + inline: >- + python .../prune_minitron.py ... + task_2: + reqs_file: modules/Model-Optimizer/examples/llm_eval/requirements.txt + inline: >- + python .../lm_eval_hf.py ... +``` + +Both work with `inline` and `script` tasks and resolve `<>`. + ### Multi-task Pipeline Tasks run sequentially — `task_1` starts only after `task_0` completes. diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yaml new file mode 100644 index 00000000000..f3402cbacf3 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yaml @@ -0,0 +1,56 @@ +# Nemotron-3-Nano-30B-A3B (MoE) pruning to 3B active via Megatron-Bridge (4 GPUs), then vLLM gen. +# +# Slurm: uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yaml --yes +# Local: uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yaml hf_local=/mnt/hf-local --yes + +# NOTE: sized for fast CI; bump for production, e.g. --calib_num_samples 1024 --seq_length 8192 --top_k 10. +# May need to reduce batch size if running out of memory at large seq_length. +job_name: Nemotron-3-Nano-30B-A3B_mbridge_prune +pipeline: + note: "Prune Nemotron-3-Nano-30B-A3B -> 3B active (Megatron-Bridge) with MMLU gate, then vLLM gen" + + global_vars: + # Per-run scratch (fresh cicd_ dir) so each run prunes fresh + output_dir: /scratchspace/Nemotron-3-Nano-30B-A3B-Pruned-A3.0B + + # 1) Prune to a pruned HF checkpoint (HF save needs transformers<5). + # --score_lower_bound fails the job if the pruned model's MMLU drops below the floor. + task_0: + reqs: "transformers<5" + environment: + - LAUNCH_SCRIPT: torchrun --nproc_per_node 4 + inline: >- + $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/prune_minitron.py + --hf_model_name_or_path nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + --trust_remote_code + --pp_size 4 + --calib_batch_size 8 + --calib_num_samples 256 + --seq_length 512 + --prune_target_active_params 3e9 + --prune_target_params 24e9 + --prune_score_func mmlu_10pct_bs32 + --max_width_pruning 0.30 + --max_depth_pruning 0.15 + --hparams_to_skip num_attention_heads + --top_k 5 + --score_lower_bound 0.50 + --output_hf_path <> + slurm_config: &sc + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.04 # 26.06 drops transformers<5, needed for the pruned-HF save + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + docker_user: root + ntasks_per_node: 4 + gpus_per_node: 4 + + # 2) vLLM sanity generation on the pruned checkpoint. + task_1: + inline: >- + python modules/Model-Optimizer/examples/megatron_bridge/generate_vllm.py + --model <> + --trust_remote_code + --tensor_parallel_size 4 + slurm_config: + <<: *sc + ntasks_per_node: 1 diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yaml new file mode 100644 index 00000000000..ba8a82681e5 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yaml @@ -0,0 +1,60 @@ +# Nemotron-3-Nano-30B-A3B FP8 quantization + unified-HF export via Megatron-Bridge (4 GPUs). +# +# Slurm: uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yaml --yes +# Local: uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yaml hf_local=/mnt/hf-local --yes + +# NOTE: sized for fast run; bump for production, e.g. --calib_num_samples 512 --seq_length 8192. +job_name: Nemotron-3-Nano-30B-A3B_mbridge_quantize +pipeline: + note: "FP8 PTQ Nemotron-3-Nano-30B-A3B (Megatron-Bridge), then unified-HF export" + + global_vars: + # Per-run scratch (fresh cicd_ dir) so each run quantizes fresh. + output_dir: /scratchspace + + # 1) FP8 Quantize and export it to a deployable unified-HF checkpoint. + task_0: + environment: + - LAUNCH_SCRIPT: torchrun --nproc_per_node 4 + inline: >- + $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/quantize.py + --hf_model_name_or_path nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + --trust_remote_code + --tp_size 4 + --quant_cfg MAMBA_MOE_FP8_CONSERVATIVE_CFG + --calib_batch_size 8 + --calib_num_samples 256 + --seq_length 512 + --skip_generate + --export_megatron_path <>/Nemotron-3-Nano-30B-A3B-FP8-megatron + && + $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/export_quantized_megatron_to_hf.py + --hf_model_name_or_path nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + --megatron_path <>/Nemotron-3-Nano-30B-A3B-FP8-megatron + --trust_remote_code + --pp_size 4 + --export_unified_hf_path <>/Nemotron-3-Nano-30B-A3B-FP8-hf + slurm_config: &sc + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.06 + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + docker_user: root + ntasks_per_node: 4 + gpus_per_node: 4 + + # 2) MMLU (10% sample) on the exported FP8 checkpoint via vLLM, gated on a lower bound. + task_1: + reqs_file: modules/Model-Optimizer/examples/llm_eval/requirements.txt + inline: >- + python modules/Model-Optimizer/examples/llm_eval/lm_eval_hf.py + --model vllm + --model_args pretrained=<>/Nemotron-3-Nano-30B-A3B-FP8-hf,tensor_parallel_size=4 + --trust_remote_code + --tasks mmlu + --limit 0.1 + --batch_size auto + --output_path /scratchspace/mmlu_results + --accuracy_lower_bound 0.69 + slurm_config: + <<: *sc + ntasks_per_node: 1 diff --git a/tools/launcher/slurm_config.py b/tools/launcher/slurm_config.py index 516789c680f..2647a7ba5ea 100644 --- a/tools/launcher/slurm_config.py +++ b/tools/launcher/slurm_config.py @@ -48,6 +48,9 @@ class SlurmConfig: requeue: bool = False nodes: int = 1 ntasks_per_node: int = 1 + # Docker-only: user for local `docker run` (e.g. "root" to read root-owned image + # paths like /opt/Megatron-Bridge). None -> host uid:gid. Ignored on Slurm. + docker_user: Optional[str] = None # None means omit GPU GRES entirely. Some clusters expose GPU nodes without # Slurm GRES, so requesting --gpus-per-node would make valid jobs fail. gpus_per_node: Optional[int] = 1 @@ -71,6 +74,7 @@ def slurm_factory( nodes: int = 1, ntasks_per_node: int = 1, gpus_per_node: Optional[int] = 1, + docker_user: Optional[str] = None, container: str = "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20", modelopt_install_path: str = "/usr/local/lib/python3.12/dist-packages/modelopt", container_mounts: list[str] = [ @@ -92,6 +96,7 @@ def slurm_factory( nodes=nodes, ntasks_per_node=ntasks_per_node, gpus_per_node=gpus_per_node, + docker_user=docker_user, container=container, modelopt_install_path=modelopt_install_path, container_mounts=container_mounts, diff --git a/tools/launcher/tests/test_core.py b/tools/launcher/tests/test_core.py index 5c0ec64ea2d..463e41f412a 100644 --- a/tools/launcher/tests/test_core.py +++ b/tools/launcher/tests/test_core.py @@ -204,7 +204,10 @@ class MockSlurmConfig: class TestGetDefaultEnv: """Tests for get_default_env utility.""" - def test_default_title(self): + def test_default_title(self, monkeypatch): + # get_default_env honors HF_HOME / TRITON_CACHE_DIR overrides, so isolate the ambient env. + monkeypatch.delenv("HF_HOME", raising=False) + monkeypatch.delenv("TRITON_CACHE_DIR", raising=False) slurm_env, local_env = get_default_env() assert slurm_env["TRITON_CACHE_DIR"] == "/cicd/triton-cache" assert slurm_env["HF_HOME"] == "/cicd/hf-cache" @@ -213,7 +216,9 @@ def test_default_title(self): assert local_env["TRITON_CACHE_DIR"] == "/cicd/triton-cache" assert "LAUNCH_SCRIPT" not in local_env - def test_custom_title(self): + def test_custom_title(self, monkeypatch): + monkeypatch.delenv("HF_HOME", raising=False) + monkeypatch.delenv("TRITON_CACHE_DIR", raising=False) slurm_env, local_env = get_default_env("modelopt") assert slurm_env["TRITON_CACHE_DIR"] == "/modelopt/triton-cache" assert slurm_env["HF_HOME"] == "/modelopt/hf-cache" diff --git a/tools/launcher/tests/test_docker_execution.py b/tools/launcher/tests/test_docker_execution.py index 01b418125c1..7b7b92850eb 100644 --- a/tools/launcher/tests/test_docker_execution.py +++ b/tools/launcher/tests/test_docker_execution.py @@ -128,6 +128,48 @@ def test_modelopt_mount(self, tmp_path): volumes = executor.volumes assert any("/custom/modelopt:/opt/modelopt" in v for v in volumes) + def test_docker_user_defaults_to_host_uid(self, tmp_path): + executor = build_docker_executor( + hf_local="/tmp/hf", + slurm_config=MagicMock( + local=False, + container="test:latest", + modelopt_install_path="/opt/modelopt", + container_mounts=None, + srun_args=None, + array=None, + docker_user=None, + ), + experiment_id="exp_u1", + job_dir=str(tmp_path / "experiments"), + task_name="task_0", + packager=MagicMock(), + modelopt_src_path="/tmp/modelopt", + experiment_title="cicd", + ) + assert executor.additional_kwargs["user"] == f"{os.getuid()}:{os.getgid()}" + + def test_docker_user_override(self, tmp_path): + executor = build_docker_executor( + hf_local="/tmp/hf", + slurm_config=MagicMock( + local=False, + container="test:latest", + modelopt_install_path="/opt/modelopt", + container_mounts=None, + srun_args=None, + array=None, + docker_user="root", + ), + experiment_id="exp_u2", + job_dir=str(tmp_path / "experiments"), + task_name="task_0", + packager=MagicMock(), + modelopt_src_path="/tmp/modelopt", + experiment_title="cicd", + ) + assert executor.additional_kwargs["user"] == "root" + def test_experiment_title_mount(self, tmp_path): job_dir = str(tmp_path / "experiments") executor = build_docker_executor( diff --git a/tools/launcher/tests/test_examples_resolve.py b/tools/launcher/tests/test_examples_resolve.py index b0fa16f37ad..5d0bc7c2138 100644 --- a/tools/launcher/tests/test_examples_resolve.py +++ b/tools/launcher/tests/test_examples_resolve.py @@ -56,7 +56,7 @@ def _tasks(cfg): for key, val in pipeline.items(): if key.startswith("task_") and isinstance(val, dict): yield key, val - elif "script" in cfg: + elif "script" in cfg or "inline" in cfg: yield "task", cfg @@ -76,15 +76,28 @@ def test_example_yaml_valid(path): for name, task in _tasks(cfg): script = task.get("script") - assert isinstance(script, str) and script.strip(), f"{path}:{name}: missing `script`" + inline = task.get("inline") + # A task runs either a `script:` wrapper or an `inline:` command — exactly one. + assert (isinstance(script, str) and script.strip()) or ( + isinstance(inline, str) and inline.strip() + ), f"{path}:{name}: task needs a `script` or `inline`" + assert not (script and inline), f"{path}:{name}: set only one of `script`/`inline`" + + # nemo-run's --yaml CLI layer rejects multi-line override values, so an + # inline command must stay single-line (use a folded scalar `>-`). + if inline: + assert "\n" not in inline.strip(), ( + f"{path}:{name}: `inline` must be single-line (use folded `>-`, chain with `&&`)" + ) # Scripts shipped with the launcher live under common/; verify the path is # real so a renamed/typo'd wrapper is caught at unit-test time. - first_token = script.split()[0] - if first_token.startswith("common/") and first_token not in _KNOWN_MISSING_SCRIPTS: - assert os.path.exists(os.path.join(_LAUNCHER_DIR, first_token)), ( - f"{path}:{name}: script not found: {first_token}" - ) + if script: + first_token = script.split()[0] + if first_token.startswith("common/") and first_token not in _KNOWN_MISSING_SCRIPTS: + assert os.path.exists(os.path.join(_LAUNCHER_DIR, first_token)), ( + f"{path}:{name}: script not found: {first_token}" + ) slurm_config = task.get("slurm_config") if slurm_config is not None: diff --git a/tools/launcher/tests/test_yaml_formats.py b/tools/launcher/tests/test_yaml_formats.py index 86a4863156f..9ba09550bb2 100644 --- a/tools/launcher/tests/test_yaml_formats.py +++ b/tools/launcher/tests/test_yaml_formats.py @@ -142,6 +142,33 @@ def test_global_vars_across_multiple_tasks(self, tmp_yaml): assert pipeline.tasks[0].environment == [{"HF_MODEL": "/hf-local/Qwen/Qwen3-8B"}] assert pipeline.tasks[1].environment == [{"HF_MODEL": "/hf-local/Qwen/Qwen3-8B"}] + def test_inline_task_and_global_vars(self): + """A task can use `inline` instead of `script`, with global_vars resolved.""" + task = SandboxTask0( + inline="torchrun prune_minitron.py --output_hf_path <>/pruned", + ) + pipeline = SandboxPipeline( + task_0=task, + global_vars=GlobalVariables(output_dir="/cicd/megatron-bridge"), + ) + assert pipeline.tasks[0].script is None + assert ( + pipeline.tasks[0].inline + == "torchrun prune_minitron.py --output_hf_path /cicd/megatron-bridge/pruned" + ) + + def test_task_reqs_fields(self): + """A task can carry `reqs` / `reqs_file`, both resolved for global_vars.""" + task = SandboxTask0( + reqs="transformers<5", + reqs_file="<>/requirements.txt", + inline="python eval.py --out <>", + ) + pipeline = SandboxPipeline(task_0=task, global_vars=GlobalVariables(output_dir="/cicd/out")) + assert pipeline.tasks[0].reqs == "transformers<5" + assert pipeline.tasks[0].reqs_file == "/cicd/out/requirements.txt" + assert pipeline.tasks[0].inline == "python eval.py --out /cicd/out" + class TestTestYamlFormat: """Tests for the test YAML format used by run_test_yaml.sh."""