Add Megatron-Bridge prune & quantize launcher pipelines

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,31 +240,100 @@ 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):
     """Move ModelOpt args from the argparse namespace into args.model_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).
         datasets.config.HF_DATASETS_TRUST_REMOTE_CODE = True
         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 <task>/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
+    # <model_name>/ 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)

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!")
 
 

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()

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")
 

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,

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] && <command>`). `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: