From fbce1d12e2c37276488371fa81bd96c52791e83c Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Wed, 29 Jul 2026 16:35:55 -0700 Subject: [PATCH 01/16] added downstream evaluation Signed-off-by: Grzegorz Karch --- CHANGELOG.rst | 1 + .../orchestration/adapters/post_mip.py | 2 +- .../puzzletron/orchestration/compiler.py | 6 +- .../puzzletron/orchestration/progress.py | 1 + modelopt/torch/puzzletron/post_mip/builtin.py | 12 +- .../torch/puzzletron/post_mip/reporting.py | 11 + modelopt/torch/puzzletron/post_mip/runner.py | 304 +++++++++++++++++- puzzletron_setup/bundle.py | 9 +- puzzletron_setup/v2/parallel_validation.py | 4 +- puzzletron_setup/v2/post_mip.py | 3 +- puzzletron_setup/v2/validation.py | 2 +- puzzletron_setup/v2/wizard.py | 95 +++++- puzzletron_setup/wizard.py | 124 ++++++- .../puzzletron/test_orchestration_compiler.py | 59 ++++ .../torch/puzzletron/test_post_mip_runner.py | 96 ++++++ .../torch/puzzletron/test_setup_bundle.py | 59 ++++ .../puzzletron/test_setup_v2_post_mip.py | 33 +- .../test_setup_v2_state_validation.py | 32 ++ 18 files changed, 830 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ad0d4acdfac..7383b652554 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,7 @@ Changelog **New Features** +- Add Puzzletron dynamic post-MIP downstream evaluation through ``lmms-eval`` with vLLM-backed checkpoint evaluation, setup-wizard topology/resource prompts, and an opt-in Nemotron-3 Nano 30B A3B BF16 example flow. - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred. - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``. diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index ba46a5b8178..90f1f745f9a 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py @@ -111,7 +111,7 @@ def plan(self, plan: CampaignPlan, node: StagePlanNode) -> WorkPlan: config = _node_config(plan, node.stage_id) node_type = str(config.get("type")) count = 1 if node_type in {"filter", "manual_filter"} else node.instances - if node_type == "evaluation": + if node_type in {"evaluation", "downstream_evaluation"}: available = _available_evaluation_candidates(plan, node.stage_id, config) if available is not None: if available < 1: diff --git a/modelopt/torch/puzzletron/orchestration/compiler.py b/modelopt/torch/puzzletron/orchestration/compiler.py index 0b12f049d59..5c76c3ce087 100644 --- a/modelopt/torch/puzzletron/orchestration/compiler.py +++ b/modelopt/torch/puzzletron/orchestration/compiler.py @@ -62,7 +62,6 @@ def _mapping(value: Any) -> dict[str, Any]: "aiperf": ExecutionStrategy.SHARDED, } - _POST_MIP_NODE_METADATA = { "filter": {"kind": "selector", "accepts": {"config", "checkpoint"}}, "manual_filter": {"kind": "selector", "accepts": {"config", "checkpoint"}}, @@ -87,7 +86,6 @@ def _mapping(value: Any) -> dict[str, Any]: "downstream_evaluation": { "kind": "evaluator", "accepts": {"checkpoint"}, - "implemented": False, }, } @@ -489,7 +487,7 @@ def compile_campaign_plan( parallel[key] = node_config[key] elif key in global_kd and key not in parallel: parallel[key] = global_kd[key] - if dynamic["node_type"] == "aiperf": + if dynamic["node_type"] in {"aiperf", "downstream_evaluation"}: topology = _mapping(node_config.get("topology")) topology_mesh = vllm_topology_to_mesh(topology) if override: @@ -498,7 +496,7 @@ def compile_campaign_plan( if ParallelMesh.from_mapping(overridden) != topology_mesh: raise ValueError( f"{stage_id} execution parallel override conflicts with " - "its AIPerf topology" + "its vLLM topology" ) mesh = topology_mesh else: diff --git a/modelopt/torch/puzzletron/orchestration/progress.py b/modelopt/torch/puzzletron/orchestration/progress.py index e395b9477ea..a9c943357e6 100644 --- a/modelopt/torch/puzzletron/orchestration/progress.py +++ b/modelopt/torch/puzzletron/orchestration/progress.py @@ -434,6 +434,7 @@ def _post_mip_progress( labels = { "evaluation": "evaluated", + "downstream_evaluation": "evaluated", "aiperf": "benchmarked", "global_kd": "distilled", "materialize": "materialized", diff --git a/modelopt/torch/puzzletron/post_mip/builtin.py b/modelopt/torch/puzzletron/post_mip/builtin.py index 8c62a1d5a11..050352af445 100644 --- a/modelopt/torch/puzzletron/post_mip/builtin.py +++ b/modelopt/torch/puzzletron/post_mip/builtin.py @@ -10,7 +10,12 @@ from .base import NodeCapabilities, NodeKind, PostMIPNode, post_mip_node from .filters import filter_metric_references, validate_filter_config from .records import ArtifactKind -from .reporting import render_aiperf_report, render_evaluation_report, render_global_kd_report +from .reporting import ( + render_aiperf_report, + render_downstream_evaluation_report, + render_evaluation_report, + render_global_kd_report, +) if TYPE_CHECKING: from collections.abc import Mapping @@ -115,6 +120,9 @@ class DownstreamEvaluationNode(PostMIPNode): NodeKind.EVALUATOR, frozenset({ArtifactKind.CHECKPOINT}), distributed=True, - implemented=False, default_strategy="sharded", ) + + @classmethod + def render_report(cls, node, payload): + return render_downstream_evaluation_report(str(payload["section_id"]), payload) diff --git a/modelopt/torch/puzzletron/post_mip/reporting.py b/modelopt/torch/puzzletron/post_mip/reporting.py index 86e3c11f44e..bc287990805 100644 --- a/modelopt/torch/puzzletron/post_mip/reporting.py +++ b/modelopt/torch/puzzletron/post_mip/reporting.py @@ -16,6 +16,7 @@ __all__ = [ "build_post_mip_report_payloads", "render_aiperf_report", + "render_downstream_evaluation_report", "render_evaluation_report", "render_global_kd_report", ] @@ -356,6 +357,16 @@ def render_aiperf_report(section_id: str, payload: Mapping[str, Any]) -> str: ) +def render_downstream_evaluation_report(section_id: str, payload: Mapping[str, Any]) -> str: + """Render lmms-eval task metrics for downstream-evaluation nodes.""" + + return render_evaluation_report(section_id, payload).replace( + "

Candidate evaluation

", + "

Downstream evaluation

", + 1, + ) + + def render_global_kd_report(section_id: str, payload: Mapping[str, Any]) -> str: """Render several candidate KD histories on shared, lineage-colored plots.""" diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 819cbbce010..3a1ee61afab 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -25,6 +25,7 @@ import math import os import subprocess +import sys import traceback import uuid from contextlib import contextmanager @@ -33,6 +34,7 @@ from typing import Any, Iterator, Mapping, Sequence from ..identity import canonicalize, stable_hash +from ..orchestration.mesh import normalize_vllm_topology from .base import CompiledPostMIPNode, NodeKind, compile_post_mip_flows from .filters import apply_filter from .records import ArtifactKind, CandidateLedger, CandidateSet, NodeObservation @@ -523,6 +525,292 @@ def _aiperf( } +_LMMS_EVAL_MODEL_ARG_FIELDS = frozenset( + { + "dtype", + "gpu_memory_utilization", + "max_model_len", + "trust_remote_code", + "tokenizer", + "tokenizer_mode", + "enforce_eager", + "limit_mm_per_prompt", + "reasoning_parser", + } +) + + +def _as_cli_bool(value: bool) -> str: + return "True" if value else "False" + + +def _as_lmms_eval_arg(value: Any) -> str: + if isinstance(value, bool): + return _as_cli_bool(value) + if isinstance(value, (int, float)) and not isinstance(value, bool): + return str(value) + if isinstance(value, (list, tuple, dict)): + return json.dumps(value, sort_keys=True, separators=(",", ":")) + return str(value) + + +def _join_cli_values(value: Any, *, path: str) -> str: + if isinstance(value, str): + text = value.strip() + if not text: + raise ValueError(f"{path} must not be empty") + return text + if not isinstance(value, Sequence): + raise TypeError(f"{path} must be a string or sequence") + values = [str(item).strip() for item in value] + if not values or any(not item for item in values): + raise ValueError(f"{path} must contain at least one non-empty value") + return ",".join(values) + + +def _model_arg_string(values: Mapping[str, Any]) -> str: + parts = [] + for key, value in values.items(): + if value is None: + continue + key_text = str(key).strip() + if not key_text or "," in key_text or "=" in key_text: + raise ValueError(f"invalid lmms-eval model_args key: {key!r}") + rendered = _as_lmms_eval_arg(value) + if "," in rendered: + raise ValueError( + f"lmms-eval model_args value for {key_text!r} contains a comma; " + "provide model_args as a preformatted string instead" + ) + parts.append(f"{key_text}={rendered}") + if not parts: + raise ValueError("lmms-eval model_args must contain at least the checkpoint path") + return ",".join(parts) + + +def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> str: + raw = settings.get("model_args") + checkpoint_arg = str(settings.get("checkpoint_arg", "model")) + topology = dict(settings.get("topology") or {}) + canonical_topology = normalize_vllm_topology(topology) if topology else {} + derived = { + checkpoint_arg: checkpoint, + } + if canonical_topology: + derived.update( + { + "tensor_parallel_size": canonical_topology["tp"], + "pipeline_parallel_size": canonical_topology["pp"], + "data_parallel_size": canonical_topology["dp"], + "enable_expert_parallel": canonical_topology["enable_expert_parallel"], + "distributed_executor_backend": canonical_topology[ + "distributed_executor_backend" + ], + } + ) + for key in _LMMS_EVAL_MODEL_ARG_FIELDS: + if key in settings: + derived[key] = settings[key] + + if isinstance(raw, str): + prefix = raw.strip().strip(",") + suffix = _model_arg_string(derived) + return ",".join(part for part in (prefix, suffix) if part) + if raw is not None and not isinstance(raw, Mapping): + raise TypeError("downstream_evaluation.config.model_args must be a mapping or string") + merged = dict(raw or {}) + for key, value in derived.items(): + merged.setdefault(key, value) + return _model_arg_string(merged) + + +def _command_prefix(settings: Mapping[str, Any]) -> list[str]: + raw = settings.get("command_prefix") + if raw is None: + return [sys.executable, "-m", "lmms_eval"] + if isinstance(raw, str): + values = [raw] + else: + values = [str(item) for item in raw] + if not values or any(not value for value in values): + raise ValueError("downstream_evaluation.config.command_prefix must not be empty") + return values + + +def _lmms_eval_command( + settings: Mapping[str, Any], + *, + checkpoint: str, + output_path: Path, +) -> tuple[list[str], dict[str, str], float | None]: + """Build a deterministic lmms-eval CLI invocation for one realized checkpoint.""" + + tasks = _join_cli_values(settings.get("tasks"), path="downstream_evaluation.config.tasks") + argv = [ + *_command_prefix(settings), + "--model", + str(settings.get("model", "vllm")), + "--model_args", + _merge_lmms_eval_model_args(settings, checkpoint), + "--tasks", + tasks, + "--batch_size", + str(settings.get("batch_size", 1)), + "--output_path", + str(output_path), + ] + optional_fields = { + "limit": "--limit", + "num_fewshot": "--num_fewshot", + "seed": "--seed", + "verbosity": "--verbosity", + "device": "--device", + "use_cache": "--use_cache", + } + for key, flag in optional_fields.items(): + value = settings.get(key) + if value is not None: + argv.extend([flag, str(value)]) + if settings.get("gen_kwargs") is not None: + argv.extend( + [ + "--gen_kwargs", + ( + settings["gen_kwargs"] + if isinstance(settings["gen_kwargs"], str) + else _model_arg_string(dict(settings["gen_kwargs"])) + ), + ] + ) + if bool(settings.get("log_samples", False)): + argv.append("--log_samples") + argv.extend(str(item) for item in settings.get("extra_args") or ()) + + env = os.environ.copy() + for key, value in dict(settings.get("env") or {}).items(): + if value is not None: + env[str(key)] = str(value) + if settings.get("cache_dir") is not None: + env.setdefault("LMMS_EVAL_HOME", str(settings["cache_dir"])) + timeout = settings.get("timeout_seconds", settings.get("timeout")) + return argv, env, (float(timeout) if timeout is not None else None) + + +def _metric_key(value: Any) -> str: + return ( + str(value) + .strip() + .replace(" ", "_") + .replace(",", "_") + .replace("/", "_") + .replace("\\", "_") + ) + + +def _flatten_lmms_eval_metrics(payload: Mapping[str, Any]) -> dict[str, float]: + results = payload.get("results") + if not isinstance(results, Mapping): + return {} + metrics = {} + for task_name, task_payload in results.items(): + if not isinstance(task_payload, Mapping): + continue + for metric_name, value in task_payload.items(): + if ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and math.isfinite(value) + ): + metrics[f"{_metric_key(task_name)}.{_metric_key(metric_name)}"] = float(value) + return metrics + + +def _lmms_eval_result_payload(output_path: Path) -> tuple[dict[str, Any], Path]: + candidates = [] + for path in sorted(output_path.rglob("*.json")): + try: + payload = json.loads(path.read_text()) + except (OSError, ValueError): + continue + if isinstance(payload, Mapping) and isinstance(payload.get("results"), Mapping): + candidates.append((path.stat().st_mtime_ns, path, dict(payload))) + if not candidates: + raise FileNotFoundError(f"lmms-eval wrote no JSON results below {output_path}") + _mtime, path, payload = max(candidates, key=lambda item: item[0]) + return payload, path + + +def _downstream_evaluation( + config: dict[str, Any], + node: CompiledPostMIPNode, + source, + execution_identity: str, +) -> dict[str, Any]: + if source.artifact_kind is not ArtifactKind.CHECKPOINT: + raise ValueError("downstream_evaluation requires materialized checkpoint artifacts") + settings = dict(node.config.get("config") or {}) + output_root = ( + _execution_root(config, node, execution_identity) + / "raw" + / source.architecture_id + / "lmms_eval" + ) + output = output_root / f"attempt_{uuid.uuid4().hex}" + output.mkdir(parents=True, exist_ok=True) + argv, env, timeout = _lmms_eval_command( + settings, + checkpoint=str(source.artifact["checkpoint"]), + output_path=output, + ) + command_path = output / "command.json" + _atomic_json( + command_path, + { + "argv": argv, + "env_overrides": sorted(str(key) for key in dict(settings.get("env") or {})), + "timeout": timeout, + }, + ) + # Campaign config controls the executable and arguments, but subprocess receives + # an argv list directly; no shell parsing is involved. + result = subprocess.run( + argv, + cwd=str(output), + env=env, + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + if result.returncode: + detail = (result.stderr or result.stdout).strip().splitlines() + tail = "\n".join(detail[-20:]) + raise RuntimeError( + f"lmms-eval failed with exit code {result.returncode}" + + (f": {tail}" if tail else "") + ) + payload, result_path = _lmms_eval_result_payload(output) + metrics = _flatten_lmms_eval_metrics(payload) + if not metrics: + raise RuntimeError(f"lmms-eval result has no numeric task metrics: {result_path}") + summary_path = output / "summary.json" + _atomic_json( + summary_path, + { + "architecture_id": source.architecture_id, + "checkpoint": source.artifact["checkpoint"], + "metrics": metrics, + "result_path": str(result_path), + }, + ) + return { + "metrics": metrics, + "result_path": str(summary_path), + "raw_result_path": str(result_path), + "command_path": str(command_path), + } + + def _post_mip_kd_settings( config: Mapping[str, Any], node_settings: Mapping[str, Any], @@ -582,6 +870,8 @@ def _run_candidate( result = _evaluate(config, node, source, execution_identity) elif node.node_type == "aiperf": result = _aiperf(config, node, source, execution_identity) + elif node.node_type == "downstream_evaluation": + result = _downstream_evaluation(config, node, source, execution_identity) elif node.node_type == "global_kd": result = _global_kd(config, node, source, execution_identity) else: @@ -649,7 +939,7 @@ def run_post_mip_node_shard( row = _run_candidate(config, node, ledger, revision_id, execution_identity) row = {**row, "execution_identity": execution_identity} except Exception as error: - timed_out = node.node_type == "aiperf" and isinstance( + timed_out = node.node_type in {"aiperf", "downstream_evaluation"} and isinstance( error, (subprocess.TimeoutExpired, TimeoutError) ) row = { @@ -661,12 +951,14 @@ def run_post_mip_node_shard( **_exception_diagnostics(error), } if timed_out: - timeout_field = ( - "benchmark_timeout" - if isinstance(error, subprocess.TimeoutExpired) - else "readiness_timeout" + timeout_field = "benchmark_timeout" + if node.node_type == "downstream_evaluation": + timeout_field = "timeout_seconds" + elif not isinstance(error, subprocess.TimeoutExpired): + timeout_field = "readiness_timeout" + default_timeout = 3600 if node.node_type == "downstream_evaluation" else ( + 600 if timeout_field == "benchmark_timeout" else 1200 ) - default_timeout = 600 if timeout_field == "benchmark_timeout" else 1200 row["timeout_seconds"] = float( getattr(error, "timeout", None) or (node.config.get("config") or {}).get(timeout_field, default_timeout) diff --git a/puzzletron_setup/bundle.py b/puzzletron_setup/bundle.py index ba091af39a8..3ee3e8f0936 100644 --- a/puzzletron_setup/bundle.py +++ b/puzzletron_setup/bundle.py @@ -346,6 +346,13 @@ def _post_mip_flows( if smoke: config["minimum_request_count"] = 4 config["requests_per_concurrency"] = 1 + elif node_type == "downstream_evaluation": + config.setdefault("model", "vllm") + config.setdefault("batch_size", 1) + config.setdefault("log_samples", True) + config.setdefault("topology", deepcopy(dict(default_serving_topology))) + if smoke: + config["limit"] = min(int(config.get("limit", 8) or 8), 8) elif node_type == "global_kd": config.setdefault("automodel", {})["parallel"] = _parallel(global_kd_mesh) config["local_batch_size"] = _aligned_batch_size( @@ -816,7 +823,7 @@ def _dynamic_stage_entries( entry.update(resource="cpu", partition=cpu_partition) if node_type == "evaluation": entry["parallel"] = dict(common) - elif node_type == "aiperf": + elif node_type in {"aiperf", "downstream_evaluation"}: config = _mapping(node.get("config")) entry["parallel"] = _serving_parallel(_mapping(config.get("topology"))) elif node_type == "materialize": diff --git a/puzzletron_setup/v2/parallel_validation.py b/puzzletron_setup/v2/parallel_validation.py index 0ce29e95297..c22058fd9f5 100644 --- a/puzzletron_setup/v2/parallel_validation.py +++ b/puzzletron_setup/v2/parallel_validation.py @@ -50,7 +50,9 @@ "vllm_stats", } ) -_CANDIDATE_POST_MIP_TYPES = frozenset({"evaluation", "global_kd", "aiperf"}) +_CANDIDATE_POST_MIP_TYPES = frozenset( + {"evaluation", "global_kd", "aiperf", "downstream_evaluation"} +) @dataclass(frozen=True) diff --git a/puzzletron_setup/v2/post_mip.py b/puzzletron_setup/v2/post_mip.py index 14d34f6b8dc..067dc18b459 100644 --- a/puzzletron_setup/v2/post_mip.py +++ b/puzzletron_setup/v2/post_mip.py @@ -29,9 +29,10 @@ "materialize", "evaluation", "aiperf", + "downstream_evaluation", "global_kd", ) -RESERVED_NODE_TYPES = ("ptq", "downstream_evaluation") +RESERVED_NODE_TYPES = ("ptq",) @dataclass(frozen=True) diff --git a/puzzletron_setup/v2/validation.py b/puzzletron_setup/v2/validation.py index f5c8dff0b4a..d0cf216d693 100644 --- a/puzzletron_setup/v2/validation.py +++ b/puzzletron_setup/v2/validation.py @@ -298,7 +298,7 @@ def validate_state(state: WizardState) -> tuple[ValidationIssue, ...]: for stage_id, node in post_mip_nodes.items(): node_type = str(node.get("type", "")) config = _mapping(node.get("config")) - if node_type == "aiperf": + if node_type in {"aiperf", "downstream_evaluation"}: topology = _mapping(config.get("topology")) if topology: issues.extend( diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index 67ca924aabe..ded6bbbff4c 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -3811,7 +3811,7 @@ def post_mip_section(session: WizardSession, resolver: DefaultsResolver, context strategy=_post_mip_strategy(node), batch=1, ) - elif node.node_type == "aiperf": + elif node.node_type in {"aiperf", "downstream_evaluation"}: node_preview["resources"] = { "instances": int(session.state.get_field("infrastructure.gpus_per_node", 8)), "topology": node.config.get("topology", {}), @@ -3908,7 +3908,7 @@ def post_mip_section(session: WizardSession, resolver: DefaultsResolver, context "aiperf", "global_kd", ("PTQ — unavailable", "unavailable"), - ("Downstream evaluation — unavailable", "unavailable"), + "downstream_evaluation", ], default="evaluation", ) @@ -3982,6 +3982,18 @@ def post_mip_section(session: WizardSession, resolver: DefaultsResolver, context "concurrency": list(configured["concurrency"]), "benchmark_timeout": 900, } + elif node_type == "downstream_evaluation": + configured = _downstream_evaluation_setting_prompt( + session, + f"post_mip.{run_id}.{node_id}", + {}, + inventory=context["model"].inventory, + pruning=_mapping_copy(_pruning_payload(session.state)), + stage_id=f"post.{run_id}.{node_id}", + ) + if configured is BACK: + return False + config = configured elif node_type == "global_kd": max_steps = session.integer( f"post_mip.{run_id}.{node_id}.max_steps", @@ -4249,6 +4261,83 @@ def validate_concurrency(value: str) -> bool | str: return values +def _downstream_evaluation_setting_prompt( + session: WizardSession, + prefix: str, + defaults: Mapping[str, Any], + *, + inventory: Any, + pruning: Mapping[str, Any], + stage_id: str, +) -> Any: + """Ask lmms-eval task settings and the vLLM topology used to run them.""" + + def validate_tasks(value: str) -> bool | str: + tasks = [item.strip() for item in value.split(",") if item.strip()] + return True if tasks else "Enter at least one lmms-eval task." + + raw_default_tasks = defaults.get("tasks", ("ifeval", "gsm8k")) + default_tasks = ( + str(raw_default_tasks) + if isinstance(raw_default_tasks, str) + else ",".join(str(item) for item in raw_default_tasks) + ) + default_model_args = _mapping_copy(defaults.get("model_args")) + tasks = session.text( + f"{prefix}.tasks", + "lmms-eval tasks (comma-separated):", + default=default_tasks, + validate=validate_tasks, + ) + if tasks is BACK: + return BACK + limit = session.integer( + f"{prefix}.limit", + "lmms-eval sample limit:", + default=int(defaults.get("limit", 128)), + minimum=1, + ) + batch_size = session.integer( + f"{prefix}.batch_size", + "lmms-eval batch size:", + default=int(defaults.get("batch_size", 1)), + minimum=1, + ) + timeout = session.integer( + f"{prefix}.timeout_seconds", + "Per-candidate lmms-eval timeout (seconds):", + default=int(defaults.get("timeout_seconds", 3600)), + minimum=1, + ) + if BACK in (limit, batch_size, timeout): + return BACK + topology = _vllm_topology_prompt( + session, + f"{prefix}.topology", + _mapping_copy(defaults.get("topology")), + inventory=inventory, + pruning=pruning, + stage_id=stage_id, + label_prefix="lmms-eval vLLM", + ) + if topology is BACK: + return BACK + return { + "model": str(defaults.get("model", "vllm")), + "tasks": [item.strip() for item in str(tasks).split(",") if item.strip()], + "limit": int(limit), + "batch_size": int(batch_size), + "log_samples": bool(defaults.get("log_samples", True)), + "topology": topology, + "model_args": { + **default_model_args, + "dtype": default_model_args.get("dtype", "bfloat16"), + "gpu_memory_utilization": default_model_args.get("gpu_memory_utilization", 0.85), + }, + "timeout_seconds": int(timeout), + } + + def _configure_dynamic_resources( session: WizardSession, editor: PostMIPFlowEditor, @@ -4301,7 +4390,7 @@ def _configure_dynamic_resources( "resource": "gpu", "gpus_per_node": gpus_per_node, } - if node.node_type == "aiperf": + if node.node_type in {"aiperf", "downstream_evaluation"}: topology = _mapping_copy(node.config.get("topology")) allocation_mesh = vllm_topology_to_mesh(topology) entry["parallel"] = { diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index a9e153e5260..2750be52f6f 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -733,6 +733,114 @@ def _ask_aiperf_config( return config +def _ask_downstream_evaluation_config( + prompts: PromptSession, + *, + detailed: bool, + moe: bool, + defaults: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Ask for lmms-eval task and vLLM settings.""" + + defaults = defaults or {} + tasks = prompts.text( + "lmms-eval tasks (comma-separated):", + default=str(defaults.get("tasks", "ifeval,gsm8k")), + ) + limit = prompts.integer( + "lmms-eval sample limit:", + default=int(defaults.get("limit", 128)), + ) + batch_size = prompts.integer( + "lmms-eval batch size:", + default=int(defaults.get("batch_size", 1)), + ) + topology_defaults = dict(defaults.get("topology") or {}) + checkpoint = prompts.checkpoint() + while True: + tp = prompts.integer( + "lmms-eval vLLM tensor parallel (TP):", + default=int(topology_defaults.get("tensor_parallel_size", 1)), + ) + pp = prompts.integer( + "lmms-eval vLLM pipeline parallel (PP):", + default=int(topology_defaults.get("pipeline_parallel_size", 1)), + ) + dp = prompts.integer( + "lmms-eval vLLM data parallel (DP):", + default=int(topology_defaults.get("data_parallel_size", 1)), + ) + prefill_cp = prompts.integer( + "lmms-eval vLLM prefill context parallel (CP):", + default=int(topology_defaults.get("prefill_context_parallel_size", 1)), + ) + decode_cp = prompts.integer( + "lmms-eval vLLM decode context parallel (CP):", + default=int(topology_defaults.get("decode_context_parallel_size", 1)), + ) + enable_expert_parallel = ( + prompts.confirm( + ( + "Enable lmms-eval vLLM expert parallelism? " + "vLLM effective EP is TP * DP." + ), + default=bool( + topology_defaults.get("enable_expert_parallel") + or int(topology_defaults.get("expert_parallel_size", 1)) > 1 + ), + ) + if moe + else False + ) + dimensions = { + "TP": tp, + "PP": pp, + "DP": dp, + "prefill CP": prefill_cp, + "decode CP": decode_cp, + } + error = None + if any(int(value) < 1 for value in dimensions.values()): + error = f"lmms-eval vLLM parallel dimensions must be positive: {dimensions}" + elif decode_cp > tp or tp % decode_cp: + error = f"lmms-eval vLLM decode CP={decode_cp} must divide TP={tp}." + if error is None: + break + print(error) + prompts.rewind(checkpoint) + topology = { + "tensor_parallel_size": tp, + "pipeline_parallel_size": pp, + "prefill_context_parallel_size": prefill_cp, + "decode_context_parallel_size": decode_cp, + "data_parallel_size": dp, + "enable_expert_parallel": bool(enable_expert_parallel), + "distributed_executor_backend": "mp", + "gpu_group_size": tp * pp * prefill_cp * dp, + } + timeout = ( + prompts.integer( + "Per-candidate lmms-eval timeout (seconds):", + default=int(defaults.get("timeout_seconds", 3600)), + ) + if detailed + else int(defaults.get("timeout_seconds", 3600)) + ) + return { + "model": "vllm", + "tasks": [item.strip() for item in str(tasks).split(",") if item.strip()], + "limit": int(limit), + "batch_size": int(batch_size), + "log_samples": True, + "topology": topology, + "model_args": { + "dtype": "bfloat16", + "gpu_memory_utilization": 0.85, + }, + "timeout_seconds": int(timeout), + } + + def _default_flow( run_id: str, run: Mapping[str, Any], @@ -916,7 +1024,7 @@ def _custom_flow( "global_kd", "manual_filter", ("PTQ (reserved; not executable yet)", "ptq"), - ("Downstream evaluation (reserved; not executable yet)", "downstream_evaluation"), + "downstream_evaluation", ], default="filter", ) @@ -963,9 +1071,16 @@ def _custom_flow( runtime=runtime, ) available_metrics.append(f"{node_id}.request_throughput") + elif node_type == "downstream_evaluation": + node["config"] = _ask_downstream_evaluation_config( + prompts, + detailed=detailed, + moe=moe, + ) + available_metrics.append(f"{node_id}.gsm8k.exact_match") elif node_type == "global_kd": node["config"] = {"max_steps": prompts.integer("Global KD steps:", default=128)} - elif node_type in {"ptq", "downstream_evaluation"}: + elif node_type == "ptq": print( f"{node_type} records the reserved interface, but current orchestration " "validation will report it as unimplemented." @@ -1164,6 +1279,11 @@ def post_mip_gpus_per_instance(node_type: str, default: int) -> int: post_mip_gpus_per_instance("aiperf", 1), post_mip_instances("aiperf", sharded_workers, sharded_workers), ), + ( + "downstream eval", + post_mip_gpus_per_instance("downstream_evaluation", 1), + post_mip_instances("downstream_evaluation", sharded_workers, sharded_workers), + ), ( "evaluation", _mesh_product(common), diff --git a/tests/unit/torch/puzzletron/test_orchestration_compiler.py b/tests/unit/torch/puzzletron/test_orchestration_compiler.py index f49d2253891..b6efc1a6fd9 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_compiler.py +++ b/tests/unit/torch/puzzletron/test_orchestration_compiler.py @@ -194,3 +194,62 @@ def test_post_mip_compiler_topologically_orders_serialized_nodes() -> None: stages = _post_mip_stage_metadata(config) assert [stage["node_id"] for stage in stages] == ["initial", "final_eval", "best"] + + +def test_compile_campaign_plan_allocates_downstream_evaluation_from_vllm_topology( + tmp_configs, +) -> None: + experiment_path, runner_path, execution_path = tmp_configs + experiment = yaml.safe_load(experiment_path.read_text()) + experiment.update( + { + "mip": {"runs": {"runtime": {}}}, + "post_mip": { + "flows": { + "runtime": { + "source": {"run": "runtime"}, + "nodes": { + "materialized": {"type": "materialize"}, + "lmms_eval": { + "type": "downstream_evaluation", + "input": "materialized", + "config": { + "tasks": ["ifeval"], + "topology": { + "tensor_parallel_size": 4, + "pipeline_parallel_size": 2, + "data_parallel_size": 1, + "prefill_context_parallel_size": 1, + "decode_context_parallel_size": 1, + "enable_expert_parallel": False, + "gpu_group_size": 8, + }, + }, + }, + }, + } + } + }, + } + ) + experiment_path.write_text(yaml.safe_dump(experiment)) + execution = yaml.safe_load(execution_path.read_text()) + execution["execution"]["stages"]["post.runtime.lmms_eval"] = { + "strategy": "sharded", + "instances": 2, + } + execution_path.write_text(yaml.safe_dump(execution)) + + plan = compile_campaign_plan( + experiment_config_path=experiment_path, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + stage_filter="post.runtime.lmms_eval", + ) + node = plan.stages[0] + + assert node.stage_id == "post.runtime.lmms_eval" + assert node.parents == ("post.runtime.materialized",) + assert node.gpus_per_instance == 8 + assert node.instances == 2 + assert node.nodes == 2 diff --git a/tests/unit/torch/puzzletron/test_post_mip_runner.py b/tests/unit/torch/puzzletron/test_post_mip_runner.py index aa12508ecb6..3ea3f722468 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -16,12 +16,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import json +import subprocess from pathlib import Path from types import SimpleNamespace from omegaconf import OmegaConf from modelopt.torch.puzzletron.post_mip import runner +from modelopt.torch.puzzletron.post_mip.records import ArtifactKind from modelopt.torch.puzzletron.post_mip.runner import ( _exception_diagnostics, _needs_puzzletron_process_group, @@ -184,3 +187,96 @@ def fake_run_aiperf_sweep(checkpoint, **settings): assert "requests_per_concurrency" not in captured assert "best_selection_mode" not in captured assert result["metrics"] == {} + + +def test_lmms_eval_command_maps_checkpoint_and_vllm_topology(tmp_path): + argv, env, timeout = runner._lmms_eval_command( + { + "command_prefix": ["python", "-m", "lmms_eval"], + "tasks": ["ifeval", "gsm8k"], + "batch_size": 2, + "limit": 8, + "cache_dir": tmp_path / "cache", + "timeout_seconds": 123, + "topology": { + "tensor_parallel_size": 4, + "pipeline_parallel_size": 2, + "data_parallel_size": 1, + "prefill_context_parallel_size": 1, + "decode_context_parallel_size": 1, + "enable_expert_parallel": False, + "gpu_group_size": 8, + }, + "model_args": {"dtype": "bfloat16"}, + }, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + + model_args = argv[argv.index("--model_args") + 1] + assert argv[:5] == ["python", "-m", "lmms_eval", "--model", "vllm"] + assert argv[argv.index("--tasks") + 1] == "ifeval,gsm8k" + assert argv[argv.index("--batch_size") + 1] == "2" + assert argv[argv.index("--limit") + 1] == "8" + assert "model=/ckpts/candidate" in model_args + assert "tensor_parallel_size=4" in model_args + assert "pipeline_parallel_size=2" in model_args + assert "gpu_group_size" not in model_args + assert env["LMMS_EVAL_HOME"] == str(tmp_path / "cache") + assert timeout == 123 + + +def test_downstream_evaluation_runs_lmms_eval_and_flattens_metrics(monkeypatch, tmp_path): + captured = {} + + def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): + del env, capture_output, text, timeout, check + captured["argv"] = argv + output = Path(cwd) / "nested" + output.mkdir(parents=True) + (output / "results.json").write_text( + json.dumps( + { + "results": { + "ifeval": {"prompt_level_strict_acc,none": 0.5}, + "gsm8k": {"exact_match,strict-match": 0.75}, + } + } + ) + ) + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(runner.subprocess, "run", fake_run) + node = SimpleNamespace( + node_id="lmms_eval", + flow_id="runtime", + stage_id="post.runtime.lmms_eval", + config={ + "config": { + "command_prefix": ["python", "-m", "lmms_eval"], + "tasks": ["ifeval", "gsm8k"], + "limit": 4, + "topology": {"gpu_group_size": 1}, + } + }, + ) + source = SimpleNamespace( + architecture_id="architecture", + artifact_kind=ArtifactKind.CHECKPOINT, + artifact={"checkpoint": str(tmp_path / "checkpoint")}, + ) + + result = runner._downstream_evaluation( + {"puzzle_dir": str(tmp_path)}, + node, + source, + "execution", + ) + + assert captured["argv"][:3] == ["python", "-m", "lmms_eval"] + assert result["metrics"] == { + "gsm8k.exact_match_strict-match": 0.75, + "ifeval.prompt_level_strict_acc_none": 0.5, + } + assert Path(result["result_path"]).is_file() + assert Path(result["raw_result_path"]).name == "results.json" diff --git a/tests/unit/torch/puzzletron/test_setup_bundle.py b/tests/unit/torch/puzzletron/test_setup_bundle.py index f453aa55d49..51f2862282e 100644 --- a/tests/unit/torch/puzzletron/test_setup_bundle.py +++ b/tests/unit/torch/puzzletron/test_setup_bundle.py @@ -549,6 +549,65 @@ def test_render_execution_uses_common_mesh_for_post_mip_evaluation_only() -> Non assert execution["post.run.materialized"]["instances"] == 1 +def test_render_execution_uses_vllm_mesh_for_post_mip_downstream_evaluation() -> None: + state = { + "answers": { + "infrastructure": { + "gpus_per_node": 8, + "workers": {"pool": 8, "sharded": 8}, + "runner": {"slurm": {}}, + "meshes": { + "common": {"tp": 1, "cp": 1, "pp": 1, "dp_shard": 2, "ep": 1}, + "bypass": {"tp": 1, "cp": 1, "pp": 1, "dp_shard": 1, "ep": 1}, + "global_kd": {"tp": 1, "cp": 1, "pp": 1, "dp_shard": 1, "ep": 1}, + }, + } + }, + } + experiment = { + "embedding_pruning": {"widths": []}, + "vllm_stats": {"runtime_stats": {"topology": {"gpu_group_size": 1}}}, + "post_mip": { + "flows": { + "run": { + "nodes": { + "materialized": {"type": "materialize"}, + "lmms_eval": { + "type": "downstream_evaluation", + "input": "materialized", + "config": { + "topology": { + "tensor_parallel_size": 4, + "pipeline_parallel_size": 2, + "data_parallel_size": 1, + "prefill_context_parallel_size": 1, + "decode_context_parallel_size": 1, + "enable_expert_parallel": False, + "gpu_group_size": 8, + } + }, + }, + } + } + } + }, + } + + stages = render_execution(state, experiment, "production")["execution"]["stages"] + + assert stages["post.run.lmms_eval"]["strategy"] == "sharded" + assert stages["post.run.lmms_eval"]["instances"] == 8 + assert stages["post.run.lmms_eval"]["parallel"] == { + "tp": 4, + "cp": 1, + "pp": 2, + "ep": 1, + "dp_shard": 1, + "dp_replicate": 1, + "sequence_parallel": False, + } + + def test_render_execution_caps_post_mip_workers_at_upstream_top_k() -> None: common = { "tp": 1, diff --git a/tests/unit/torch/puzzletron/test_setup_v2_post_mip.py b/tests/unit/torch/puzzletron/test_setup_v2_post_mip.py index 459011a3bc3..247f0e85ebf 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_post_mip.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_post_mip.py @@ -1,7 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from puzzletron_setup.v2.post_mip import recommended_flow +from collections import OrderedDict + +from puzzletron_setup.v2.post_mip import FlowDraft, NodeDraft, PostMIPFlowEditor, recommended_flow def test_recommended_flow_propagates_aiperf_sweep_selection_mode(): @@ -45,3 +47,32 @@ def test_recommended_flow_accepts_a_single_concurrency_value(): assert flow.nodes["serving"].config["concurrency"] == [2] assert flow.nodes["fastest"].selector["best_selection_mode"] == "individual_best" + + +def test_downstream_evaluation_node_is_configurable_after_materialization(): + flow = FlowDraft( + "runtime", + "runtime", + nodes=OrderedDict( + ( + ("materialized", NodeDraft("materialized", "materialize")), + ( + "lmms_eval", + NodeDraft( + "lmms_eval", + "downstream_evaluation", + input_id="materialized", + config={"tasks": ["ifeval"]}, + ), + ), + ) + ), + ) + editor = PostMIPFlowEditor({"runtime": {}}) + editor.add_flow(flow) + + review = editor.review("runtime") + + assert review.node_order == ("materialized", "lmms_eval") + assert review.parents["lmms_eval"] == ("materialized",) + assert review.artifacts["lmms_eval"] == "checkpoint" diff --git a/tests/unit/torch/puzzletron/test_setup_v2_state_validation.py b/tests/unit/torch/puzzletron/test_setup_v2_state_validation.py index 049905630a8..78be65d9916 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_state_validation.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_state_validation.py @@ -181,6 +181,38 @@ def test_persisted_post_mip_aiperf_topology_is_candidate_checked(tmp_path): assert "expert counts [48, 64]" in messages +def test_persisted_post_mip_downstream_eval_topology_is_candidate_checked(tmp_path): + state = _state(tmp_path) + topology = { + "tensor_parallel_size": 4, + "pipeline_parallel_size": 1, + "data_parallel_size": 8, + "prefill_context_parallel_size": 1, + "decode_context_parallel_size": 1, + "enable_expert_parallel": True, + "gpu_group_size": 32, + } + state.set_collection( + "post_mip_flows", + { + "run": { + "source": {"run": "run"}, + "nodes": { + "lmms_eval": { + "type": "downstream_evaluation", + "config": {"tasks": ["ifeval"], "topology": topology}, + } + }, + } + }, + ) + + messages = _messages(state) + + assert "effective EP=32 (TP * DP)" in messages + assert "expert counts [48, 64]" in messages + + def test_persisted_vllm_measurement_topology_is_candidate_checked(tmp_path): state = _state(tmp_path) state.set_collection( From ff4631777a131fe2ffd47b88fb13f6555d0ff0a7 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Wed, 29 Jul 2026 16:36:52 -0700 Subject: [PATCH 02/16] added yaml for nano eval Signed-off-by: Grzegorz Karch --- .../nano_30b_a3b_bf16/runs/lmms_eval.yaml | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) create mode 100644 examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml diff --git a/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml b/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml new file mode 100644 index 00000000000..4bd8c6b2954 --- /dev/null +++ b/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml @@ -0,0 +1,63 @@ +# @package _global_ + +defaults: + - default + - _self_ + +# Opt-in downstream lmms-eval workflow for the realized runtime-075 candidate. +# Non-empty post_mip.flows replaces the legacy post-MIP tail in the v2 orchestrator. +zero_shot_evaluation: + enabled: false +aiperf: + enabled: false +global_distillation_sanity: + enabled: false +global_distillation: + enabled: false +post_distillation_evaluation: + enabled: false + +post_mip: + flows: + runtime-075-lmms-eval: + source: + run: runtime-075 + variants: all + objectives: all + nodes: + best_mip: + type: filter + mode: top_k + metric: mip.score + direction: minimize + top_k: 1 + materialized: + type: materialize + input: best_mip + lmms_eval: + type: downstream_evaluation + input: materialized + config: + model: vllm + checkpoint_arg: model + tasks: + - ifeval + - gsm8k + limit: 128 + batch_size: 1 + log_samples: true + timeout_seconds: 7200 + topology: + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + data_parallel_size: 1 + prefill_context_parallel_size: 1 + decode_context_parallel_size: 1 + enable_expert_parallel: false + distributed_executor_backend: mp + gpu_group_size: 8 + model_args: + dtype: bfloat16 + gpu_memory_utilization: 0.85 + max_model_len: 262144 + trust_remote_code: ${model.trust_remote_code} From 701dad988d5f1a64eefd7a68aa645959ada44e23 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Tue, 4 Aug 2026 10:48:25 -0700 Subject: [PATCH 03/16] Label downstream evaluation dashboard stages Signed-off-by: Grzegorz Karch --- .../puzzletron/orchestration/controller.py | 26 ++++++++++++++++++- .../test_orchestration_controller.py | 23 +++++++++++++++- 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/puzzletron/orchestration/controller.py b/modelopt/torch/puzzletron/orchestration/controller.py index 4c0c3477c1d..71cf6f63cf5 100644 --- a/modelopt/torch/puzzletron/orchestration/controller.py +++ b/modelopt/torch/puzzletron/orchestration/controller.py @@ -74,6 +74,29 @@ def create_executor(plan: CampaignPlan, *, local: bool = False) -> Executor: raise ValueError(f"Unsupported runner kind: {plan.runner.kind}") +def _stage_dashboard_display_name( + config: Mapping[str, Any], + stage_id: str, + *, + granularity: str | None = None, +) -> str: + if stage_id.startswith("post."): + parts = stage_id.split(".", 2) + if len(parts) == 3: + _prefix, flow_id, node_id = parts + node = ( + (config.get("post_mip") or {}) + .get("flows", {}) + .get(flow_id, {}) + .get("nodes", {}) + .get(node_id, {}) + ) + node_type = str(node.get("type") or "") if isinstance(node, Mapping) else "" + if node_type == "downstream_evaluation": + return "Downstream Evaluation" + return stage_display_name(stage_id, granularity=granularity) + + @dataclass class DryRunSubmission: stage_id: str @@ -748,7 +771,8 @@ def _stage_views(self) -> list[StageView]: views.append( StageView( stage_id=node.stage_id, - display_name=stage_display_name( + display_name=_stage_dashboard_display_name( + self.plan.experiment_config, node.stage_id, granularity=str(granularity) if granularity is not None else None, ), diff --git a/tests/unit/torch/puzzletron/test_orchestration_controller.py b/tests/unit/torch/puzzletron/test_orchestration_controller.py index a0390c30fb6..8b204706ff1 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_controller.py +++ b/tests/unit/torch/puzzletron/test_orchestration_controller.py @@ -14,7 +14,7 @@ load_execution_config, load_runner_config, ) -from puzzletron_orchestrator.controller import dry_run_plan +from puzzletron_orchestrator.controller import _stage_dashboard_display_name, dry_run_plan from puzzletron_orchestrator.executors.baremetal import GpuLeaseManager from puzzletron_orchestrator.schema import ( AttemptSpec, @@ -179,6 +179,27 @@ def test_adapter_registry_selects_sharded_adapter(tmp_path: Path): assert [item.work_id for item in work_plan.items] == ["vllm_stats:default:gang"] +def test_stage_dashboard_display_name_uses_downstream_evaluation_node_type(): + config = { + "post_mip": { + "flows": { + "params-75": { + "nodes": { + "lmms_eval": { + "type": "downstream_evaluation", + } + } + } + } + } + } + + assert ( + _stage_dashboard_display_name(config, "post.params-75.lmms_eval") + == "Downstream Evaluation" + ) + + def test_vllm_stats_rejects_conflicting_execution_mesh(tmp_path: Path): experiment, runner_path, execution_path = _write_configs(tmp_path) experiment_config = yaml.safe_load(experiment.read_text()) From 179e5216eef80b5ab705a82cdf091d641aef4c75 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Thu, 6 Aug 2026 14:35:08 -0700 Subject: [PATCH 04/16] nano/smoke pipeline works with downstream eval Signed-off-by: Grzegorz Karch --- .../distributed_eval/run_coordinator.sh | 5 +- .../distributed_eval/run_depth_coordinator.sh | 5 +- .../distributed_eval/run_depth_pool.sh | 7 +- .../distributed_eval/run_replacement_pool.sh | 7 +- .../torch/puzzletron/benchmarks/aiperf.py | 44 ++- .../puzzletron/orchestration/adapters/pool.py | 4 +- .../orchestration/executors/slurm.py | 27 ++ .../puzzletron/orchestration/task_launcher.py | 29 ++ modelopt/torch/puzzletron/post_mip/runner.py | 34 ++- modelopt/torch/puzzletron/stages/pipeline.py | 159 ++++++++++- .../subblock_stats/calc_subblock_stats.py | 127 +++++++-- .../test_aiperf_context_capacity.py | 44 +++ .../test_orchestration_executors.py | 51 ++++ .../test_orchestration_task_topology.py | 50 ++++ .../torch/puzzletron/test_post_mip_runner.py | 51 ++++ .../puzzletron/test_sparse_runtime_stats.py | 266 ++++++++++++++++++ 16 files changed, 872 insertions(+), 38 deletions(-) diff --git a/examples/puzzletron/distributed_eval/run_coordinator.sh b/examples/puzzletron/distributed_eval/run_coordinator.sh index ac22b6774fd..62ea0fc9ed1 100755 --- a/examples/puzzletron/distributed_eval/run_coordinator.sh +++ b/examples/puzzletron/distributed_eval/run_coordinator.sh @@ -3,7 +3,8 @@ set -Eeuo pipefail : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}" : "${CONFIG_PATH:?set CONFIG_PATH}" -: "${WORLD_SIZE:?set WORLD_SIZE to one torchrun worker-group world size}" +: "${WORKER_WORLD_SIZE:=${WORLD_SIZE:-}}" +: "${WORKER_WORLD_SIZE:?set WORKER_WORLD_SIZE to one torchrun worker-group world size}" : "${SOLUTIONS_PATH:?set SOLUTIONS_PATH}" : "${OUTPUT_DIR:?set OUTPUT_DIR}" @@ -42,7 +43,7 @@ if [[ ! -f "${CAMPAIGN_DIR}/manifest.json" ]]; then "${PYTHON_BIN}" -m modelopt.torch.puzzletron.distributed_eval.cli init \ --campaign-dir "${CAMPAIGN_DIR}" \ --config "${CONFIG_PATH}" \ - --world-size "${WORLD_SIZE}" \ + --world-size "${WORKER_WORLD_SIZE}" \ --evaluator-revision "${EVALUATOR_REVISION}" \ "${override_args[@]}" fi diff --git a/examples/puzzletron/distributed_eval/run_depth_coordinator.sh b/examples/puzzletron/distributed_eval/run_depth_coordinator.sh index dbc7682f931..dc37dee7b36 100755 --- a/examples/puzzletron/distributed_eval/run_depth_coordinator.sh +++ b/examples/puzzletron/distributed_eval/run_depth_coordinator.sh @@ -3,7 +3,8 @@ set -Eeuo pipefail : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}" : "${CONFIG_PATH:?set CONFIG_PATH}" -: "${WORLD_SIZE:?set WORLD_SIZE to one torchrun worker-group world size}" +: "${WORKER_WORLD_SIZE:=${WORLD_SIZE:-}}" +: "${WORKER_WORLD_SIZE:?set WORKER_WORLD_SIZE to one torchrun worker-group world size}" PYTHON_BIN="${PYTHON_BIN:-python}" @@ -18,7 +19,7 @@ if [[ ! -f "${CAMPAIGN_DIR}/manifest.json" ]]; then "${PYTHON_BIN}" -m modelopt.torch.puzzletron.distributed_eval.cli init \ --campaign-dir "${CAMPAIGN_DIR}" \ --config "${CONFIG_PATH}" \ - --world-size "${WORLD_SIZE}" \ + --world-size "${WORKER_WORLD_SIZE}" \ --stage depth \ --evaluator-revision "${EVALUATOR_REVISION:-puzzletron-depth-v1}" \ "${override_args[@]}" diff --git a/examples/puzzletron/distributed_eval/run_depth_pool.sh b/examples/puzzletron/distributed_eval/run_depth_pool.sh index 3b8b663c9b4..c38044f87e3 100644 --- a/examples/puzzletron/distributed_eval/run_depth_pool.sh +++ b/examples/puzzletron/distributed_eval/run_depth_pool.sh @@ -6,7 +6,8 @@ set -Eeuo pipefail : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}" : "${CONFIG_PATH:?set CONFIG_PATH}" -: "${WORLD_SIZE:?set WORLD_SIZE to one worker-group world size}" +: "${WORKER_WORLD_SIZE:=${WORLD_SIZE:-}}" +: "${WORKER_WORLD_SIZE:?set WORKER_WORLD_SIZE to one worker-group world size}" : "${WORKER_COUNT:?set WORKER_COUNT to the number of worker groups}" : "${PUZZLETRON_GROUP_INDEX:=${PUZZLETRON_TASK_INDEX:-${SLURM_PROCID:-}}}" : "${PUZZLETRON_GROUP_INDEX:?run this script as one orchestrator worker-group task}" @@ -55,7 +56,7 @@ if [[ "${GROUP_INDEX}" == "0" && ! -f "${MANIFEST_PATH}" ]]; then -m modelopt.torch.puzzletron.distributed_eval.cli init \ --campaign-dir "${CAMPAIGN_DIR}" \ --config "${CONFIG_PATH}" \ - --world-size "${WORLD_SIZE}" \ + --world-size "${WORKER_WORLD_SIZE}" \ --stage depth \ --evaluator-revision "${EVALUATOR_REVISION:-puzzletron-depth-v1}" \ "${override_args[@]}" @@ -74,7 +75,7 @@ done # independent worker groups may share a node. export NNODES=1 export NODE_RANK=0 -export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORLD_SIZE}}" +export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORKER_WORLD_SIZE}}" export WORKER_GROUP_INDEX="${GROUP_INDEX}" export WORKER_ID="${WORKER_PREFIX}${GROUP_INDEX}" export WORKER_HOST="${WORKER_HOST:-$(hostname -f)}" diff --git a/examples/puzzletron/distributed_eval/run_replacement_pool.sh b/examples/puzzletron/distributed_eval/run_replacement_pool.sh index 283a9ef11a3..46d5336a940 100755 --- a/examples/puzzletron/distributed_eval/run_replacement_pool.sh +++ b/examples/puzzletron/distributed_eval/run_replacement_pool.sh @@ -6,7 +6,8 @@ set -Eeuo pipefail : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}" : "${CONFIG_PATH:?set CONFIG_PATH}" -: "${WORLD_SIZE:?set WORLD_SIZE to one worker-group world size}" +: "${WORKER_WORLD_SIZE:=${WORLD_SIZE:-}}" +: "${WORKER_WORLD_SIZE:?set WORKER_WORLD_SIZE to one worker-group world size}" : "${WORKER_COUNT:?set WORKER_COUNT to the number of worker groups}" : "${PUZZLETRON_GROUP_INDEX:=${PUZZLETRON_TASK_INDEX:-${SLURM_PROCID:-}}}" : "${PUZZLETRON_GROUP_INDEX:?run this script as one orchestrator worker-group task}" @@ -54,7 +55,7 @@ if [[ "${GROUP_INDEX}" == "0" && ! -f "${MANIFEST_PATH}" ]]; then -m modelopt.torch.puzzletron.distributed_eval.cli init \ --campaign-dir "${CAMPAIGN_DIR}" \ --config "${CONFIG_PATH}" \ - --world-size "${WORLD_SIZE}" \ + --world-size "${WORKER_WORLD_SIZE}" \ --stage replace_block \ --evaluator-revision "${EVALUATOR_REVISION:-puzzletron-distributed-replace-block-v1}" \ "${override_args[@]}" @@ -71,7 +72,7 @@ done export NNODES=1 export NODE_RANK=0 -export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORLD_SIZE}}" +export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORKER_WORLD_SIZE}}" export WORKER_GROUP_INDEX="${GROUP_INDEX}" export WORKER_ID="${WORKER_PREFIX}${GROUP_INDEX}" export WORKER_HOST="${WORKER_HOST:-$(hostname -f)}" diff --git a/modelopt/torch/puzzletron/benchmarks/aiperf.py b/modelopt/torch/puzzletron/benchmarks/aiperf.py index 872822777bd..5824247f8be 100644 --- a/modelopt/torch/puzzletron/benchmarks/aiperf.py +++ b/modelopt/torch/puzzletron/benchmarks/aiperf.py @@ -143,6 +143,38 @@ def _topology_vllm_args(topology: dict[str, Any]) -> list[str]: return args +def _extra_vllm_args(topology: dict[str, Any]) -> tuple[str, ...]: + """Return caller-provided vLLM args without treating a string as characters.""" + + raw = topology.get("extra_vllm_args", ()) + if raw is None: + return () + if isinstance(raw, str): + return (raw,) + return tuple(str(arg) for arg in raw) + + +def _has_vllm_option(args: Iterable[str], *options: str) -> bool: + """Return whether an option is present as ``--flag`` or ``--flag=value``.""" + + option_set = set(options) + return any(str(arg).split("=", 1)[0] in option_set for arg in args) + + +def _server_vllm_args( + checkpoint_dir: Path, topology: dict[str, Any], concurrency_values: Iterable[int] +) -> list[str]: + """Build stable vLLM server args derived from topology and benchmark demand.""" + + args = _topology_vllm_args(topology) + args.extend(_descriptor_vllm_args(checkpoint_dir)) + extra_args = _extra_vllm_args(topology) + if not _has_vllm_option(extra_args, "--max-num-seqs", "--max_num_seqs"): + args.extend(("--max-num-seqs", str(max(concurrency_values)))) + args.extend(extra_args) + return args + + def _exact_length_extra_inputs( extra_inputs: dict[str, Any] | None, output_tokens: int ) -> dict[str, Any]: @@ -405,6 +437,8 @@ def run_aiperf_sweep( model_name = f"puzzletron-{architecture_id[:16]}" tokenizer_dir = _short_tokenizer_alias(checkpoint_dir, artifact_dir) server_log = artifact_dir / "vllm_server.log" + server_max_model_len = _server_max_model_len(input_tokens, output_tokens, topology) + server_vllm_args = _server_vllm_args(checkpoint_dir, topology, concurrency_values) server_cmd = [ "vllm", "serve", @@ -416,12 +450,10 @@ def run_aiperf_sweep( "--served-model-name", model_name, "--max-model-len", - str(_server_max_model_len(input_tokens, output_tokens, topology)), + str(server_max_model_len), "--trust-remote-code", ] - server_cmd.extend(_topology_vllm_args(topology)) - server_cmd.extend(_descriptor_vllm_args(checkpoint_dir)) - server_cmd.extend(str(arg) for arg in topology.get("extra_vllm_args", ())) + server_cmd.extend(server_vllm_args) executable = _resolve_executable(executable) env = _clean_subprocess_environment( gpu_ids, @@ -463,6 +495,10 @@ def run_aiperf_sweep( "endpoint_type": endpoint_type, "extra_inputs": _exact_length_extra_inputs(extra_inputs, output_tokens), "use_server_token_count": use_server_token_count, + "server": { + "max_model_len": server_max_model_len, + "vllm_args": tuple(server_vllm_args), + }, "revisions": revisions, }, prefix="aiperf_result", diff --git a/modelopt/torch/puzzletron/orchestration/adapters/pool.py b/modelopt/torch/puzzletron/orchestration/adapters/pool.py index 73a489fe9d3..7abee0a5263 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/pool.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/pool.py @@ -242,8 +242,8 @@ def command( "CAMPAIGN_DIR": str(campaign_dir), "CONFIG_PATH": plan.experiment_config_path, "PUZZLE_DIR": str(replacement_puzzle_dir), - "WORLD_SIZE": str(node.gpus_per_instance), "NPROC_PER_NODE": str(node.gpus_per_instance), + "WORKER_WORLD_SIZE": str(node.gpus_per_instance), "WORKER_COUNT": str(worker_count), } if node.stage_id == "depth_importance": @@ -317,7 +317,7 @@ def command( "CAMPAIGN_DIR": str(campaign_dir), "CONFIG_PATH": plan.experiment_config_path, "PUZZLE_DIR": str(replacement_puzzle_dir), - "WORLD_SIZE": str(node.gpus_per_instance), + "WORKER_WORLD_SIZE": str(node.gpus_per_instance), "WORKER_ID": str(item.metadata.get("worker_id", 0)), "WORKER_COUNT": str(node.instances), } diff --git a/modelopt/torch/puzzletron/orchestration/executors/slurm.py b/modelopt/torch/puzzletron/orchestration/executors/slurm.py index d33b865c327..047c2a0ce9f 100644 --- a/modelopt/torch/puzzletron/orchestration/executors/slurm.py +++ b/modelopt/torch/puzzletron/orchestration/executors/slurm.py @@ -90,6 +90,31 @@ def render_hook_lines(commands: Sequence[str]) -> str: return "\n".join(lines) +def _render_host_container_env(repository: str) -> str: + """Render host-side container runtime defaults for Pyxis/Enroot.""" + + cache_root = Path(repository) / ".cache" / "enroot" + lines = [ + 'if [[ -z "${ENROOT_CACHE_PATH:-}" ]]; then', + f" export ENROOT_CACHE_PATH={shlex.quote(str(cache_root / 'cache'))}", + "fi", + 'if [[ -z "${ENROOT_DATA_PATH:-}" ]]; then', + f" export ENROOT_DATA_PATH={shlex.quote(str(cache_root / 'data'))}", + "fi", + 'if [[ -z "${ENROOT_TEMP_PATH:-}" ]]; then', + ' export ENROOT_TEMP_PATH="/tmp/puzzletron-enroot-${USER:-$(id -u)}/tmp"', + "fi", + 'if [[ -z "${ENROOT_RUNTIME_PATH:-}" ]]; then', + ' export ENROOT_RUNTIME_PATH="/tmp/puzzletron-enroot-${USER:-$(id -u)}/runtime"', + "fi", + ( + 'mkdir -p "$ENROOT_CACHE_PATH" "$ENROOT_DATA_PATH" ' + '"$ENROOT_TEMP_PATH" "$ENROOT_RUNTIME_PATH"' + ), + ] + return "\n".join(lines) + + def render_sbatch_script( *, attempt: AttemptSpec, @@ -178,6 +203,8 @@ def render_sbatch_script( header_lines.append(f"#SBATCH --gpus-per-node={step_gpus_per_node}") header_lines.append(f"#SBATCH --output={log_path}") script = "\n".join(header_lines) + "\n" + if container_image: + script += _render_host_container_env(repository) + "\n" prologue_parts = [ "set -Eeuo pipefail", postrun_trap, diff --git a/modelopt/torch/puzzletron/orchestration/task_launcher.py b/modelopt/torch/puzzletron/orchestration/task_launcher.py index 4d4b21c414a..42606b6180b 100644 --- a/modelopt/torch/puzzletron/orchestration/task_launcher.py +++ b/modelopt/torch/puzzletron/orchestration/task_launcher.py @@ -29,9 +29,17 @@ TASK_IDENTITY_ENV_KEYS = frozenset( { "CUDA_VISIBLE_DEVICES", + "GROUP_RANK", + "GROUP_WORLD_SIZE", + "LOCAL_RANK", + "LOCAL_WORLD_SIZE", + "MASTER_ADDR", + "MASTER_PORT", "SLURM_LOCALID", "SLURM_NTASKS", "SLURM_PROCID", + "RANK", + "WORLD_SIZE", "PUZZLETRON_GROUP_INDEX", "PUZZLETRON_GROUP_RANK", "PUZZLETRON_GROUP_SIZE", @@ -146,6 +154,25 @@ def build_task_command( ) +def _direct_distributed_env(binding: TaskBinding) -> dict[str, str]: + """Return torch.distributed env for direct payloads that initialize env://.""" + + master_addr = "127.0.0.1" if binding.group_size == 1 else binding.master_addr + return { + "RANK": str(binding.group_rank), + "WORLD_SIZE": str(binding.group_size), + # The launcher slices CUDA_VISIBLE_DEVICES per task, so every direct + # payload has a local single-process view even when multiple tasks share + # a physical host. + "LOCAL_RANK": "0", + "LOCAL_WORLD_SIZE": "1", + "GROUP_RANK": str(binding.group_index), + "GROUP_WORLD_SIZE": str(binding.group_size), + "MASTER_ADDR": master_addr, + "MASTER_PORT": str(binding.master_port), + } + + def _required_index(env: Mapping[str, str], primary: str, fallback: str) -> int: value = env.get(primary, env.get(fallback)) if value is None: @@ -236,6 +263,8 @@ def main(argv: Sequence[str] | None = None) -> int: PUZZLETRON_MASTER_PORT=str(binding.master_port), PUZZLETRON_RENDEZVOUS_ID=binding.rendezvous_id, ) + if TaskLauncher(args.launcher) is TaskLauncher.DIRECT: + env.update(_direct_distributed_env(binding)) print( "puzzletron binding " f"host={binding.hostname} task={binding.task_index} " diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 3a1ee61afab..1bbe463d0b7 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -740,6 +740,29 @@ def _lmms_eval_result_payload(output_path: Path) -> tuple[dict[str, Any], Path]: return payload, path +def _write_lmms_eval_streams( + output_path: Path, result: subprocess.CompletedProcess[str] +) -> dict[str, str]: + stream_paths = {} + for stream_name, text in (("stdout", result.stdout), ("stderr", result.stderr)): + if not text: + continue + stream_path = output_path / f"{stream_name}.txt" + stream_path.write_text(text) + stream_paths[f"{stream_name}_path"] = str(stream_path) + return stream_paths + + +def _lmms_eval_output_tail(result: subprocess.CompletedProcess[str], *, max_lines: int = 20) -> str: + sections = [] + for stream_name, text in (("stderr", result.stderr), ("stdout", result.stdout)): + lines = (text or "").strip().splitlines() + if lines: + sections.append(f"{stream_name} tail:") + sections.extend(lines[-max_lines:]) + return "\n".join(sections) + + def _downstream_evaluation( config: dict[str, Any], node: CompiledPostMIPNode, @@ -782,14 +805,18 @@ def _downstream_evaluation( timeout=timeout, check=False, ) + stream_paths = _write_lmms_eval_streams(output, result) if result.returncode: - detail = (result.stderr or result.stdout).strip().splitlines() - tail = "\n".join(detail[-20:]) + tail = _lmms_eval_output_tail(result) raise RuntimeError( f"lmms-eval failed with exit code {result.returncode}" + (f": {tail}" if tail else "") ) - payload, result_path = _lmms_eval_result_payload(output) + try: + payload, result_path = _lmms_eval_result_payload(output) + except FileNotFoundError as error: + tail = _lmms_eval_output_tail(result) + raise FileNotFoundError(str(error) + (f": {tail}" if tail else "")) from error metrics = _flatten_lmms_eval_metrics(payload) if not metrics: raise RuntimeError(f"lmms-eval result has no numeric task metrics: {result_path}") @@ -808,6 +835,7 @@ def _downstream_evaluation( "result_path": str(summary_path), "raw_result_path": str(result_path), "command_path": str(command_path), + **stream_paths, } diff --git a/modelopt/torch/puzzletron/stages/pipeline.py b/modelopt/torch/puzzletron/stages/pipeline.py index 6c54623fd77..26e0c059416 100644 --- a/modelopt/torch/puzzletron/stages/pipeline.py +++ b/modelopt/torch/puzzletron/stages/pipeline.py @@ -33,6 +33,7 @@ from ..rpc_eval import EvaluationCache, EvaluationRequest, EvaluationResult from ..scoring_parent import ensure_scoring_parent from ..subblock_stats.measurements import apply_vllm_measurement, normalize_vllm_measurements +from ..tools.hydra_utils import clone_hydra_config from .common import complete_stage, experiment_dir, stage_manifest_path if TYPE_CHECKING: @@ -217,10 +218,10 @@ def _calculate_static_workload_stats(config: dict[str, Any], hydra_cfg: Any) -> "osl": int(hydra_cfg.calc_subblock_stats.generation_seq_len), "batch_size": int(hydra_cfg.calc_subblock_stats.batch_sizes[0]), } - } + } for raw_workload in workloads.values(): workload = dict(raw_workload or {}) - selected = OmegaConf.create(OmegaConf.to_container(hydra_cfg, resolve=True)) + selected = clone_hydra_config(hydra_cfg) OmegaConf.set_struct(selected, False) stats_cfg = selected.calc_subblock_stats concurrency = int(workload.get("concurrency", workload.get("batch_size", 1))) @@ -241,6 +242,158 @@ def _calculate_static_workload_stats(config: dict[str, Any], hydra_cfg: Any) -> launch_calc_subblock_stats(selected) +def _scenario_hidden_width(puzzle_dir: Path) -> int | None: + manifest_path = puzzle_dir / "scenario_manifest.json" + if not manifest_path.is_file(): + return None + manifest = json.loads(manifest_path.read_text()) + hidden_width = manifest.get("hidden_width") + return int(hidden_width) if hidden_width is not None else None + + +def _has_runtime_measurement( + path: Path, + *, + hidden_width: int, + measurement: Any, + allow_missing_workload_id: bool = False, +) -> bool: + try: + payload = json.loads(path.read_text()) + except (OSError, ValueError): + return False + if not isinstance(payload, list): + return False + expected_backend = (measurement.runtime_stats or {}).get("backend") + for entry in payload: + if not isinstance(entry, dict): + continue + args = entry.get("args") or {} + if not isinstance(args, dict) or args.get("runtime_stats") is not True: + continue + if int(args.get("n_embd", -1)) != int(hidden_width): + continue + if args.get("weights_dtype") != "torch.bfloat16": + continue + if int(args.get("batch_size", -1)) != int(measurement.batch_size): + continue + if int(args.get("prefill_seq_len", -1)) != int(measurement.prefill_seq_len): + continue + if int(args.get("generation_seq_len", -1)) != int(measurement.generation_seq_len): + continue + if int(args.get("max_num_seqs", -1)) != int(measurement.max_num_seqs): + continue + if args.get("runtime_granularity", "subblock") != measurement.granularity: + continue + if expected_backend is not None and args.get("runtime_backend") != expected_backend: + continue + workload_id = args.get("workload_id") + if workload_id is None and not allow_missing_workload_id: + continue + if workload_id is not None and workload_id != measurement.measurement_id: + continue + return True + return False + + +def _runtime_measurement_candidate_paths( + *, + config: dict[str, Any], + puzzle_dir: Path, + stats_path: Path, + measurement: Any, +) -> list[tuple[Path, bool]]: + stats_name = str( + (config.get("vllm_stats") or {}).get("subblock_stats_filename", stats_path.name) + ) + candidates: list[tuple[Path, bool]] = [(stats_path, False)] + root_dir = puzzle_dir + if ( + puzzle_dir.name == "depth-00" + and puzzle_dir.parent.name.startswith("width-") + and puzzle_dir.parent.parent.name == "scenarios" + ): + root_dir = puzzle_dir.parent.parent.parent + candidates.append((root_dir / stats_name, False)) + relative_stats_path = getattr(measurement, "relative_stats_path", None) + if relative_stats_path is not None: + candidates.append((root_dir / Path(relative_stats_path), True)) + return candidates + + +def _runtime_reuse_source_path( + config: dict[str, Any], + *, + puzzle_dir: Path, + stats_path: Path, + hidden_width: int, + measurement: Any, +) -> Path: + candidates = _runtime_measurement_candidate_paths( + config=config, + puzzle_dir=puzzle_dir, + stats_path=stats_path, + measurement=measurement, + ) + for candidate, allow_missing_workload_id in dict.fromkeys(candidates): + if _has_runtime_measurement( + candidate, + hidden_width=hidden_width, + measurement=measurement, + allow_missing_workload_id=allow_missing_workload_id, + ): + return candidate + raise RuntimeError( + "build-library cannot refresh width-scenario runtime stats because no " + f"reusable vLLM measurement for width {hidden_width}, workload " + f"{measurement.measurement_id!r} was found in {[str(path) for path, _ in candidates]}" + ) + + +def _refresh_scenario_runtime_workload_stats( + config: dict[str, Any], + hydra_cfg: Any, + stats_path: Path, +) -> None: + """Refresh width-scenario runtime rows with the local parameter inventory identity.""" + from ..subblock_stats.calc_subblock_stats import launch_calc_subblock_stats + + puzzle_dir = _puzzle_dir(config, hydra_cfg) + hidden_width = _scenario_hidden_width(puzzle_dir) + if hidden_width is None: + return + measurements = normalize_vllm_measurements(config) + for measurement in measurements.values(): + if measurement.model_hidden_sizes and hidden_width not in measurement.model_hidden_sizes: + continue + source_path = _runtime_reuse_source_path( + config, + puzzle_dir=puzzle_dir, + stats_path=stats_path, + hidden_width=hidden_width, + measurement=measurement, + ) + selected = clone_hydra_config(hydra_cfg) + stats_cfg = selected.calc_subblock_stats + stats_cfg.model_hidden_sizes = [hidden_width] + stats_cfg.batch_sizes = [measurement.batch_size] + stats_cfg.prefill_seq_len = measurement.prefill_seq_len + stats_cfg.generation_seq_len = measurement.generation_seq_len + if stats_cfg.get("runtime_stats") is None: + stats_cfg.runtime_stats = {} + runtime_cfg = stats_cfg.runtime_stats + for key, value in dict(measurement.runtime_stats).items(): + runtime_cfg[key] = _plain(value) + runtime_cfg.enabled = True + runtime_cfg.reuse_stats_path = str(source_path) + runtime_cfg.workload_id = measurement.measurement_id + runtime_cfg.reuse_workload_id_if_missing = measurement.measurement_id + runtime_cfg.max_num_seqs = measurement.max_num_seqs + runtime_cfg.granularity = measurement.granularity + stats_cfg.merge_with_existing_stats = True + launch_calc_subblock_stats(selected) + + def _write_runtime_subblock_library(path: Path, block_configs: tuple[Any, ...]) -> None: """Write the legacy subblock-library input without assembling a replacement library.""" rows = [] @@ -1000,6 +1153,8 @@ def build_library_stage(config: dict[str, Any], manifest: StageManifest): ) launch_build_replacement_library(hydra_cfg) + if _vllm_stats_is_explicit(config): + _refresh_scenario_runtime_workload_stats(config, hydra_cfg, stats_path) _calculate_static_workload_stats(config, hydra_cfg) # Keep this stage-local so handler registration and config preflight do not # load build-library implementation dependencies before the stage executes. diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py index 8f44bc69e85..38580f156b7 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py @@ -259,6 +259,7 @@ def _runtime_measurement_fields( "latency_difference_negative", ) _REUSABLE_RUNTIME_ARG_FIELDS = ( + "workload_id", "runtime_granularity", "runtime_backend", "num_iters", @@ -270,7 +271,67 @@ def _runtime_measurement_fields( ) -def _reuse_runtime_stats(target: dict, source: dict, *, source_path: str) -> dict: +def _runtime_reuse_key_from_args( + args: Mapping, + *, + fallback_workload_id: str | None = None, +) -> tuple | None: + """Return the exact runtime-reuse identity represented by persisted args.""" + + if not args.get("runtime_stats") or args.get("weights_dtype") != str(torch.bfloat16): + return None + workload_id = args.get("workload_id") + if workload_id is None: + workload_id = fallback_workload_id + return ( + int(args["n_embd"]), + int(args["batch_size"]), + int(args.get("prefill_seq_len")), + int(args.get("generation_seq_len")), + args.get("max_num_seqs"), + args.get("runtime_granularity", "subblock"), + args.get("runtime_backend"), + args.get("num_iters"), + args.get("num_warmup_iters"), + args.get("repeat_block_n_times"), + _freeze_stats_args(args.get("vllm_args")), + workload_id, + ) + + +def _runtime_reuse_key( + *, + width: int, + batch_size: int, + prefill_seq_len: int, + generation_seq_len: int, + runtime_stats_config: Mapping, +) -> tuple: + """Return the exact runtime-reuse identity requested by this calculation.""" + + return ( + int(width), + int(batch_size), + int(prefill_seq_len), + int(generation_seq_len), + runtime_stats_config.get("max_num_seqs"), + runtime_stats_config.get("granularity", "subblock"), + runtime_stats_config.get("backend"), + runtime_stats_config.get("num_iters", 30), + runtime_stats_config.get("num_warmup_iters", 10), + max(2, int(runtime_stats_config.get("repeat_block_n_times", 10))), + _freeze_stats_args([str(arg) for arg in runtime_stats_config.get("vllm_args", [])]), + runtime_stats_config.get("workload_id"), + ) + + +def _reuse_runtime_stats( + target: dict, + source: dict, + *, + source_path: str, + fallback_workload_id: str | None = None, +) -> dict: """Overlay immutable measured latency onto refreshed static statistics.""" # Synthetic vLLM timings are layer-independent: collection benchmarks the @@ -302,7 +363,10 @@ def _reuse_runtime_stats(target: dict, source: dict, *, source_path: str) -> dic target_args["runtime_stats"] = True target_args["runtime_reuse_source"] = str(source_path) for field in _REUSABLE_RUNTIME_ARG_FIELDS: - target_args[field] = source_args.get(field) + value = source_args.get(field) + if field == "workload_id" and value is None: + value = target_args.get("workload_id") or fallback_workload_id + target_args[field] = value source_non_block = source.get("non_block", {}) target_non_block = target.setdefault("non_block", {}) @@ -812,6 +876,7 @@ def calculate_subblock_stats( max_num_seqs=( runtime_stats_config.get("max_num_seqs") if runtime_stats_enabled else None ), + workload_id=runtime_stats_config.get("workload_id") if runtime_stats_enabled else None, repeat_block_n_times=( max(2, int(runtime_stats_config.get("repeat_block_n_times", 10))) if runtime_stats_enabled @@ -1160,6 +1225,7 @@ def _arg_signature(args: dict) -> tuple: runtime_stats, args.get("runtime_granularity") if runtime_stats else None, args.get("max_num_seqs") if runtime_stats else None, + args.get("workload_id") if runtime_stats else None, args.get("runtime_selection_identity"), args.get("parameter_inventory_identity"), ) @@ -1174,6 +1240,7 @@ def _subblock_stats_already_complete( runtime_stats_enabled: bool, runtime_granularity: str = "subblock", runtime_max_num_seqs: int | None = None, + runtime_workload_id: str | None = None, runtime_selection_identity: str | None = None, parameter_inventory_identities: Mapping[int, str] | None = None, prefill_seq_len: int = 2048, @@ -1221,6 +1288,7 @@ def _entry_subblock_keys(entry: dict) -> set[tuple[SubblockConfig, int]]: runtime_expected, runtime_granularity if runtime_expected else None, runtime_max_num_seqs if runtime_expected else None, + runtime_workload_id if runtime_expected else None, ( runtime_selection_identity if runtime_expected @@ -1290,6 +1358,8 @@ def calculate_subblock_stats_for_puzzle_dir( batch_sizes = [ int(batch_size) for batch_size in batch_sizes.strip("[]").replace(" ", "").split(",") ] + else: + batch_sizes = list(batch_sizes) master_puzzle_dir = Path(master_puzzle_dir) teacher_dir = ( @@ -1333,7 +1403,7 @@ def calculate_subblock_stats_for_puzzle_dir( teacher_hidden_size = int(lm_config.hidden_size) model_hidden_sizes = _unique_hidden_sizes(model_hidden_sizes, teacher_hidden_size) runtime_reuse_path = runtime_stats_config.get("reuse_stats_path") - runtime_reuse_by_width: dict[int, dict] = {} + runtime_reuse_by_key: dict[tuple, dict] = {} if runtime_stats_enabled and runtime_reuse_path: runtime_reuse_path = Path(str(runtime_reuse_path)) if not runtime_reuse_path.is_file(): @@ -1345,17 +1415,30 @@ def calculate_subblock_stats_for_puzzle_dir( f"Reusable runtime stats file does not exist: {runtime_reuse_path}" ) reusable_entries = json.loads(runtime_reuse_path.read_text()) - runtime_reuse_by_width = { - int(entry["args"]["n_embd"]): entry - for entry in reusable_entries - if entry.get("args", {}).get("runtime_stats") - and entry.get("args", {}).get("weights_dtype") == str(torch.bfloat16) + fallback_workload_id = runtime_stats_config.get("reuse_workload_id_if_missing") + for entry in reusable_entries: + key = _runtime_reuse_key_from_args( + entry.get("args", {}), + fallback_workload_id=fallback_workload_id, + ) + if key is not None: + runtime_reuse_by_key[key] = entry + requested_runtime_keys = { + _runtime_reuse_key( + width=int(width), + batch_size=int(batch_size), + prefill_seq_len=int(prefill_seq_len), + generation_seq_len=int(generation_seq_len), + runtime_stats_config=runtime_stats_config, + ) + for width in model_hidden_sizes + for batch_size in batch_sizes } - missing_runtime_widths = set(model_hidden_sizes) - set(runtime_reuse_by_width) - if missing_runtime_widths: + missing_runtime_keys = requested_runtime_keys - set(runtime_reuse_by_key) + if missing_runtime_keys: raise ValueError( - f"Reusable runtime stats {runtime_reuse_path} are missing widths " - f"{sorted(missing_runtime_widths)}" + f"Reusable runtime stats {runtime_reuse_path} are missing requested " + f"runtime identities {sorted(missing_runtime_keys)}" ) runtime_selection_identity = "reuse-" + hashlib.sha256( runtime_reuse_path.read_bytes() @@ -1410,6 +1493,9 @@ def calculate_subblock_stats_for_puzzle_dir( runtime_max_num_seqs=calc_subblock_stats_config.get("runtime_stats", {}).get( "max_num_seqs" ), + runtime_workload_id=calc_subblock_stats_config.get("runtime_stats", {}).get( + "workload_id" + ), runtime_selection_identity=runtime_selection_identity, parameter_inventory_identities=parameter_inventory_identities, prefill_seq_len=prefill_seq_len, @@ -1455,11 +1541,17 @@ def calculate_subblock_stats_for_puzzle_dir( curr_runtime_stats_enabled = ( runtime_stats_enabled if weights_dtype == torch.bfloat16 else False ) - reused_runtime_stats = ( - runtime_reuse_by_width.get(int(model_hidden_size)) - if curr_runtime_stats_enabled - else None - ) + reused_runtime_stats = None + if curr_runtime_stats_enabled: + reused_runtime_stats = runtime_reuse_by_key.get( + _runtime_reuse_key( + width=int(model_hidden_size), + batch_size=int(batch_size), + prefill_seq_len=int(prefill_seq_len), + generation_seq_len=int(generation_seq_len), + runtime_stats_config=runtime_stats_config, + ) + ) curr_subblock_stats = calculate_subblock_stats( calc_subblock_stats_config, @@ -1489,6 +1581,7 @@ def calculate_subblock_stats_for_puzzle_dir( curr_subblock_stats, reused_runtime_stats, source_path=str(runtime_reuse_path), + fallback_workload_id=runtime_stats_config.get("workload_id"), ) curr_subblock_stats["args"]["runtime_selection_identity"] = ( runtime_selection_identity diff --git a/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py b/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py index 616b4d24d9b..ee127c27999 100644 --- a/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py +++ b/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py @@ -27,6 +27,7 @@ _prepare_vllm_checkpoint, _profile_command, _server_max_model_len, + _server_vllm_args, _topology_vllm_args, ) @@ -155,6 +156,49 @@ def test_vllm_topology_args_enable_dp_and_expert_parallel_only_when_requested(): assert "--expert-parallel-size" not in ep_args +def test_server_vllm_args_default_max_num_seqs_follows_concurrency(monkeypatch, tmp_path): + monkeypatch.setattr( + "modelopt.torch.puzzletron.benchmarks.aiperf._descriptor_vllm_args", + lambda checkpoint_dir: ["--descriptor-arg"], + ) + + args = _server_vllm_args( + tmp_path, + { + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + "gpu_group_size": 1, + }, + (1, 4), + ) + + assert args[args.index("--max-num-seqs") + 1] == "4" + assert "--descriptor-arg" in args + + +def test_server_vllm_args_respects_explicit_max_num_seqs(monkeypatch, tmp_path): + monkeypatch.setattr( + "modelopt.torch.puzzletron.benchmarks.aiperf._descriptor_vllm_args", + lambda checkpoint_dir: [], + ) + + args = _server_vllm_args( + tmp_path, + { + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + "gpu_group_size": 1, + "extra_vllm_args": ["--max_num_seqs=8"], + }, + (1,), + ) + + assert "--max-num-seqs" not in args + assert "--max_num_seqs=8" in args + + def test_profile_command_maps_each_workload_answer_to_aiperf_cli(tmp_path): command = _profile_command( executable=Path("/opt/aiperf/bin/aiperf"), diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index c2bfab06d8f..9b628faefd6 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -194,6 +194,47 @@ def test_render_sbatch_script_omits_gpu_requests_for_cpu_stage(): assert "--gpu-bind" not in srun +def test_render_sbatch_script_sets_writable_enroot_paths_before_srun(): + runner = RunnerEnvironment( + kind="slurm", + contract=ExecutionContract( + repository="/repo", + venv="/repo/.venv", + container="/images/pytorch.sqsh", + container_mounts="/repo:/repo", + ), + slurm=SlurmRunnerConfig(account="acct", partition_cpu="cpu"), + ) + attempt = AttemptSpec( + attempt_id="a1", + work_id="build_library:0", + stage_id="build_library", + command=CommandSpec(argv=("python", "worker.py"), log_path="/tmp/out.log"), + allocation_nodes=1, + allocation_gpus=0, + metadata={"gpus_per_node": 0, "partition": "cpu"}, + task_topology=TaskTopology(task_count=1, gpus_per_task=0), + ) + + script = render_sbatch_script( + attempt=attempt, + runner=runner, + partition="cpu", + account="acct", + time_limit="4:00:00", + qos=None, + job_name="pt-build", + ) + + enroot_cache = "export ENROOT_CACHE_PATH=/repo/.cache/enroot/cache" + enroot_runtime = ( + 'export ENROOT_RUNTIME_PATH="/tmp/puzzletron-enroot-${USER:-$(id -u)}/runtime"' + ) + assert enroot_cache in script + assert enroot_runtime in script + assert script.index(enroot_cache) < script.index("srun ") + + def test_vllm_aggregation_uses_slurm_execution_contract(tmp_path: Path, monkeypatch): """Controller-side merges must run in the same container/venv as workers.""" @@ -488,6 +529,8 @@ def test_depth_pool_uses_one_four_node_gang_allocation(tmp_path: Path): assert attempt.allocation_nodes == 4 assert attempt.allocation_gpus == 32 assert attempt.metadata["kill_on_bad_exit"] is True + assert attempt.command.env["WORKER_WORLD_SIZE"] == "8" + assert "WORLD_SIZE" not in attempt.command.env assert attempt.command.argv[-1].endswith("run_depth_pool.sh") script = render_sbatch_script( @@ -503,6 +546,8 @@ def test_depth_pool_uses_one_four_node_gang_allocation(tmp_path: Path): assert "#SBATCH --ntasks=4" in script assert "#SBATCH --gpus-per-node=8" in script assert "--kill-on-bad-exit=1" in script + assert "export WORKER_WORLD_SIZE=8" in script + assert "export WORLD_SIZE=8" not in script assert "run_depth_pool.sh" in script @@ -551,6 +596,8 @@ def test_depth_pool_packs_four_two_gpu_workers_per_node(tmp_path: Path): assert attempt.allocation_gpus == 16 assert attempt.task_topology.task_count == 8 assert attempt.task_topology.gpus_per_task == 2 + assert attempt.command.env["WORKER_WORLD_SIZE"] == "2" + assert "WORLD_SIZE" not in attempt.command.env script = render_sbatch_script( attempt=attempt, @@ -665,6 +712,8 @@ def test_replacement_pool_uses_one_four_node_gang_allocation(tmp_path: Path): assert attempt.allocation_gpus == 32 assert attempt.metadata["kill_on_bad_exit"] is True assert attempt.metadata["partition"] == "batch" + assert attempt.command.env["WORKER_WORLD_SIZE"] == "8" + assert "WORLD_SIZE" not in attempt.command.env assert attempt.command.argv[-1].endswith("run_replacement_pool.sh") @@ -724,6 +773,8 @@ def test_replacement_pool_splits_workers_across_embedding_widths(tmp_path: Path) assert [attempt.allocation_gpus for attempt in attempts] == [16, 16] assert [attempt.task_topology.task_count for attempt in attempts] == [4, 4] assert [attempt.task_topology.gpus_per_task for attempt in attempts] == [4, 4] + assert [attempt.command.env["WORKER_WORLD_SIZE"] for attempt in attempts] == ["4", "4"] + assert all("WORLD_SIZE" not in attempt.command.env for attempt in attempts) assert [attempt.command.env["WORKER_COUNT"] for attempt in attempts] == ["4", "4"] assert [ attempt.command.env["FINALIZE_EXPECTED_COMPLETIONS"] for attempt in attempts diff --git a/tests/unit/torch/puzzletron/test_orchestration_task_topology.py b/tests/unit/torch/puzzletron/test_orchestration_task_topology.py index 22791749112..fe15e657068 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_task_topology.py +++ b/tests/unit/torch/puzzletron/test_orchestration_task_topology.py @@ -156,6 +156,56 @@ def fake_execvpe(executable, command, env) -> None: assert result == 0 assert captured["env"]["CUDA_VISIBLE_DEVICES"] == expected assert captured["env"]["PUZZLETRON_TASK_LAUNCHER"] == "direct" + assert captured["env"]["RANK"] == "0" + assert captured["env"]["WORLD_SIZE"] == "1" + assert captured["env"]["LOCAL_RANK"] == "0" + assert captured["env"]["LOCAL_WORLD_SIZE"] == "1" + assert captured["env"]["MASTER_ADDR"] == "127.0.0.1" + assert captured["env"]["MASTER_PORT"].isdigit() + + +def test_cpu_direct_task_launcher_sets_single_rank_env(monkeypatch) -> None: + captured: dict[str, object] = {} + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") + monkeypatch.setenv("PUZZLETRON_TASK_INDEX", "0") + monkeypatch.setenv("PUZZLETRON_LOCAL_TASK_INDEX", "0") + monkeypatch.setenv("PUZZLETRON_TASK_HOSTS", "cpu-a") + + def fake_execvpe(executable, command, env) -> None: + captured.update(executable=executable, command=command, env=env) + + monkeypatch.setattr(task_launcher.os, "execvpe", fake_execvpe) + + result = task_launcher.main( + [ + "--attempt-id", + "attempt-a", + "--nodes", + "1", + "--gpus-per-node", + "0", + "--task-count", + "1", + "--gpus-per-task", + "0", + "--tasks-per-group", + "1", + "--launcher", + "direct", + "--", + "python", + "worker.py", + ] + ) + + assert result == 0 + assert captured["env"]["CUDA_VISIBLE_DEVICES"] == "" + assert captured["env"]["RANK"] == "0" + assert captured["env"]["WORLD_SIZE"] == "1" + assert captured["env"]["LOCAL_RANK"] == "0" + assert captured["env"]["LOCAL_WORLD_SIZE"] == "1" + assert captured["env"]["MASTER_ADDR"] == "127.0.0.1" + assert captured["env"]["MASTER_PORT"].isdigit() def _task_binding(*, group_size: int) -> task_launcher.TaskBinding: diff --git a/tests/unit/torch/puzzletron/test_post_mip_runner.py b/tests/unit/torch/puzzletron/test_post_mip_runner.py index 3ea3f722468..432e9a89f1c 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -280,3 +280,54 @@ def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): } assert Path(result["result_path"]).is_file() assert Path(result["raw_result_path"]).name == "results.json" + + +def test_downstream_evaluation_reports_lmms_eval_output_when_results_are_missing( + monkeypatch, tmp_path +): + def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): + del cwd, env, capture_output, text, timeout, check + return subprocess.CompletedProcess( + argv, + 0, + stdout="Saving results aggregated\nCould not save results aggregated\n", + stderr="", + ) + + monkeypatch.setattr(runner.subprocess, "run", fake_run) + node = SimpleNamespace( + node_id="lmms_eval", + flow_id="runtime", + stage_id="post.runtime.lmms_eval", + config={ + "config": { + "command_prefix": ["python", "-m", "lmms_eval"], + "tasks": ["ifeval"], + "limit": 1, + "topology": {"gpu_group_size": 1}, + } + }, + ) + source = SimpleNamespace( + architecture_id="architecture", + artifact_kind=ArtifactKind.CHECKPOINT, + artifact={"checkpoint": str(tmp_path / "checkpoint")}, + ) + + try: + runner._downstream_evaluation( + {"puzzle_dir": str(tmp_path)}, node, source, "execution" + ) + except FileNotFoundError as error: + message = str(error) + else: + raise AssertionError("expected missing lmms-eval results to fail") + + assert "lmms-eval wrote no JSON results" in message + assert "stdout tail:" in message + assert "Could not save results aggregated" in message + stream_root = ( + tmp_path + / "artifacts/post_mip/nodes/lmms_eval/executions/execution/raw/architecture/lmms_eval" + ) + assert list(stream_root.glob("attempt_*/stdout.txt")) diff --git a/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py b/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py index 94c160c95a5..4e9a6e39c6d 100644 --- a/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py +++ b/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py @@ -22,6 +22,7 @@ from types import SimpleNamespace import pytest +import torch from immutabledict import immutabledict from omegaconf import OmegaConf @@ -66,6 +67,7 @@ _reuse_runtime_stats, _runtime_measurement_fields, _select_runtime_subblock_configs, + _subblock_stats_already_complete, _unique_hidden_sizes, _validate_sparse_runtime_settings, ) @@ -187,6 +189,7 @@ def test_runtime_stats_can_be_reused_while_static_metrics_are_refreshed(): "runtime_stats": True, "runtime_granularity": "subblock", "runtime_backend": "vllm", + "workload_id": "serving-default", }, "subblocks": [ { @@ -211,6 +214,7 @@ def test_runtime_stats_can_be_reused_while_static_metrics_are_refreshed(): assert result["args"]["runtime_stats"] is True assert result["args"]["runtime_reuse_source"] == "existing.json" + assert result["args"]["workload_id"] == "serving-default" assert result["subblocks"][0]["runtime_ms"] == 7.0 assert result["subblocks"][1]["runtime_ms"] == 7.0 assert result["subblocks"][0]["num_params"] == 123 @@ -221,6 +225,42 @@ def test_runtime_stats_can_be_reused_while_static_metrics_are_refreshed(): assert result["non_block"]["runtime_ms"] == 1.0 +def test_runtime_stats_reuse_preserves_target_workload_when_source_is_legacy(): + config = FFNConfig(intermediate_size=16) + target = { + "args": {"runtime_stats": False, "workload_id": None}, + "subblocks": [{"subblock_config": config.to_dict(), "parent_layer_index": 0}], + "non_block": {}, + } + source = { + "args": { + "runtime_stats": True, + "runtime_granularity": "subblock", + "runtime_backend": "vllm", + "workload_id": None, + }, + "subblocks": [ + { + "subblock_config": config.to_dict(), + "parent_layer_index": 0, + "runtime_ms": 7.0, + } + ], + "non_block": {}, + } + + result = _reuse_runtime_stats( + target, + source, + source_path="measurement.json", + fallback_workload_id="serving-default", + ) + + assert result["args"]["runtime_stats"] is True + assert result["args"]["workload_id"] == "serving-default" + assert result["subblocks"][0]["runtime_ms"] == 7.0 + + def test_runtime_block_candidates_load_converted_teacher_without_replacement_library(tmp_path): attention = AttentionConfig(num_query_heads=8, num_kv_heads=2) teacher_block = { @@ -1145,6 +1185,232 @@ def raise_candidate_library_import_error(*_args, **_kwargs): assert not (tmp_path / "manifests" / "build_library.json").exists() +def test_static_workload_stats_accept_instantiated_activation_pass_objects(monkeypatch): + class DummyMixin: + pass + + hydra_cfg = OmegaConf.create( + { + "calc_subblock_stats": { + "batch_sizes": [1], + "prefill_seq_len": 64, + "generation_seq_len": 16, + "runtime_stats": {"enabled": True}, + "merge_with_existing_stats": False, + }, + "pruning": { + "activation_passes": [ + { + "name": "attention_grouped", + "pruning_mixin": DummyMixin, + } + ] + }, + }, + flags={"allow_objects": True}, + ) + captured = [] + monkeypatch.setattr( + "modelopt.torch.puzzletron.subblock_stats.calc_subblock_stats.launch_calc_subblock_stats", + lambda cfg: captured.append(cfg), + ) + + pipeline_stages._calculate_static_workload_stats( + { + "mip": { + "workloads": { + "short": {"batch_size": 2, "isl": 32, "osl": 8}, + } + } + }, + hydra_cfg, + ) + + assert len(captured) == 1 + stats_cfg = captured[0].calc_subblock_stats + assert stats_cfg.batch_sizes == [2] + assert stats_cfg.prefill_seq_len == 32 + assert stats_cfg.generation_seq_len == 8 + assert stats_cfg.runtime_stats.enabled is False + assert stats_cfg.merge_with_existing_stats is True + assert ( + captured[0].pruning.activation_passes[0].pruning_mixin + is DummyMixin + ) + assert hydra_cfg.calc_subblock_stats.batch_sizes == [1] + assert hydra_cfg.calc_subblock_stats.runtime_stats.enabled is True + + +def test_width_scenario_runtime_stats_reuse_root_measurement(tmp_path, monkeypatch): + scenario = tmp_path / "scenarios" / "width-2688" / "depth-00" + scenario.mkdir(parents=True) + (scenario / "scenario_manifest.json").write_text( + json.dumps({"status": "complete", "hidden_width": 2688}) + ) + (tmp_path / "subblock_stats.json").write_text( + json.dumps( + [ + { + "args": { + "runtime_stats": True, + "runtime_granularity": "subblock", + "runtime_backend": "vllm", + "weights_dtype": "torch.bfloat16", + "batch_size": 1, + "prefill_seq_len": 4096, + "generation_seq_len": 1024, + "max_num_seqs": 1, + "n_embd": 2688, + "workload_id": "serving-default", + }, + "subblocks": [], + } + ] + ) + ) + hydra_cfg = OmegaConf.create( + { + "puzzle_dir": str(scenario), + "calc_subblock_stats": { + "model_hidden_sizes": [2688, 2560], + "batch_sizes": [1], + "prefill_seq_len": 4096, + "generation_seq_len": 1024, + "subblock_stats_filename": "subblock_stats.json", + "merge_with_existing_stats": True, + "runtime_stats": { + "enabled": True, + "backend": "vllm", + "granularity": "subblock", + "max_num_seqs": 1, + "num_iters": 30, + "num_warmup_iters": 10, + "repeat_block_n_times": 4, + "topology": {"gpu_group_size": 1}, + }, + }, + } + ) + config = { + "puzzle_dir": str(scenario), + "vllm_stats": { + "enabled": True, + "subblock_stats_filename": "subblock_stats.json", + "measurements": { + "serving-default": { + "prefill_seq_len": 4096, + "generation_seq_len": 1024, + "batch_size": 1, + "max_num_seqs": 1, + "granularity": "subblock", + "runtime_stats": { + "backend": "vllm", + "granularity": "subblock", + "max_num_seqs": 1, + "num_iters": 30, + "num_warmup_iters": 10, + "repeat_block_n_times": 4, + "topology": {"gpu_group_size": 1}, + }, + } + }, + }, + } + calls = [] + + def launch(cfg): + calls.append(cfg) + assert list(cfg.calc_subblock_stats.model_hidden_sizes) == [2688] + assert cfg.calc_subblock_stats.runtime_stats.enabled is True + assert cfg.calc_subblock_stats.runtime_stats.workload_id == "serving-default" + assert ( + cfg.calc_subblock_stats.runtime_stats.reuse_workload_id_if_missing + == "serving-default" + ) + assert cfg.calc_subblock_stats.runtime_stats.reuse_stats_path == str( + tmp_path / "subblock_stats.json" + ) + + monkeypatch.setattr( + "modelopt.torch.puzzletron.subblock_stats.calc_subblock_stats.launch_calc_subblock_stats", + launch, + ) + + pipeline_stages._refresh_scenario_runtime_workload_stats( + config, + hydra_cfg, + scenario / "subblock_stats.json", + ) + + assert len(calls) == 1 + + +def test_runtime_stats_resume_signature_includes_workload_id(): + config = FFNConfig(intermediate_size=16) + runtime_fields = { + "runtime_ms": 1.0, + "prefill_runtime_ms": 1.0, + "decode_runtime_ms": 1.0, + "decode_runtime_ms_per_token": 1.0, + "weight_memory_mib": 1.0, + "kv_cache_bytes_per_token": 1.0, + "state_cache_bytes_per_sequence": 1.0, + "prefill_flops": 1.0, + "decode_flops": 1.0, + } + existing = { + "args": { + "batch_size": 1, + "prefill_seq_len": 4096, + "generation_seq_len": 1024, + "weights_dtype": "torch.bfloat16", + "activations_dtype": "torch.bfloat16", + "kv_cache_dtype": "torch.bfloat16", + "n_embd": 2688, + "runtime_stats": True, + "runtime_granularity": "subblock", + "max_num_seqs": 1, + "workload_id": "serving-default", + "runtime_selection_identity": "reuse-root-aggregate", + "parameter_inventory_identity": "scenario-inventory", + }, + "subblocks": [ + { + "subblock_config": config.to_dict(), + "parent_layer_index": 0, + **runtime_fields, + "additive_metric_provenance": { + field: "test" for field in runtime_fields + }, + } + ], + } + kwargs = dict( + existing_stats=[existing], + subblock_configs=[_indexed(config, 0)], + batch_sizes=[1], + data_types=[(torch.bfloat16, torch.bfloat16, torch.bfloat16)], + model_hidden_sizes=[2688], + runtime_stats_enabled=True, + runtime_granularity="subblock", + runtime_max_num_seqs=1, + runtime_selection_identity="reuse-root-aggregate", + parameter_inventory_identities={2688: "scenario-inventory"}, + prefill_seq_len=4096, + generation_seq_len=1024, + ) + + assert _subblock_stats_already_complete( + **kwargs, + runtime_workload_id="serving-default", + ) + assert not _subblock_stats_already_complete( + **kwargs, + runtime_workload_id="different-workload", + ) + assert hydra_cfg.calc_subblock_stats.merge_with_existing_stats is False + + def test_sparse_runtime_selection_is_unique_and_layer_independent(): teacher_ffn = FFNConfig(intermediate_size=16) reduced_ffn = FFNConfig(intermediate_size=8) From d135c1f47dc3d5f28ccb84405dcdd7d8e063d777 Mon Sep 17 00:00:00 2001 From: "coderabbitai[bot]" <136622811+coderabbitai[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 12:16:08 +0200 Subject: [PATCH 05/16] =?UTF-8?q?=F0=9F=93=9D=20Add=20docstrings=20to=20`g?= =?UTF-8?q?karch/add=5Fdownstream=5Feval`=20(#2105)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Docstrings generation was requested by @grzegorz-k-karch. * https://github.com/NVIDIA/Model-Optimizer/pull/2104#issuecomment-5215641305 The following files were modified: * `modelopt/torch/puzzletron/benchmarks/aiperf.py` * `modelopt/torch/puzzletron/orchestration/adapters/pool.py` * `modelopt/torch/puzzletron/orchestration/adapters/post_mip.py` * `modelopt/torch/puzzletron/orchestration/compiler.py` * `modelopt/torch/puzzletron/orchestration/controller.py` * `modelopt/torch/puzzletron/orchestration/executors/slurm.py` * `modelopt/torch/puzzletron/orchestration/progress.py` * `modelopt/torch/puzzletron/orchestration/task_launcher.py` * `modelopt/torch/puzzletron/post_mip/builtin.py` * `modelopt/torch/puzzletron/post_mip/reporting.py` * `modelopt/torch/puzzletron/post_mip/runner.py` * `modelopt/torch/puzzletron/stages/pipeline.py` * `modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py` * `puzzletron_setup/bundle.py` * `puzzletron_setup/v2/validation.py` * `puzzletron_setup/v2/wizard.py` * `puzzletron_setup/wizard.py`
These files were kept as they were * `tests/unit/torch/puzzletron/test_aiperf_context_capacity.py` * `tests/unit/torch/puzzletron/test_orchestration_compiler.py` * `tests/unit/torch/puzzletron/test_orchestration_controller.py` * `tests/unit/torch/puzzletron/test_orchestration_executors.py` * `tests/unit/torch/puzzletron/test_orchestration_task_topology.py` * `tests/unit/torch/puzzletron/test_post_mip_runner.py` * `tests/unit/torch/puzzletron/test_setup_bundle.py` * `tests/unit/torch/puzzletron/test_setup_v2_post_mip.py` * `tests/unit/torch/puzzletron/test_setup_v2_state_validation.py` * `tests/unit/torch/puzzletron/test_sparse_runtime_stats.py`
These file types are not supported * `CHANGELOG.rst` * `examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml`
ℹ️ Note
CodeRabbit cannot perform edits on its own pull requests yet.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Johannes Rausch --- .../torch/puzzletron/benchmarks/aiperf.py | 60 +++++- .../puzzletron/orchestration/adapters/pool.py | 14 ++ .../orchestration/adapters/post_mip.py | 13 ++ .../puzzletron/orchestration/compiler.py | 17 +- .../puzzletron/orchestration/controller.py | 23 +++ .../orchestration/executors/slurm.py | 30 ++- .../puzzletron/orchestration/progress.py | 11 ++ .../puzzletron/orchestration/task_launcher.py | 37 +++- modelopt/torch/puzzletron/post_mip/builtin.py | 1 + .../torch/puzzletron/post_mip/reporting.py | 23 ++- modelopt/torch/puzzletron/post_mip/runner.py | 183 +++++++++++++++++- modelopt/torch/puzzletron/stages/pipeline.py | 82 +++++++- .../subblock_stats/calc_subblock_stats.py | 111 ++++++++++- puzzletron_setup/bundle.py | 28 +++ puzzletron_setup/v2/validation.py | 7 +- puzzletron_setup/v2/wizard.py | 62 +++++- puzzletron_setup/wizard.py | 68 ++++++- 17 files changed, 740 insertions(+), 30 deletions(-) diff --git a/modelopt/torch/puzzletron/benchmarks/aiperf.py b/modelopt/torch/puzzletron/benchmarks/aiperf.py index 5824247f8be..ebbeea66fde 100644 --- a/modelopt/torch/puzzletron/benchmarks/aiperf.py +++ b/modelopt/torch/puzzletron/benchmarks/aiperf.py @@ -114,7 +114,14 @@ def _canonical_topology(topology: dict[str, Any]) -> dict[str, Any]: def _topology_vllm_args(topology: dict[str, Any]) -> list[str]: - """Translate the canonical TP/PP/DP/EP/CP contract to vLLM CLI arguments.""" + """Convert canonical parallelism settings into vLLM command-line arguments. + + Parameters: + topology (dict[str, Any]): Topology configuration to normalize and convert. + + Returns: + list[str]: vLLM command-line arguments for the configured tensor, pipeline, context, data, and expert parallelism. + """ canonical = _canonical_topology(topology) args = [ @@ -164,7 +171,17 @@ def _has_vllm_option(args: Iterable[str], *options: str) -> bool: def _server_vllm_args( checkpoint_dir: Path, topology: dict[str, Any], concurrency_values: Iterable[int] ) -> list[str]: - """Build stable vLLM server args derived from topology and benchmark demand.""" + """ + Build vLLM server arguments from checkpoint configuration, topology, and concurrency demand. + + Parameters: + checkpoint_dir (Path): Directory containing the model checkpoint. + topology (dict[str, Any]): Topology and extra vLLM argument configuration. + concurrency_values (Iterable[int]): Requested concurrency levels used to set the default maximum number of sequences. + + Returns: + list[str]: Command-line arguments for the vLLM server. + """ args = _topology_vllm_args(topology) args.extend(_descriptor_vllm_args(checkpoint_dir)) @@ -178,7 +195,16 @@ def _server_vllm_args( def _exact_length_extra_inputs( extra_inputs: dict[str, Any] | None, output_tokens: int ) -> dict[str, Any]: - """Guarantee the measured OSL unless the caller chose an explicit policy.""" + """ + Configure extra input settings for exact output-length measurements. + + Parameters: + extra_inputs (dict[str, Any] | None): Optional caller-provided input settings. + + Returns: + dict[str, Any]: The input settings with ``ignore_eos`` enabled when neither + ``ignore_eos`` nor ``min_tokens`` was explicitly provided. + """ resolved = dict(extra_inputs or {}) if "ignore_eos" not in resolved and "min_tokens" not in resolved: resolved["ignore_eos"] = True @@ -409,7 +435,33 @@ def run_aiperf_sweep( benchmark_timeout: float = 600, gpu_telemetry: str | None = "pynvml", ) -> list[BenchmarkResult]: - """Run multiple concurrencies against one persistent vLLM server.""" + """ + Run multiple concurrency benchmarks against a persistent vLLM server. + + Parameters: + checkpoint_dir (str | Path): Model checkpoint directory. + artifact_dir (str | Path): Directory for benchmark outputs and logs. + concurrencies (Iterable[int]): Unique positive concurrency levels to benchmark. + input_tokens (int): Synthetic input length for each request. + output_tokens (int): Synthetic output length for each request. + gpu_ids (str): GPU visibility specification for the benchmark processes. + topology (dict[str, Any]): vLLM parallelism and environment configuration. + request_counts (dict[int, int] | None): Optional request count for each concurrency level. + solution_id (str): Identifier for the benchmark solution. + profile_id (str): Identifier for the benchmark profile. + topology_id (str | None): Optional topology identifier. + executable (str | Path): AIPerf executable to run. + endpoint_type (str): Endpoint type used by AIPerf. + extra_inputs (dict[str, Any] | None): Additional AIPerf input settings. + use_server_token_count (bool): Whether to use token counts reported by the server. + seed (int): Seed for synthetic request generation. + readiness_timeout (float): Maximum time to wait for vLLM readiness, in seconds. + benchmark_timeout (float): Maximum time allowed for each benchmark, in seconds. + gpu_telemetry (str | None): GPU telemetry backend, or None to disable telemetry. + + Returns: + list[BenchmarkResult]: Benchmark results in the original concurrency order. + """ checkpoint_dir = Path(checkpoint_dir).resolve() artifact_dir = Path(artifact_dir).resolve() diff --git a/modelopt/torch/puzzletron/orchestration/adapters/pool.py b/modelopt/torch/puzzletron/orchestration/adapters/pool.py index 7abee0a5263..098c05261b6 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/pool.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/pool.py @@ -224,6 +224,20 @@ def command( runner, overrides: list[str] | None = None, ) -> AttemptSpec: + """ + Build the execution command and resource allocation for a planned work item. + + Parameters: + plan (CampaignPlan): Campaign configuration and execution context. + node (StagePlanNode): Stage and resource configuration for the work item. + item (WorkItem): Work item metadata, role, and local GPU assignments. + attempt_id (str): Identifier for the execution attempt. + runner: Runner context containing the repository location. + overrides (list[str] | None): Optional configuration overrides to apply. + + Returns: + AttemptSpec: Command, environment, resource allocation, and execution metadata for the work item. + """ repo = Path(runner.contract.repository) role = item.metadata.get("role", "worker") log_dir = plan.puzzle_dir / "logs" diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index 90f1f745f9a..6887a744beb 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py @@ -108,6 +108,19 @@ class PostMIPAdapter(WorkAdapter): strategy = ExecutionStrategy.SHARDED def plan(self, plan: CampaignPlan, node: StagePlanNode) -> WorkPlan: + """ + Plan sharded execution for a post-MIP node and mark aggregation as required. + + Parameters: + plan (CampaignPlan): Campaign execution plan containing node configuration and candidate information. + node (StagePlanNode): Post-MIP node to plan. + + Returns: + WorkPlan: Work plan containing the node's work item and execution strategy. + + Raises: + RuntimeError: If an evaluation node has no candidate architectures to evaluate. + """ config = _node_config(plan, node.stage_id) node_type = str(config.get("type")) count = 1 if node_type in {"filter", "manual_filter"} else node.instances diff --git a/modelopt/torch/puzzletron/orchestration/compiler.py b/modelopt/torch/puzzletron/orchestration/compiler.py index 5c76c3ce087..a9b80ed2de7 100644 --- a/modelopt/torch/puzzletron/orchestration/compiler.py +++ b/modelopt/torch/puzzletron/orchestration/compiler.py @@ -431,7 +431,22 @@ def compile_campaign_plan( overrides: list[str] | None = None, stage_filter: str | None = None, ) -> CampaignPlan: - """Compile one campaign plan from experiment + runner + execution configs.""" + """ + Compile a campaign plan from experiment, runner, and execution configurations. + + Parameters: + experiment_config_path: Path to the experiment configuration file. + runner: Runner environment used to execute the campaign. + execution: Execution defaults and per-stage settings. + overrides: Optional experiment configuration overrides. + stage_filter: Optional stage identifier limiting the plan to one enabled stage. + + Returns: + A compiled campaign plan containing stage meshes, dependencies, resources, and GPU allocations. + + Raises: + ValueError: If the selected stage is disabled, a CPU stage requests multiple instances, or a mesh override conflicts with its topology. + """ experiment_path = Path(experiment_config_path) experiment_config = load_experiment_config(experiment_path, overrides=overrides or []) diff --git a/modelopt/torch/puzzletron/orchestration/controller.py b/modelopt/torch/puzzletron/orchestration/controller.py index 71cf6f63cf5..aa6a3c04877 100644 --- a/modelopt/torch/puzzletron/orchestration/controller.py +++ b/modelopt/torch/puzzletron/orchestration/controller.py @@ -63,6 +63,19 @@ def create_executor(plan: CampaignPlan, *, local: bool = False) -> Executor: + """ + Create an executor for the campaign plan's configured runner. + + Parameters: + plan (CampaignPlan): Campaign plan containing runner configuration. + local (bool): Whether to use a local executor instead of the configured runner. + + Returns: + Executor: Executor configured for local, Slurm, or bare-metal SSH execution. + + Raises: + ValueError: If the configured runner kind is unsupported. + """ if local: return LocalExecutor(plan.runner) if plan.runner.kind == "slurm": @@ -80,6 +93,16 @@ def _stage_dashboard_display_name( *, granularity: str | None = None, ) -> str: + """Resolve the dashboard display name for a campaign stage. + + Parameters: + config (Mapping[str, Any]): Campaign configuration containing post-MIP flow definitions. + stage_id (str): Stage identifier to format. + granularity (str | None): Optional naming granularity. + + Returns: + str: ``"Downstream Evaluation"`` for downstream evaluation post-MIP stages; otherwise, the formatted stage name. + """ if stage_id.startswith("post."): parts = stage_id.split(".", 2) if len(parts) == 3: diff --git a/modelopt/torch/puzzletron/orchestration/executors/slurm.py b/modelopt/torch/puzzletron/orchestration/executors/slurm.py index 047c2a0ce9f..b94e2c7319f 100644 --- a/modelopt/torch/puzzletron/orchestration/executors/slurm.py +++ b/modelopt/torch/puzzletron/orchestration/executors/slurm.py @@ -91,7 +91,15 @@ def render_hook_lines(commands: Sequence[str]) -> str: def _render_host_container_env(repository: str) -> str: - """Render host-side container runtime defaults for Pyxis/Enroot.""" + """ + Render shell commands that configure default Pyxis/Enroot paths and create their directories. + + Parameters: + repository (str): Repository path used to derive default Enroot cache and data paths. + + Returns: + str: Shell commands for configuring and preparing the container runtime environment. + """ cache_root = Path(repository) / ".cache" / "enroot" lines = [ @@ -125,7 +133,25 @@ def render_sbatch_script( qos: str | None, job_name: str, ) -> str: - """Render one sbatch script for an attempt.""" + """ + Render an executable Slurm batch script for an attempt, including resource + allocations, environment setup, hooks, logging, and optional container + configuration. + + Parameters: + attempt (AttemptSpec): Attempt specification containing the command and + requested task topology. + runner (RunnerEnvironment): Runner configuration used for repository, + environment, and container settings. + partition (str): Slurm partition for the job. + account (str): Slurm account for the job. + time_limit (str): Slurm time limit. + qos (str | None): Optional Slurm quality-of-service name. + job_name (str): Name assigned to the Slurm job. + + Returns: + str: The generated executable sbatch script. + """ contract = runner.contract topology = resolve_task_topology(attempt) diff --git a/modelopt/torch/puzzletron/orchestration/progress.py b/modelopt/torch/puzzletron/orchestration/progress.py index a9c943357e6..c89c276d66a 100644 --- a/modelopt/torch/puzzletron/orchestration/progress.py +++ b/modelopt/torch/puzzletron/orchestration/progress.py @@ -387,6 +387,17 @@ def _post_mip_progress( stage_id: str, config: Mapping[str, Any] | None, ) -> str | None: + """ + Summarize post-MIP candidate processing progress for a configured node. + + Parameters: + puzzle_dir (Path): Root directory containing post-MIP artifacts. + stage_id (str): Identifier of the post-MIP stage and node. + config (Mapping[str, Any] | None): Configuration containing post-MIP flow and node definitions. + + Returns: + str | None: Progress summary with completed, failed, and timed-out candidate counts, or None when progress data is unavailable or the node is not applicable. + """ parts = stage_id.split(".", 2) if len(parts) != 3: return None diff --git a/modelopt/torch/puzzletron/orchestration/task_launcher.py b/modelopt/torch/puzzletron/orchestration/task_launcher.py index 42606b6180b..6428ada478e 100644 --- a/modelopt/torch/puzzletron/orchestration/task_launcher.py +++ b/modelopt/torch/puzzletron/orchestration/task_launcher.py @@ -134,7 +134,18 @@ def build_task_command( binding: TaskBinding, gpus_per_task: int, ) -> tuple[str, ...]: - """Wrap an application payload in torchrun when the topology requests it.""" + """ + Build the command used to launch the application for the selected launcher and task topology. + + Parameters: + payload (Sequence[str]): Application command and its arguments. + launcher (TaskLauncher): Launcher mode that determines whether to wrap the payload. + binding (TaskBinding): Resolved task placement and rendezvous information. + gpus_per_task (int): Number of processes to launch per node when using distributed execution. + + Returns: + tuple[str, ...]: The original payload for direct execution, or a torchrun command configured for the task topology. + """ command = tuple(str(part) for part in payload) if launcher is TaskLauncher.DIRECT: @@ -174,6 +185,20 @@ def _direct_distributed_env(binding: TaskBinding) -> dict[str, str]: def _required_index(env: Mapping[str, str], primary: str, fallback: str) -> int: + """ + Read a task index from the primary environment variable or its fallback. + + Parameters: + env (Mapping[str, str]): Environment variables containing the task index. + primary (str): Preferred environment variable name. + fallback (str): Alternate environment variable name. + + Returns: + int: The task index parsed from the selected environment variable. + + Raises: + RuntimeError: If neither environment variable is set. + """ value = env.get(primary, env.get(fallback)) if value is None: raise RuntimeError(f"missing task identity: set {primary} or {fallback}") @@ -196,7 +221,15 @@ def _parser() -> argparse.ArgumentParser: def main(argv: Sequence[str] | None = None) -> int: - """Resolve this task's binding and replace the launcher with its payload.""" + """ + Resolve the task's distributed binding, prepare its execution environment, and replace the current process with the payload command. + + Parameters: + argv (Sequence[str] | None): Optional command-line arguments to parse instead of the process arguments. + + Returns: + int: Zero after replacing the current process with the payload command. + """ args = _parser().parse_args(argv) payload = tuple(args.payload[1:] if args.payload[:1] == ["--"] else args.payload) diff --git a/modelopt/torch/puzzletron/post_mip/builtin.py b/modelopt/torch/puzzletron/post_mip/builtin.py index 050352af445..52ab60509d5 100644 --- a/modelopt/torch/puzzletron/post_mip/builtin.py +++ b/modelopt/torch/puzzletron/post_mip/builtin.py @@ -125,4 +125,5 @@ class DownstreamEvaluationNode(PostMIPNode): @classmethod def render_report(cls, node, payload): + """Render the downstream evaluation report for the payload's section.""" return render_downstream_evaluation_report(str(payload["section_id"]), payload) diff --git a/modelopt/torch/puzzletron/post_mip/reporting.py b/modelopt/torch/puzzletron/post_mip/reporting.py index bc287990805..c21d61ce4a9 100644 --- a/modelopt/torch/puzzletron/post_mip/reporting.py +++ b/modelopt/torch/puzzletron/post_mip/reporting.py @@ -319,7 +319,17 @@ def render_evaluation_report(section_id: str, payload: Mapping[str, Any]) -> str def render_aiperf_report(section_id: str, payload: Mapping[str, Any]) -> str: - """Render AIPerf throughput/latency observations and timeout evidence.""" + """ + Render AIPerf candidate status, performance metrics, selection markers, and errors. + + Parameters: + section_id (str): Identifier used to scope the throughput chart element. + payload (Mapping[str, Any]): AIPerf observations and status data. + + Returns: + str: HTML fragment containing the status summary, throughput chart placeholder, + and candidate metrics table. + """ observations = list(payload.get("observations") or ()) rows = [] @@ -368,7 +378,16 @@ def render_downstream_evaluation_report(section_id: str, payload: Mapping[str, A def render_global_kd_report(section_id: str, payload: Mapping[str, Any]) -> str: - """Render several candidate KD histories on shared, lineage-colored plots.""" + """ + Render the Short KD comparison with candidate statuses, loss plots, and run summaries. + + Parameters: + section_id (str): Identifier used to generate unique plot element IDs. + payload (Mapping[str, Any]): Short KD runs and status data to display. + + Returns: + str: HTML fragment containing the comparison summary, plot placeholders, and run table. + """ runs = list(payload.get("runs") or ()) rows = [] diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 1bbe463d0b7..3d23e0d05b0 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -470,6 +470,17 @@ def _evaluate( def _aiperf( config: dict[str, Any], node: CompiledPostMIPNode, source, execution_identity: str ) -> dict[str, Any]: + """Run an AI performance sweep for a checkpoint across configured concurrency levels. + + Parameters: + config (dict[str, Any]): Workflow configuration used to determine the execution directory. + node (CompiledPostMIPNode): Compiled post-MIP node containing benchmark settings. + source: Candidate source containing the checkpoint and architecture identifier. + execution_identity (str): Identifier for the current node execution. + + Returns: + dict[str, Any]: Benchmark metrics and paths to the raw result artifacts. + """ from ..benchmarks import run_aiperf_sweep settings = dict(node.config.get("config") or {}) @@ -541,10 +552,27 @@ def _aiperf( def _as_cli_bool(value: bool) -> str: + """Convert a Boolean value to the CLI-compatible ``"True"`` or ``"False"`` string. + + Parameters: + value (bool): The Boolean value to convert. + + Returns: + str: ``"True"`` for true values and ``"False"`` for false values. + """ return "True" if value else "False" def _as_lmms_eval_arg(value: Any) -> str: + """ + Convert a value to the command-line argument format expected by lmms-eval. + + Parameters: + value (Any): The value to convert. + + Returns: + str: The formatted command-line argument value. + """ if isinstance(value, bool): return _as_cli_bool(value) if isinstance(value, (int, float)) and not isinstance(value, bool): @@ -555,6 +583,20 @@ def _as_lmms_eval_arg(value: Any) -> str: def _join_cli_values(value: Any, *, path: str) -> str: + """ + Convert a string or sequence of values into a comma-separated CLI value. + + Parameters: + value (Any): String or sequence of values to normalize. + path (str): Configuration path used in validation errors. + + Returns: + str: The normalized comma-separated value. + + Raises: + TypeError: If value is neither a string nor a sequence. + ValueError: If value is empty or contains an empty item. + """ if isinstance(value, str): text = value.strip() if not text: @@ -569,6 +611,18 @@ def _join_cli_values(value: Any, *, path: str) -> str: def _model_arg_string(values: Mapping[str, Any]) -> str: + """ + Convert model arguments to lmms-eval's comma-separated argument format. + + Parameters: + values (Mapping[str, Any]): Model argument names and values. + + Returns: + str: A comma-separated string of rendered key-value arguments. + + Raises: + ValueError: If an argument key or value is invalid, or if no arguments are provided. + """ parts = [] for key, value in values.items(): if value is None: @@ -589,6 +643,16 @@ def _model_arg_string(values: Mapping[str, Any]) -> str: def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> str: + """ + Merge checkpoint, topology, and supported model settings into lmms-eval model arguments. + + Parameters: + settings (Mapping[str, Any]): Downstream evaluation settings containing optional model arguments and configuration overrides. + checkpoint (str): Path to the checkpoint used for evaluation. + + Returns: + str: Comma-separated lmms-eval model arguments. + """ raw = settings.get("model_args") checkpoint_arg = str(settings.get("checkpoint_arg", "model")) topology = dict(settings.get("topology") or {}) @@ -625,6 +689,18 @@ def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> def _command_prefix(settings: Mapping[str, Any]) -> list[str]: + """ + Resolve the command prefix used to invoke lmms-eval. + + Parameters: + settings (Mapping[str, Any]): Downstream evaluation settings containing an optional command prefix. + + Returns: + list[str]: The configured command prefix, or the current Python interpreter followed by the lmms-eval module. + + Raises: + ValueError: If the configured command prefix is empty or contains an empty value. + """ raw = settings.get("command_prefix") if raw is None: return [sys.executable, "-m", "lmms_eval"] @@ -643,7 +719,17 @@ def _lmms_eval_command( checkpoint: str, output_path: Path, ) -> tuple[list[str], dict[str, str], float | None]: - """Build a deterministic lmms-eval CLI invocation for one realized checkpoint.""" + """ + Builds an lmms-eval command, environment, and optional timeout for a realized checkpoint. + + Parameters: + settings (Mapping[str, Any]): Downstream evaluation settings. + checkpoint (str): Path to the realized checkpoint. + output_path (Path): Directory for lmms-eval output. + + Returns: + tuple[list[str], dict[str, str], float | None]: The command arguments, environment variables, and timeout in seconds. + """ tasks = _join_cli_values(settings.get("tasks"), path="downstream_evaluation.config.tasks") argv = [ @@ -697,6 +783,15 @@ def _lmms_eval_command( def _metric_key(value: Any) -> str: + """ + Normalize a metric name component for use in metric keys. + + Parameters: + value (Any): The value to convert into a normalized metric name component. + + Returns: + str: The stripped string representation with spaces, commas, and slashes replaced by underscores. + """ return ( str(value) .strip() @@ -708,6 +803,15 @@ def _metric_key(value: Any) -> str: def _flatten_lmms_eval_metrics(payload: Mapping[str, Any]) -> dict[str, float]: + """ + Flatten finite numeric task metrics from an lmms-eval result payload. + + Parameters: + payload (Mapping[str, Any]): Result payload containing task metrics under the ``results`` key. + + Returns: + dict[str, float]: Metric names mapped to finite numeric values, or an empty dictionary when no valid results are present. + """ results = payload.get("results") if not isinstance(results, Mapping): return {} @@ -726,6 +830,18 @@ def _flatten_lmms_eval_metrics(payload: Mapping[str, Any]) -> dict[str, float]: def _lmms_eval_result_payload(output_path: Path) -> tuple[dict[str, Any], Path]: + """ + Finds the newest valid lmms-eval result payload under an output directory. + + Parameters: + output_path (Path): Directory containing lmms-eval output files. + + Returns: + tuple[dict[str, Any], Path]: The result payload and path of the newest JSON file containing a `results` mapping. + + Raises: + FileNotFoundError: If no valid result JSON file is found. + """ candidates = [] for path in sorted(output_path.rglob("*.json")): try: @@ -743,6 +859,16 @@ def _lmms_eval_result_payload(output_path: Path) -> tuple[dict[str, Any], Path]: def _write_lmms_eval_streams( output_path: Path, result: subprocess.CompletedProcess[str] ) -> dict[str, str]: + """ + Persist non-empty lmms-eval subprocess output streams and return their artifact paths. + + Parameters: + output_path (Path): Directory where stream files are written. + result (subprocess.CompletedProcess[str]): Completed subprocess result containing captured output. + + Returns: + dict[str, str]: Mapping of stream path keys to the paths of written output files. + """ stream_paths = {} for stream_name, text in (("stdout", result.stdout), ("stderr", result.stderr)): if not text: @@ -754,6 +880,16 @@ def _write_lmms_eval_streams( def _lmms_eval_output_tail(result: subprocess.CompletedProcess[str], *, max_lines: int = 20) -> str: + """ + Format the most recent subprocess output lines from stderr and stdout. + + Parameters: + result (subprocess.CompletedProcess[str]): Completed process containing captured output. + max_lines (int): Maximum number of lines to include from each stream. + + Returns: + str: Formatted stderr and stdout output tails. + """ sections = [] for stream_name, text in (("stderr", result.stderr), ("stdout", result.stdout)): lines = (text or "").strip().splitlines() @@ -769,6 +905,23 @@ def _downstream_evaluation( source, execution_identity: str, ) -> dict[str, Any]: + """ + Run downstream lmms-eval benchmarking for a materialized checkpoint. + + Parameters: + config (dict[str, Any]): Campaign configuration used to determine execution paths. + node (CompiledPostMIPNode): Post-MIP node containing lmms-eval settings. + source: Checkpoint artifact to evaluate. + execution_identity (str): Identity of the current node execution. + + Returns: + dict[str, Any]: Paths to the evaluation summary, raw result, command record, and captured streams, together with numeric metrics. + + Raises: + ValueError: If the source is not a checkpoint artifact. + RuntimeError: If lmms-eval fails or produces no numeric task metrics. + FileNotFoundError: If no valid lmms-eval result file is produced. + """ if source.artifact_kind is not ArtifactKind.CHECKPOINT: raise ValueError("downstream_evaluation requires materialized checkpoint artifacts") settings = dict(node.config.get("config") or {}) @@ -891,6 +1044,22 @@ def _run_candidate( input_revision_id: str, execution_identity: str, ) -> dict[str, Any]: + """ + Execute a candidate according to the node type and return its execution result. + + Parameters: + config (dict[str, Any]): Runtime configuration for the candidate execution. + node (CompiledPostMIPNode): Compiled node defining the execution type and model source. + ledger (CandidateLedger): Ledger containing the input candidate revision. + input_revision_id (str): Identifier of the candidate revision to execute. + execution_identity (str): Identifier for the current node execution. + + Returns: + dict[str, Any]: A successful result containing the input and source revision identifiers, architecture identifier, and executor-specific metadata. + + Raises: + ValueError: If the node type is not a supported candidate executor. + """ source = ledger.source_revision(input_revision_id, node.model_source) if node.node_type == "materialize": result = _materialize(config, node, ledger, input_revision_id, source, execution_identity) @@ -934,6 +1103,18 @@ def _distributed_shard(config: dict[str, Any], node: CompiledPostMIPNode) -> Ite def run_post_mip_node_shard( config: dict[str, Any], stage_id: str, *, shard_index: int = 0, shard_count: int = 1 ) -> Path: + """ + Execute the assigned candidate revisions for a post-MIP node shard and persist the results. + + Parameters: + config (dict[str, Any]): Post-MIP configuration. + stage_id (str): Identifier of the compiled node to execute. + shard_index (int): Zero-based index of this shard. + shard_count (int): Total number of shards distributing the candidate revisions. + + Returns: + Path: Path to the shard result artifact. + """ node = _compiled_node(config, stage_id) ledger = _ledger(config) ledger.ingest_mip(_puzzle_dir(config)) diff --git a/modelopt/torch/puzzletron/stages/pipeline.py b/modelopt/torch/puzzletron/stages/pipeline.py index 26e0c059416..279aec34744 100644 --- a/modelopt/torch/puzzletron/stages/pipeline.py +++ b/modelopt/torch/puzzletron/stages/pipeline.py @@ -207,7 +207,13 @@ def _vllm_stats_is_explicit(config: dict[str, Any]) -> bool: def _calculate_static_workload_stats(config: dict[str, Any], hydra_cfg: Any) -> None: - """Append one analytical memory profile for every configured MIP workload.""" + """ + Append an analytical memory profile for each configured MIP workload. + + Parameters: + config (dict[str, Any]): Pipeline configuration containing optional MIP workloads. + hydra_cfg (Any): Base subblock-statistics configuration to customize for each workload. + """ from ..subblock_stats.calc_subblock_stats import launch_calc_subblock_stats workloads = dict((config.get("mip") or {}).get("workloads") or {}) @@ -243,6 +249,15 @@ def _calculate_static_workload_stats(config: dict[str, Any], hydra_cfg: Any) -> def _scenario_hidden_width(puzzle_dir: Path) -> int | None: + """ + Read the hidden width from a scenario manifest. + + Parameters: + puzzle_dir (Path): Directory containing the scenario manifest. + + Returns: + int | None: The manifest's hidden width, or `None` when the manifest or value is absent. + """ manifest_path = puzzle_dir / "scenario_manifest.json" if not manifest_path.is_file(): return None @@ -258,6 +273,18 @@ def _has_runtime_measurement( measurement: Any, allow_missing_workload_id: bool = False, ) -> bool: + """ + Determine whether a statistics file contains a compatible runtime measurement. + + Parameters: + path (Path): Statistics file to inspect. + hidden_width (int): Model hidden width expected by the measurement. + measurement (Any): Runtime measurement configuration to match. + allow_missing_workload_id (bool): Whether entries without a workload identifier may match. + + Returns: + bool: `True` if a compatible runtime measurement is present, `False` otherwise. + """ try: payload = json.loads(path.read_text()) except (OSError, ValueError): @@ -303,6 +330,18 @@ def _runtime_measurement_candidate_paths( stats_path: Path, measurement: Any, ) -> list[tuple[Path, bool]]: + """ + Builds candidate paths for locating reusable runtime measurement statistics. + + Parameters: + config (dict[str, Any]): Configuration containing the statistics filename. + puzzle_dir (Path): Directory associated with the current scenario. + stats_path (Path): Primary statistics file path. + measurement (Any): Measurement configuration that may specify a relative statistics path. + + Returns: + list[tuple[Path, bool]]: Candidate statistics paths paired with a flag indicating whether each path came from a configured relative path. + """ stats_name = str( (config.get("vllm_stats") or {}).get("subblock_stats_filename", stats_path.name) ) @@ -329,6 +368,22 @@ def _runtime_reuse_source_path( hidden_width: int, measurement: Any, ) -> Path: + """ + Finds a reusable vLLM measurement file matching the requested hidden width and workload. + + Parameters: + config (dict[str, Any]): Runtime configuration used to resolve candidate measurement paths. + puzzle_dir (Path): Experiment directory containing scenario-specific measurement files. + stats_path (Path): Configured statistics file path. + hidden_width (int): Hidden width required for the reusable measurement. + measurement (Any): Workload measurement whose identity must match. + + Returns: + Path: The first candidate measurement file containing a compatible runtime measurement. + + Raises: + RuntimeError: If no candidate contains a matching reusable measurement. + """ candidates = _runtime_measurement_candidate_paths( config=config, puzzle_dir=puzzle_dir, @@ -355,7 +410,7 @@ def _refresh_scenario_runtime_workload_stats( hydra_cfg: Any, stats_path: Path, ) -> None: - """Refresh width-scenario runtime rows with the local parameter inventory identity.""" + """Refresh scenario-specific runtime statistics using measurements for the local hidden width.""" from ..subblock_stats.calc_subblock_stats import launch_calc_subblock_stats puzzle_dir = _puzzle_dir(config, hydra_cfg) @@ -395,7 +450,13 @@ def _refresh_scenario_runtime_workload_stats( def _write_runtime_subblock_library(path: Path, block_configs: tuple[Any, ...]) -> None: - """Write the legacy subblock-library input without assembling a replacement library.""" + """ + Write runtime subblock configurations to a JSON library file. + + Parameters: + path (Path): Destination path for the library file. + block_configs (tuple[Any, ...]): Block configurations to serialize. + """ rows = [] for block_config in block_configs: row = { @@ -1124,6 +1185,21 @@ def bypass_overfit_stage(config: dict[str, Any], manifest: StageManifest): def build_library_stage(config: dict[str, Any], manifest: StageManifest): + """ + Build the replacement and candidate libraries and record their associated statistics. + + The stage validates and shares the resolved scoring parent, optionally refreshes runtime + statistics, calculates static workload statistics, and publishes the resulting artifact + paths and execution metadata. + + Parameters: + config (dict[str, Any]): Pipeline configuration. + manifest (StageManifest): Manifest used to record stage completion and outputs. + + Returns: + StageManifest: Updated manifest containing the generated library paths, statistics + metadata, and scoring-parent information. + """ hydra_cfg = load_runtime_hydra_config(config) puzzle_dir = _puzzle_dir(config, hydra_cfg) candidate_library_path = puzzle_dir / "candidate_library.json" diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py index 38580f156b7..9677b5c55c1 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py @@ -276,7 +276,16 @@ def _runtime_reuse_key_from_args( *, fallback_workload_id: str | None = None, ) -> tuple | None: - """Return the exact runtime-reuse identity represented by persisted args.""" + """ + Builds the identity used to match reusable runtime measurements. + + Parameters: + args (Mapping): Persisted calculation arguments containing runtime and workload settings. + fallback_workload_id (str | None): Workload identifier to use when `args` does not provide one. + + Returns: + tuple | None: Runtime-reuse identity, or `None` when runtime statistics are unavailable or the dtype is not bfloat16. + """ if not args.get("runtime_stats") or args.get("weights_dtype") != str(torch.bfloat16): return None @@ -307,7 +316,15 @@ def _runtime_reuse_key( generation_seq_len: int, runtime_stats_config: Mapping, ) -> tuple: - """Return the exact runtime-reuse identity requested by this calculation.""" + """ + Builds an exact identity for matching reusable runtime measurements. + + Parameters: + runtime_stats_config (Mapping): Runtime settings that determine measurement compatibility. + + Returns: + tuple: Identity containing model dimensions, runtime settings, vLLM arguments, and workload identity. + """ return ( int(width), @@ -332,7 +349,21 @@ def _reuse_runtime_stats( source_path: str, fallback_workload_id: str | None = None, ) -> dict: - """Overlay immutable measured latency onto refreshed static statistics.""" + """ + Reuse measured runtime statistics from a compatible source entry in refreshed statistics. + + Parameters: + target (dict): Statistics entry to update with reusable runtime data. + source (dict): Statistics entry containing the measured runtime data. + source_path (str): Path identifying the source statistics entry. + fallback_workload_id (str | None): Workload identifier to use when the source does not provide one. + + Returns: + dict: The updated target statistics entry. + + Raises: + KeyError: If a target subblock has no matching runtime statistics in the source entry. + """ # Synthetic vLLM timings are layer-independent: collection benchmarks the # set of unique subblock configs, while a post-scoring replacement library @@ -826,6 +857,33 @@ def calculate_subblock_stats( runtime_selection_identity: str | None = None, parameter_inventory: Mapping | None = None, ) -> dict: + """ + Compute parameter, memory, additive-metric, and optional runtime statistics for subblock configurations. + + Parameters: + calc_subblock_stats_config (DictConfig): Runtime measurement and calculation settings. + teacher_dir (Path): Directory containing the teacher model or checkpoint. + model_config (PretrainedConfig): Model configuration used for metric calculations. + descriptor (Type[ModelDescriptor]): Model descriptor defining architecture-specific behavior. + master_puzzle_dir (Path): Puzzle directory used for runtime measurement caches. + subblock_configs (list[immutabledict[str, SubblockConfig]]): Subblock configurations and their parent layer indices. + batch_size (int): Number of sequences in the workload. + prefill_seq_len (int): Input sequence length used for prefill calculations. + generation_seq_len (int): Number of generated tokens used for decode calculations. + n_embd (int): Model hidden size. + n_head (int): Number of attention heads. + vocab_size (int): Model vocabulary size. + runtime_stats_enabled (bool): Whether to measure runtime statistics. + use_cuda_graph (bool): Whether to use CUDA graphs during runtime measurement. + weights_dtype (torch.dtype): Data type used for model weights. + activations_dtype (torch.dtype): Data type used for activations. + kv_cache_dtype (torch.dtype): Data type used for the key-value cache. + runtime_selection_identity (str | None): Identity of the runtime subblock selection. + parameter_inventory (Mapping | None): Precomputed parameter inventory to use for parameter counts. + + Returns: + dict: Statistics for the requested workload, including calculation arguments, non-block statistics, and per-subblock metrics. + """ runtime_granularity = "subblock" runtime_stats_config = ( calc_subblock_stats_config.get("runtime_stats", {}) if runtime_stats_enabled else {} @@ -1246,12 +1304,26 @@ def _subblock_stats_already_complete( prefill_seq_len: int = 2048, generation_seq_len: int = 2048, ) -> bool: - """Whether ``existing_stats`` already covers every configuration this run would compute. - - When runtime benchmarking is enabled, the bf16 entries (the only ones for - which runtime is ever measured) must additionally already carry runtime - measurements **at the requested granularity** — switching subblock<->block must trigger a - recompute rather than silently reusing the other granularity's numbers. + """ + Determine whether existing statistics cover all requested configurations. + + Parameters: + existing_stats (list): Previously calculated statistics entries. + subblock_configs (list): Subblock configurations required for coverage. + batch_sizes (Iterable[int]): Batch sizes to verify. + data_types (list): Weight, activation, and KV-cache dtype combinations. + model_hidden_sizes (Iterable[int]): Model widths to verify. + runtime_stats_enabled (bool): Whether runtime measurements are required. + runtime_granularity (str): Required runtime measurement granularity. + runtime_max_num_seqs (int | None): Required maximum number of runtime sequences. + runtime_workload_id (str | None): Required runtime workload identity. + runtime_selection_identity (str | None): Required runtime subblock-selection identity. + parameter_inventory_identities (Mapping[int, str] | None): Inventory identity required for each model width. + prefill_seq_len (int): Required prefill sequence length. + generation_seq_len (int): Required generation sequence length. + + Returns: + bool: True if every requested configuration and required measurement is present, False otherwise. """ by_signature = {_arg_signature(entry["args"]): entry for entry in existing_stats} required_subblock_keys = { @@ -1354,6 +1426,27 @@ def calculate_subblock_stats_for_puzzle_dir( # from attach_helper import debugging_setup # debugging_setup() # You can optionally pass a name to identify the job (e.g. `debugging_setup(name="my_script")`) # ==== END === Setup for attach-helper ==== + """ + Compute and persist subblock statistics for all requested batch sizes, data types, and model widths. + + Parameters: + calc_subblock_stats_config (DictConfig): Configuration for statistics calculation and optional runtime measurement. + master_puzzle_dir (Path | str): Puzzle directory containing subblock configurations and output files. + teacher_dir (Path | str): Teacher checkpoint directory used for model metadata and parameter inventories. + descriptor (Type[ModelDescriptor]): Model descriptor defining architecture-specific behavior. + model_hidden_sizes (ListConfig): Hidden sizes to evaluate; the teacher hidden size is always included. + ffn_hidden_sizes (ListConfig): Additional FFN sizes to include in the subblock configurations. + batch_sizes (Iterable[int]): Batch sizes to evaluate. + prefill_seq_len (int): Number of prompt tokens used for runtime measurements. + generation_seq_len (int): Number of generated tokens used for runtime measurements. + runtime_stats_enabled (bool): Whether to compute or reuse runtime statistics. + merge_with_existing_stats (bool): Whether to update an existing incomplete statistics file. + subblock_stats_filename (str): Name of the JSON file used to persist statistics. + + Raises: + FileNotFoundError: If a configured runtime manifest or reusable runtime statistics file cannot be found. + ValueError: If runtime settings or reusable runtime statistics do not cover the requested configurations. + """ if isinstance(batch_sizes, str): batch_sizes = [ int(batch_size) for batch_size in batch_sizes.strip("[]").replace(" ", "").split(",") diff --git a/puzzletron_setup/bundle.py b/puzzletron_setup/bundle.py index 3ee3e8f0936..13bad12f5ef 100644 --- a/puzzletron_setup/bundle.py +++ b/puzzletron_setup/bundle.py @@ -329,6 +329,19 @@ def _post_mip_flows( global_kd_mesh: Mapping[str, Any], default_serving_topology: Mapping[str, Any], ) -> dict[str, Any]: + """ + Prepare post-MIP flow configurations with mesh settings, serving defaults, and smoke-run limits. + + Parameters: + state (Mapping[str, Any]): Campaign state containing post-MIP flow definitions. + smoke (bool): Whether to apply reduced settings for a smoke run. + common_mesh (Mapping[str, Any]): Mesh used by evaluation and materialization nodes. + global_kd_mesh (Mapping[str, Any]): Mesh used by global knowledge-distillation nodes. + default_serving_topology (Mapping[str, Any]): Default topology for serving-based nodes. + + Returns: + dict[str, Any]: The normalized post-MIP flow configurations. + """ flows = deepcopy(_mapping(_answers(state, "post_mip").get("flows"))) for flow in flows.values(): for node in _mapping(flow.get("nodes")).values(): @@ -795,6 +808,21 @@ def _dynamic_stage_entries( *, pool_source_evaluations: bool, ) -> dict[str, Any]: + """ + Builds scheduler entries for dynamic post-MIP stages. + + Parameters: + experiment (Mapping[str, Any]): Experiment configuration containing post-MIP flows. + workers (Mapping[str, Any]): Worker limits for pooled and sharded stages. + gpus_per_node (int): Number of GPUs assigned to each node. + common (Mapping[str, Any]): Parallel configuration for evaluation stages. + single_gpu (Mapping[str, Any]): Parallel configuration for materialization stages. + cpu_partition (str | None): CPU partition to assign to CPU stages. + pool_source_evaluations (bool): Whether source evaluations should use pooled workers. + + Returns: + dict[str, Any]: Scheduler entries keyed by post-MIP flow and node identifiers. + """ entries = {} candidate_limits = _post_mip_candidate_limits(experiment) for flow_id, flow in _mapping(_mapping(experiment.get("post_mip")).get("flows")).items(): diff --git a/puzzletron_setup/v2/validation.py b/puzzletron_setup/v2/validation.py index d0cf216d693..787c2989ec7 100644 --- a/puzzletron_setup/v2/validation.py +++ b/puzzletron_setup/v2/validation.py @@ -188,7 +188,12 @@ def _dataset_subset_issues(state: WizardState) -> list[ValidationIssue]: def validate_state(state: WizardState) -> tuple[ValidationIssue, ...]: - """Return actionable authoring issues before canonical compilation.""" + """ + Validate wizard state and identify issues that prevent canonical compilation. + + Returns: + tuple[ValidationIssue, ...]: Validation issues sorted by configuration section and path. + """ issues: list[ValidationIssue] = [] required = ( "model.source", diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index ded6bbbff4c..52d752d0bb5 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -3767,6 +3767,17 @@ def _post_mip_strategy(node: NodeDraft) -> str: def post_mip_section(session: WizardSession, resolver: DefaultsResolver, context: dict) -> bool: + """ + Configure post-MIP execution flows for each MIP run, using recommended or custom nodes. + + Parameters: + session (WizardSession): Wizard session used to read state and collect configuration. + resolver (DefaultsResolver): Resolver for stage resource defaults. + context (dict): Model and pruning context required to configure serving and evaluation nodes. + + Returns: + bool: `True` when post-MIP flows are configured, `False` when the section is exited through back navigation. + """ mip = _mapping_copy(session.state.collection("mip_config")) runs = _mapping_copy(mip.get("runs")) sequence = int(session.state.get_field("data.sequence_length", 4096)) @@ -4193,7 +4204,22 @@ def _serving_setting_prompt( pruning: Mapping[str, Any], stage_id: str, ) -> Any: - """Ask the complete AIPerf workload and serving-only parallel setting.""" + """ + Collect AIPerf serving workload settings and the vLLM serving topology. + + Parameters: + session (WizardSession): Wizard session used to collect and validate responses. + prefix (str): State key prefix for the serving settings. + defaults (Mapping[str, Any]): Default workload and topology values. + inventory (Any): Model inventory used to validate the topology. + pruning (Mapping[str, Any]): Pruning configuration relevant to topology validation. + stage_id (str): Pruning stage associated with the serving configuration. + + Returns: + Any: A mapping containing input and output sequence lengths, concurrency values, + request count, model selection mode, and topology, or the `BACK` sentinel when + the user navigates to the previous prompt. + """ values = {} for name, label, default in ( ("input_tokens", "Serving input sequence length (ISL):", defaults["input_tokens"]), @@ -4270,9 +4296,30 @@ def _downstream_evaluation_setting_prompt( pruning: Mapping[str, Any], stage_id: str, ) -> Any: - """Ask lmms-eval task settings and the vLLM topology used to run them.""" + """ + Collect lmms-eval tasks, execution settings, model arguments, and vLLM topology. + + Parameters: + session (WizardSession): Wizard session used to prompt for settings. + prefix (str): State-key prefix for the prompted values. + defaults (Mapping[str, Any]): Existing values used as prompt defaults. + inventory (Any): Model inventory used to validate the vLLM topology. + pruning (Mapping[str, Any]): Pruning configuration relevant to topology validation. + stage_id (str): Identifier of the stage using the evaluation settings. + + Returns: + Any: A mapping containing lmms-eval tasks, sample and batch limits, timeout, model arguments, logging settings, and vLLM topology, or `BACK` if prompting is cancelled. + """ def validate_tasks(value: str) -> bool | str: + """Validate a comma-separated list of lmms-eval tasks. + + Parameters: + value (str): Comma-separated task names. + + Returns: + bool | str: `True` if at least one task is provided, otherwise an error message. + """ tasks = [item.strip() for item in value.split(",") if item.strip()] return True if tasks else "Enter at least one lmms-eval task." @@ -4346,7 +4393,16 @@ def _configure_dynamic_resources( *, ask: bool, ) -> Any: - """Attach an independent resource/batch card to every node in one flow.""" + """ + Configure independent resource assignments for all nodes in a post-MIP flow. + + Parameters: + flow_id (str): Identifier of the flow whose nodes are configured. + ask (bool): Whether to prompt for resource and batch customizations. + + Returns: + True when configuration completes, or `BACK` when navigation is requested. + """ registry = ResourceProfileRegistry.from_dict( session.state.collection("parallel_profiles") or {} ) diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index 2750be52f6f..c6140623757 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -640,7 +640,18 @@ def _ask_aiperf_config( runtime: Mapping[str, Any], defaults: Mapping[str, Any] | None = None, ) -> dict[str, Any]: - """Ask for one AIPerf node's independent Serving topology and workload.""" + """ + Configure an AIPerf serving node's parallel topology and workload settings. + + Parameters: + detailed (bool): Whether to prompt for workload and timeout values. + moe (bool): Whether to configure expert parallelism for a mixture-of-experts model. + runtime (Mapping[str, Any]): Runtime defaults for input length, output length, and concurrency. + defaults (Mapping[str, Any] | None): Previously saved configuration values. + + Returns: + dict[str, Any]: The configured AIPerf topology, workload, and timeout settings. + """ defaults = dict(defaults or {}) topology_defaults = dict(defaults.get("topology") or {}) checkpoint = prompts.checkpoint() @@ -740,7 +751,17 @@ def _ask_downstream_evaluation_config( moe: bool, defaults: Mapping[str, Any] | None = None, ) -> dict[str, Any]: - """Ask for lmms-eval task and vLLM settings.""" + """ + Collect lmms-eval tasks, sampling settings, vLLM topology, and evaluation timeout. + + Parameters: + detailed (bool): Whether to prompt for the per-candidate timeout. + moe (bool): Whether to allow configuring expert parallelism. + defaults (Mapping[str, Any] | None): Previously saved settings used as prompt defaults. + + Returns: + dict[str, Any]: The configured downstream evaluation settings. + """ defaults = defaults or {} tasks = prompts.text( @@ -851,6 +872,21 @@ def _default_flow( objective: Mapping[str, Any] | None = None, include_initial_filter: bool = True, ) -> dict[str, Any]: + """ + Build the standard post-MIP evaluation and selection flow. + + Parameters: + run_id (str): Identifier of the MIP run. + run (Mapping[str, Any]): MIP run configuration. + runtime (Mapping[str, Any]): Runtime settings for serving evaluation. + data (Mapping[str, Any]): Dataset settings, including sequence length. + prefix (str): Prefix applied to generated node identifiers. + objective (Mapping[str, Any] | None): Objective used to configure ranking; the run's first objective is used when omitted. + include_initial_filter (bool): Whether to include the initial MIP-score filter. + + Returns: + dict[str, Any]: Flow configuration containing the source metadata and ordered post-MIP nodes. + """ def node_id(name: str) -> str: return f"{prefix}{name}" @@ -1010,6 +1046,20 @@ def _custom_flow( detailed: bool, moe: bool, ) -> dict[str, Any]: + """ + Build a custom post-MIP evaluation flow through interactive configuration. + + Parameters: + run_id (str): Identifier of the MIP run supplying candidate models. + runtime (Mapping[str, Any]): Runtime settings used by serving evaluations. + data (Mapping[str, Any]): Dataset settings used by evaluation nodes. + used_ids (set[str]): Node IDs already in use; newly configured IDs are added. + detailed (bool): Whether to collect detailed evaluation settings. + moe (bool): Whether to enable mixture-of-experts configuration options. + + Returns: + dict[str, Any]: A flow definition containing the MIP source and configured nodes. + """ nodes: OrderedDict[str, Any] = OrderedDict() available_metrics = ["mip.score"] transformer_nodes = [] @@ -1226,6 +1276,20 @@ def _resource_rows( gpus_per_node: int, workers: Mapping[str, int], ) -> list[dict[str, Any]]: + """ + Calculate resource requirements for each campaign execution stage. + + Parameters: + state (AnswerState): Campaign configuration containing post-MIP flows and execution details. + common (Mapping[str, int]): Parallel mesh dimensions shared by common stages. + bypass (Mapping[str, int]): Parallel mesh dimensions for bypass processing. + global_kd (Mapping[str, int]): Parallel mesh dimensions for global knowledge distillation. + gpus_per_node (int): Number of GPUs available on each node. + workers (Mapping[str, int]): Worker limits for pool and sharded stages. + + Returns: + list[dict[str, Any]]: Resource rows containing each stage's name, instance count, GPUs per instance, and required node count. + """ from .bundle import _post_mip_candidate_limits, _serving_parallel rows = [] From 5983fd6711f11bbe4acc7101d4963f4a979e8b4e Mon Sep 17 00:00:00 2001 From: "coderabbitai[bot]" <136622811+coderabbitai[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 12:33:43 +0200 Subject: [PATCH 06/16] fix: CodeRabbit auto-fixes for PR #2104 (#2106) This stacked PR contains CodeRabbit auto-fixes for #2104. **Files modified:** - `modelopt/torch/puzzletron/post_mip/runner.py` - `modelopt/torch/puzzletron/stages/pipeline.py` - `modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py` - `puzzletron_setup/wizard.py` - `tests/unit/torch/puzzletron/test_sparse_runtime_stats.py` Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: CodeRabbit Signed-off-by: Johannes Rausch --- modelopt/torch/puzzletron/post_mip/runner.py | 41 ++++++++++++------- modelopt/torch/puzzletron/stages/pipeline.py | 41 ++++++++----------- .../subblock_stats/calc_subblock_stats.py | 13 +++--- puzzletron_setup/wizard.py | 9 +++- .../puzzletron/test_sparse_runtime_stats.py | 5 ++- 5 files changed, 63 insertions(+), 46 deletions(-) diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 3d23e0d05b0..0511091b87e 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -45,6 +45,8 @@ "run_post_mip_node_shard", ] +_DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS = 3600.0 + def _puzzle_dir(config: Mapping[str, Any]) -> Path: return Path(config.get("puzzle_dir") or (config.get("experiment") or {})["dir"]) @@ -672,19 +674,18 @@ def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> ], } ) - for key in _LMMS_EVAL_MODEL_ARG_FIELDS: + for key in sorted(_LMMS_EVAL_MODEL_ARG_FIELDS): if key in settings: derived[key] = settings[key] if isinstance(raw, str): prefix = raw.strip().strip(",") suffix = _model_arg_string(derived) - return ",".join(part for part in (prefix, suffix) if part) + return ",".join(part for part in (suffix, prefix) if part) if raw is not None and not isinstance(raw, Mapping): raise TypeError("downstream_evaluation.config.model_args must be a mapping or string") merged = dict(raw or {}) - for key, value in derived.items(): - merged.setdefault(key, value) + merged.update(derived) return _model_arg_string(merged) @@ -779,7 +780,7 @@ def _lmms_eval_command( if settings.get("cache_dir") is not None: env.setdefault("LMMS_EVAL_HOME", str(settings["cache_dir"])) timeout = settings.get("timeout_seconds", settings.get("timeout")) - return argv, env, (float(timeout) if timeout is not None else None) + return argv, env, (float(timeout) if timeout is not None else _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS) def _metric_key(value: Any) -> str: @@ -949,15 +950,25 @@ def _downstream_evaluation( ) # Campaign config controls the executable and arguments, but subprocess receives # an argv list directly; no shell parsing is involved. - result = subprocess.run( - argv, - cwd=str(output), - env=env, - capture_output=True, - text=True, - timeout=timeout, - check=False, - ) + try: + result = subprocess.run( + argv, + cwd=str(output), + env=env, + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + except subprocess.TimeoutExpired as timeout_error: + timeout_result = subprocess.CompletedProcess( + args=argv, + returncode=-1, + stdout=timeout_error.stdout.decode("utf-8", errors="replace") if timeout_error.stdout else "", + stderr=timeout_error.stderr.decode("utf-8", errors="replace") if timeout_error.stderr else "", + ) + _write_lmms_eval_streams(output, timeout_result) + raise stream_paths = _write_lmms_eval_streams(output, result) if result.returncode: tail = _lmms_eval_output_tail(result) @@ -1165,7 +1176,7 @@ def run_post_mip_node_shard( timeout_field = "timeout_seconds" elif not isinstance(error, subprocess.TimeoutExpired): timeout_field = "readiness_timeout" - default_timeout = 3600 if node.node_type == "downstream_evaluation" else ( + default_timeout = _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS if node.node_type == "downstream_evaluation" else ( 600 if timeout_field == "benchmark_timeout" else 1200 ) row["timeout_seconds"] = float( diff --git a/modelopt/torch/puzzletron/stages/pipeline.py b/modelopt/torch/puzzletron/stages/pipeline.py index 279aec34744..8a40ad5e52c 100644 --- a/modelopt/torch/puzzletron/stages/pipeline.py +++ b/modelopt/torch/puzzletron/stages/pipeline.py @@ -275,16 +275,18 @@ def _has_runtime_measurement( ) -> bool: """ Determine whether a statistics file contains a compatible runtime measurement. - + Parameters: path (Path): Statistics file to inspect. hidden_width (int): Model hidden width expected by the measurement. measurement (Any): Runtime measurement configuration to match. allow_missing_workload_id (bool): Whether entries without a workload identifier may match. - + Returns: bool: `True` if a compatible runtime measurement is present, `False` otherwise. """ + from ..subblock_stats.calc_subblock_stats import _runtime_reuse_key, _runtime_stats_identity + try: payload = json.loads(path.read_text()) except (OSError, ValueError): @@ -292,34 +294,27 @@ def _has_runtime_measurement( if not isinstance(payload, list): return False expected_backend = (measurement.runtime_stats or {}).get("backend") + requested_key = _runtime_reuse_key( + width=hidden_width, + batch_size=measurement.batch_size, + prefill_seq_len=measurement.prefill_seq_len, + generation_seq_len=measurement.generation_seq_len, + runtime_stats_config=measurement.runtime_stats or {}, + ) for entry in payload: if not isinstance(entry, dict): continue args = entry.get("args") or {} if not isinstance(args, dict) or args.get("runtime_stats") is not True: continue - if int(args.get("n_embd", -1)) != int(hidden_width): - continue - if args.get("weights_dtype") != "torch.bfloat16": - continue - if int(args.get("batch_size", -1)) != int(measurement.batch_size): - continue - if int(args.get("prefill_seq_len", -1)) != int(measurement.prefill_seq_len): - continue - if int(args.get("generation_seq_len", -1)) != int(measurement.generation_seq_len): - continue - if int(args.get("max_num_seqs", -1)) != int(measurement.max_num_seqs): - continue - if args.get("runtime_granularity", "subblock") != measurement.granularity: - continue - if expected_backend is not None and args.get("runtime_backend") != expected_backend: - continue - workload_id = args.get("workload_id") - if workload_id is None and not allow_missing_workload_id: - continue - if workload_id is not None and workload_id != measurement.measurement_id: + persisted_key = _runtime_stats_identity( + args, + fallback_workload_id=measurement.measurement_id if allow_missing_workload_id else None, + ) + if persisted_key is None: continue - return True + if persisted_key == requested_key: + return True return False diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py index 9677b5c55c1..127f4b2d8c4 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py @@ -289,20 +289,23 @@ def _runtime_reuse_key_from_args( if not args.get("runtime_stats") or args.get("weights_dtype") != str(torch.bfloat16): return None + required_fields = ["n_embd", "batch_size", "prefill_seq_len", "generation_seq_len"] + if any(args.get(field) is None for field in required_fields): + return None workload_id = args.get("workload_id") if workload_id is None: workload_id = fallback_workload_id return ( int(args["n_embd"]), int(args["batch_size"]), - int(args.get("prefill_seq_len")), - int(args.get("generation_seq_len")), + int(args["prefill_seq_len"]), + int(args["generation_seq_len"]), args.get("max_num_seqs"), args.get("runtime_granularity", "subblock"), args.get("runtime_backend"), - args.get("num_iters"), - args.get("num_warmup_iters"), - args.get("repeat_block_n_times"), + args.get("num_iters", 30), + args.get("num_warmup_iters", 10), + max(2, int(args.get("repeat_block_n_times", 10))), _freeze_stats_args(args.get("vllm_args")), workload_id, ) diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index c6140623757..4b092569c34 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -764,9 +764,13 @@ def _ask_downstream_evaluation_config( """ defaults = defaults or {} + default_tasks = defaults.get("tasks", "ifeval,gsm8k") + if isinstance(default_tasks, list): + default_tasks = ",".join(default_tasks) tasks = prompts.text( "lmms-eval tasks (comma-separated):", - default=str(defaults.get("tasks", "ifeval,gsm8k")), + default=str(default_tasks), + validate=lambda value: bool(str(value).strip()) or "Enter at least one task.", ) limit = prompts.integer( "lmms-eval sample limit:", @@ -1127,7 +1131,8 @@ def _custom_flow( detailed=detailed, moe=moe, ) - available_metrics.append(f"{node_id}.gsm8k.exact_match") + for task_name in node["config"].get("tasks", []): + available_metrics.append(f"{node_id}.{task_name}.strict-match") elif node_type == "global_kd": node["config"] = {"max_steps": prompts.integer("Global KD steps:", default=128)} elif node_type == "ptq": diff --git a/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py b/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py index 4e9a6e39c6d..87f21abdbba 100644 --- a/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py +++ b/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py @@ -1261,6 +1261,10 @@ def test_width_scenario_runtime_stats_reuse_root_measurement(tmp_path, monkeypat "generation_seq_len": 1024, "max_num_seqs": 1, "n_embd": 2688, + "num_iters": 30, + "num_warmup_iters": 10, + "repeat_block_n_times": 10, + "vllm_args": [], "workload_id": "serving-default", }, "subblocks": [], @@ -1408,7 +1412,6 @@ def test_runtime_stats_resume_signature_includes_workload_id(): **kwargs, runtime_workload_id="different-workload", ) - assert hydra_cfg.calc_subblock_stats.merge_with_existing_stats is False def test_sparse_runtime_selection_is_unique_and_layer_independent(): From 59bcf30158ca72691f63e3bdcbeda968a312cf56 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Fri, 7 Aug 2026 12:00:16 -0700 Subject: [PATCH 07/16] drop docstring churn Signed-off-by: Grzegorz Karch --- .../torch/puzzletron/benchmarks/aiperf.py | 60 +---- .../puzzletron/orchestration/adapters/pool.py | 14 -- .../orchestration/adapters/post_mip.py | 13 - .../puzzletron/orchestration/compiler.py | 17 +- .../puzzletron/orchestration/controller.py | 23 -- .../orchestration/executors/slurm.py | 30 +-- .../puzzletron/orchestration/progress.py | 11 - .../puzzletron/orchestration/task_launcher.py | 37 +-- modelopt/torch/puzzletron/post_mip/builtin.py | 1 - .../torch/puzzletron/post_mip/reporting.py | 23 +- modelopt/torch/puzzletron/post_mip/runner.py | 224 ++---------------- modelopt/torch/puzzletron/stages/pipeline.py | 119 ++-------- .../subblock_stats/calc_subblock_stats.py | 124 ++-------- puzzletron_setup/bundle.py | 28 --- puzzletron_setup/v2/validation.py | 7 +- puzzletron_setup/v2/wizard.py | 62 +---- puzzletron_setup/wizard.py | 77 +----- .../puzzletron/test_sparse_runtime_stats.py | 5 +- 18 files changed, 74 insertions(+), 801 deletions(-) diff --git a/modelopt/torch/puzzletron/benchmarks/aiperf.py b/modelopt/torch/puzzletron/benchmarks/aiperf.py index ebbeea66fde..5824247f8be 100644 --- a/modelopt/torch/puzzletron/benchmarks/aiperf.py +++ b/modelopt/torch/puzzletron/benchmarks/aiperf.py @@ -114,14 +114,7 @@ def _canonical_topology(topology: dict[str, Any]) -> dict[str, Any]: def _topology_vllm_args(topology: dict[str, Any]) -> list[str]: - """Convert canonical parallelism settings into vLLM command-line arguments. - - Parameters: - topology (dict[str, Any]): Topology configuration to normalize and convert. - - Returns: - list[str]: vLLM command-line arguments for the configured tensor, pipeline, context, data, and expert parallelism. - """ + """Translate the canonical TP/PP/DP/EP/CP contract to vLLM CLI arguments.""" canonical = _canonical_topology(topology) args = [ @@ -171,17 +164,7 @@ def _has_vllm_option(args: Iterable[str], *options: str) -> bool: def _server_vllm_args( checkpoint_dir: Path, topology: dict[str, Any], concurrency_values: Iterable[int] ) -> list[str]: - """ - Build vLLM server arguments from checkpoint configuration, topology, and concurrency demand. - - Parameters: - checkpoint_dir (Path): Directory containing the model checkpoint. - topology (dict[str, Any]): Topology and extra vLLM argument configuration. - concurrency_values (Iterable[int]): Requested concurrency levels used to set the default maximum number of sequences. - - Returns: - list[str]: Command-line arguments for the vLLM server. - """ + """Build stable vLLM server args derived from topology and benchmark demand.""" args = _topology_vllm_args(topology) args.extend(_descriptor_vllm_args(checkpoint_dir)) @@ -195,16 +178,7 @@ def _server_vllm_args( def _exact_length_extra_inputs( extra_inputs: dict[str, Any] | None, output_tokens: int ) -> dict[str, Any]: - """ - Configure extra input settings for exact output-length measurements. - - Parameters: - extra_inputs (dict[str, Any] | None): Optional caller-provided input settings. - - Returns: - dict[str, Any]: The input settings with ``ignore_eos`` enabled when neither - ``ignore_eos`` nor ``min_tokens`` was explicitly provided. - """ + """Guarantee the measured OSL unless the caller chose an explicit policy.""" resolved = dict(extra_inputs or {}) if "ignore_eos" not in resolved and "min_tokens" not in resolved: resolved["ignore_eos"] = True @@ -435,33 +409,7 @@ def run_aiperf_sweep( benchmark_timeout: float = 600, gpu_telemetry: str | None = "pynvml", ) -> list[BenchmarkResult]: - """ - Run multiple concurrency benchmarks against a persistent vLLM server. - - Parameters: - checkpoint_dir (str | Path): Model checkpoint directory. - artifact_dir (str | Path): Directory for benchmark outputs and logs. - concurrencies (Iterable[int]): Unique positive concurrency levels to benchmark. - input_tokens (int): Synthetic input length for each request. - output_tokens (int): Synthetic output length for each request. - gpu_ids (str): GPU visibility specification for the benchmark processes. - topology (dict[str, Any]): vLLM parallelism and environment configuration. - request_counts (dict[int, int] | None): Optional request count for each concurrency level. - solution_id (str): Identifier for the benchmark solution. - profile_id (str): Identifier for the benchmark profile. - topology_id (str | None): Optional topology identifier. - executable (str | Path): AIPerf executable to run. - endpoint_type (str): Endpoint type used by AIPerf. - extra_inputs (dict[str, Any] | None): Additional AIPerf input settings. - use_server_token_count (bool): Whether to use token counts reported by the server. - seed (int): Seed for synthetic request generation. - readiness_timeout (float): Maximum time to wait for vLLM readiness, in seconds. - benchmark_timeout (float): Maximum time allowed for each benchmark, in seconds. - gpu_telemetry (str | None): GPU telemetry backend, or None to disable telemetry. - - Returns: - list[BenchmarkResult]: Benchmark results in the original concurrency order. - """ + """Run multiple concurrencies against one persistent vLLM server.""" checkpoint_dir = Path(checkpoint_dir).resolve() artifact_dir = Path(artifact_dir).resolve() diff --git a/modelopt/torch/puzzletron/orchestration/adapters/pool.py b/modelopt/torch/puzzletron/orchestration/adapters/pool.py index 098c05261b6..7abee0a5263 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/pool.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/pool.py @@ -224,20 +224,6 @@ def command( runner, overrides: list[str] | None = None, ) -> AttemptSpec: - """ - Build the execution command and resource allocation for a planned work item. - - Parameters: - plan (CampaignPlan): Campaign configuration and execution context. - node (StagePlanNode): Stage and resource configuration for the work item. - item (WorkItem): Work item metadata, role, and local GPU assignments. - attempt_id (str): Identifier for the execution attempt. - runner: Runner context containing the repository location. - overrides (list[str] | None): Optional configuration overrides to apply. - - Returns: - AttemptSpec: Command, environment, resource allocation, and execution metadata for the work item. - """ repo = Path(runner.contract.repository) role = item.metadata.get("role", "worker") log_dir = plan.puzzle_dir / "logs" diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index 6887a744beb..90f1f745f9a 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py @@ -108,19 +108,6 @@ class PostMIPAdapter(WorkAdapter): strategy = ExecutionStrategy.SHARDED def plan(self, plan: CampaignPlan, node: StagePlanNode) -> WorkPlan: - """ - Plan sharded execution for a post-MIP node and mark aggregation as required. - - Parameters: - plan (CampaignPlan): Campaign execution plan containing node configuration and candidate information. - node (StagePlanNode): Post-MIP node to plan. - - Returns: - WorkPlan: Work plan containing the node's work item and execution strategy. - - Raises: - RuntimeError: If an evaluation node has no candidate architectures to evaluate. - """ config = _node_config(plan, node.stage_id) node_type = str(config.get("type")) count = 1 if node_type in {"filter", "manual_filter"} else node.instances diff --git a/modelopt/torch/puzzletron/orchestration/compiler.py b/modelopt/torch/puzzletron/orchestration/compiler.py index a9b80ed2de7..5c76c3ce087 100644 --- a/modelopt/torch/puzzletron/orchestration/compiler.py +++ b/modelopt/torch/puzzletron/orchestration/compiler.py @@ -431,22 +431,7 @@ def compile_campaign_plan( overrides: list[str] | None = None, stage_filter: str | None = None, ) -> CampaignPlan: - """ - Compile a campaign plan from experiment, runner, and execution configurations. - - Parameters: - experiment_config_path: Path to the experiment configuration file. - runner: Runner environment used to execute the campaign. - execution: Execution defaults and per-stage settings. - overrides: Optional experiment configuration overrides. - stage_filter: Optional stage identifier limiting the plan to one enabled stage. - - Returns: - A compiled campaign plan containing stage meshes, dependencies, resources, and GPU allocations. - - Raises: - ValueError: If the selected stage is disabled, a CPU stage requests multiple instances, or a mesh override conflicts with its topology. - """ + """Compile one campaign plan from experiment + runner + execution configs.""" experiment_path = Path(experiment_config_path) experiment_config = load_experiment_config(experiment_path, overrides=overrides or []) diff --git a/modelopt/torch/puzzletron/orchestration/controller.py b/modelopt/torch/puzzletron/orchestration/controller.py index aa6a3c04877..71cf6f63cf5 100644 --- a/modelopt/torch/puzzletron/orchestration/controller.py +++ b/modelopt/torch/puzzletron/orchestration/controller.py @@ -63,19 +63,6 @@ def create_executor(plan: CampaignPlan, *, local: bool = False) -> Executor: - """ - Create an executor for the campaign plan's configured runner. - - Parameters: - plan (CampaignPlan): Campaign plan containing runner configuration. - local (bool): Whether to use a local executor instead of the configured runner. - - Returns: - Executor: Executor configured for local, Slurm, or bare-metal SSH execution. - - Raises: - ValueError: If the configured runner kind is unsupported. - """ if local: return LocalExecutor(plan.runner) if plan.runner.kind == "slurm": @@ -93,16 +80,6 @@ def _stage_dashboard_display_name( *, granularity: str | None = None, ) -> str: - """Resolve the dashboard display name for a campaign stage. - - Parameters: - config (Mapping[str, Any]): Campaign configuration containing post-MIP flow definitions. - stage_id (str): Stage identifier to format. - granularity (str | None): Optional naming granularity. - - Returns: - str: ``"Downstream Evaluation"`` for downstream evaluation post-MIP stages; otherwise, the formatted stage name. - """ if stage_id.startswith("post."): parts = stage_id.split(".", 2) if len(parts) == 3: diff --git a/modelopt/torch/puzzletron/orchestration/executors/slurm.py b/modelopt/torch/puzzletron/orchestration/executors/slurm.py index b94e2c7319f..047c2a0ce9f 100644 --- a/modelopt/torch/puzzletron/orchestration/executors/slurm.py +++ b/modelopt/torch/puzzletron/orchestration/executors/slurm.py @@ -91,15 +91,7 @@ def render_hook_lines(commands: Sequence[str]) -> str: def _render_host_container_env(repository: str) -> str: - """ - Render shell commands that configure default Pyxis/Enroot paths and create their directories. - - Parameters: - repository (str): Repository path used to derive default Enroot cache and data paths. - - Returns: - str: Shell commands for configuring and preparing the container runtime environment. - """ + """Render host-side container runtime defaults for Pyxis/Enroot.""" cache_root = Path(repository) / ".cache" / "enroot" lines = [ @@ -133,25 +125,7 @@ def render_sbatch_script( qos: str | None, job_name: str, ) -> str: - """ - Render an executable Slurm batch script for an attempt, including resource - allocations, environment setup, hooks, logging, and optional container - configuration. - - Parameters: - attempt (AttemptSpec): Attempt specification containing the command and - requested task topology. - runner (RunnerEnvironment): Runner configuration used for repository, - environment, and container settings. - partition (str): Slurm partition for the job. - account (str): Slurm account for the job. - time_limit (str): Slurm time limit. - qos (str | None): Optional Slurm quality-of-service name. - job_name (str): Name assigned to the Slurm job. - - Returns: - str: The generated executable sbatch script. - """ + """Render one sbatch script for an attempt.""" contract = runner.contract topology = resolve_task_topology(attempt) diff --git a/modelopt/torch/puzzletron/orchestration/progress.py b/modelopt/torch/puzzletron/orchestration/progress.py index c89c276d66a..a9c943357e6 100644 --- a/modelopt/torch/puzzletron/orchestration/progress.py +++ b/modelopt/torch/puzzletron/orchestration/progress.py @@ -387,17 +387,6 @@ def _post_mip_progress( stage_id: str, config: Mapping[str, Any] | None, ) -> str | None: - """ - Summarize post-MIP candidate processing progress for a configured node. - - Parameters: - puzzle_dir (Path): Root directory containing post-MIP artifacts. - stage_id (str): Identifier of the post-MIP stage and node. - config (Mapping[str, Any] | None): Configuration containing post-MIP flow and node definitions. - - Returns: - str | None: Progress summary with completed, failed, and timed-out candidate counts, or None when progress data is unavailable or the node is not applicable. - """ parts = stage_id.split(".", 2) if len(parts) != 3: return None diff --git a/modelopt/torch/puzzletron/orchestration/task_launcher.py b/modelopt/torch/puzzletron/orchestration/task_launcher.py index 6428ada478e..42606b6180b 100644 --- a/modelopt/torch/puzzletron/orchestration/task_launcher.py +++ b/modelopt/torch/puzzletron/orchestration/task_launcher.py @@ -134,18 +134,7 @@ def build_task_command( binding: TaskBinding, gpus_per_task: int, ) -> tuple[str, ...]: - """ - Build the command used to launch the application for the selected launcher and task topology. - - Parameters: - payload (Sequence[str]): Application command and its arguments. - launcher (TaskLauncher): Launcher mode that determines whether to wrap the payload. - binding (TaskBinding): Resolved task placement and rendezvous information. - gpus_per_task (int): Number of processes to launch per node when using distributed execution. - - Returns: - tuple[str, ...]: The original payload for direct execution, or a torchrun command configured for the task topology. - """ + """Wrap an application payload in torchrun when the topology requests it.""" command = tuple(str(part) for part in payload) if launcher is TaskLauncher.DIRECT: @@ -185,20 +174,6 @@ def _direct_distributed_env(binding: TaskBinding) -> dict[str, str]: def _required_index(env: Mapping[str, str], primary: str, fallback: str) -> int: - """ - Read a task index from the primary environment variable or its fallback. - - Parameters: - env (Mapping[str, str]): Environment variables containing the task index. - primary (str): Preferred environment variable name. - fallback (str): Alternate environment variable name. - - Returns: - int: The task index parsed from the selected environment variable. - - Raises: - RuntimeError: If neither environment variable is set. - """ value = env.get(primary, env.get(fallback)) if value is None: raise RuntimeError(f"missing task identity: set {primary} or {fallback}") @@ -221,15 +196,7 @@ def _parser() -> argparse.ArgumentParser: def main(argv: Sequence[str] | None = None) -> int: - """ - Resolve the task's distributed binding, prepare its execution environment, and replace the current process with the payload command. - - Parameters: - argv (Sequence[str] | None): Optional command-line arguments to parse instead of the process arguments. - - Returns: - int: Zero after replacing the current process with the payload command. - """ + """Resolve this task's binding and replace the launcher with its payload.""" args = _parser().parse_args(argv) payload = tuple(args.payload[1:] if args.payload[:1] == ["--"] else args.payload) diff --git a/modelopt/torch/puzzletron/post_mip/builtin.py b/modelopt/torch/puzzletron/post_mip/builtin.py index 52ab60509d5..050352af445 100644 --- a/modelopt/torch/puzzletron/post_mip/builtin.py +++ b/modelopt/torch/puzzletron/post_mip/builtin.py @@ -125,5 +125,4 @@ class DownstreamEvaluationNode(PostMIPNode): @classmethod def render_report(cls, node, payload): - """Render the downstream evaluation report for the payload's section.""" return render_downstream_evaluation_report(str(payload["section_id"]), payload) diff --git a/modelopt/torch/puzzletron/post_mip/reporting.py b/modelopt/torch/puzzletron/post_mip/reporting.py index c21d61ce4a9..bc287990805 100644 --- a/modelopt/torch/puzzletron/post_mip/reporting.py +++ b/modelopt/torch/puzzletron/post_mip/reporting.py @@ -319,17 +319,7 @@ def render_evaluation_report(section_id: str, payload: Mapping[str, Any]) -> str def render_aiperf_report(section_id: str, payload: Mapping[str, Any]) -> str: - """ - Render AIPerf candidate status, performance metrics, selection markers, and errors. - - Parameters: - section_id (str): Identifier used to scope the throughput chart element. - payload (Mapping[str, Any]): AIPerf observations and status data. - - Returns: - str: HTML fragment containing the status summary, throughput chart placeholder, - and candidate metrics table. - """ + """Render AIPerf throughput/latency observations and timeout evidence.""" observations = list(payload.get("observations") or ()) rows = [] @@ -378,16 +368,7 @@ def render_downstream_evaluation_report(section_id: str, payload: Mapping[str, A def render_global_kd_report(section_id: str, payload: Mapping[str, Any]) -> str: - """ - Render the Short KD comparison with candidate statuses, loss plots, and run summaries. - - Parameters: - section_id (str): Identifier used to generate unique plot element IDs. - payload (Mapping[str, Any]): Short KD runs and status data to display. - - Returns: - str: HTML fragment containing the comparison summary, plot placeholders, and run table. - """ + """Render several candidate KD histories on shared, lineage-colored plots.""" runs = list(payload.get("runs") or ()) rows = [] diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 0511091b87e..1bbe463d0b7 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -45,8 +45,6 @@ "run_post_mip_node_shard", ] -_DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS = 3600.0 - def _puzzle_dir(config: Mapping[str, Any]) -> Path: return Path(config.get("puzzle_dir") or (config.get("experiment") or {})["dir"]) @@ -472,17 +470,6 @@ def _evaluate( def _aiperf( config: dict[str, Any], node: CompiledPostMIPNode, source, execution_identity: str ) -> dict[str, Any]: - """Run an AI performance sweep for a checkpoint across configured concurrency levels. - - Parameters: - config (dict[str, Any]): Workflow configuration used to determine the execution directory. - node (CompiledPostMIPNode): Compiled post-MIP node containing benchmark settings. - source: Candidate source containing the checkpoint and architecture identifier. - execution_identity (str): Identifier for the current node execution. - - Returns: - dict[str, Any]: Benchmark metrics and paths to the raw result artifacts. - """ from ..benchmarks import run_aiperf_sweep settings = dict(node.config.get("config") or {}) @@ -554,27 +541,10 @@ def _aiperf( def _as_cli_bool(value: bool) -> str: - """Convert a Boolean value to the CLI-compatible ``"True"`` or ``"False"`` string. - - Parameters: - value (bool): The Boolean value to convert. - - Returns: - str: ``"True"`` for true values and ``"False"`` for false values. - """ return "True" if value else "False" def _as_lmms_eval_arg(value: Any) -> str: - """ - Convert a value to the command-line argument format expected by lmms-eval. - - Parameters: - value (Any): The value to convert. - - Returns: - str: The formatted command-line argument value. - """ if isinstance(value, bool): return _as_cli_bool(value) if isinstance(value, (int, float)) and not isinstance(value, bool): @@ -585,20 +555,6 @@ def _as_lmms_eval_arg(value: Any) -> str: def _join_cli_values(value: Any, *, path: str) -> str: - """ - Convert a string or sequence of values into a comma-separated CLI value. - - Parameters: - value (Any): String or sequence of values to normalize. - path (str): Configuration path used in validation errors. - - Returns: - str: The normalized comma-separated value. - - Raises: - TypeError: If value is neither a string nor a sequence. - ValueError: If value is empty or contains an empty item. - """ if isinstance(value, str): text = value.strip() if not text: @@ -613,18 +569,6 @@ def _join_cli_values(value: Any, *, path: str) -> str: def _model_arg_string(values: Mapping[str, Any]) -> str: - """ - Convert model arguments to lmms-eval's comma-separated argument format. - - Parameters: - values (Mapping[str, Any]): Model argument names and values. - - Returns: - str: A comma-separated string of rendered key-value arguments. - - Raises: - ValueError: If an argument key or value is invalid, or if no arguments are provided. - """ parts = [] for key, value in values.items(): if value is None: @@ -645,16 +589,6 @@ def _model_arg_string(values: Mapping[str, Any]) -> str: def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> str: - """ - Merge checkpoint, topology, and supported model settings into lmms-eval model arguments. - - Parameters: - settings (Mapping[str, Any]): Downstream evaluation settings containing optional model arguments and configuration overrides. - checkpoint (str): Path to the checkpoint used for evaluation. - - Returns: - str: Comma-separated lmms-eval model arguments. - """ raw = settings.get("model_args") checkpoint_arg = str(settings.get("checkpoint_arg", "model")) topology = dict(settings.get("topology") or {}) @@ -674,34 +608,23 @@ def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> ], } ) - for key in sorted(_LMMS_EVAL_MODEL_ARG_FIELDS): + for key in _LMMS_EVAL_MODEL_ARG_FIELDS: if key in settings: derived[key] = settings[key] if isinstance(raw, str): prefix = raw.strip().strip(",") suffix = _model_arg_string(derived) - return ",".join(part for part in (suffix, prefix) if part) + return ",".join(part for part in (prefix, suffix) if part) if raw is not None and not isinstance(raw, Mapping): raise TypeError("downstream_evaluation.config.model_args must be a mapping or string") merged = dict(raw or {}) - merged.update(derived) + for key, value in derived.items(): + merged.setdefault(key, value) return _model_arg_string(merged) def _command_prefix(settings: Mapping[str, Any]) -> list[str]: - """ - Resolve the command prefix used to invoke lmms-eval. - - Parameters: - settings (Mapping[str, Any]): Downstream evaluation settings containing an optional command prefix. - - Returns: - list[str]: The configured command prefix, or the current Python interpreter followed by the lmms-eval module. - - Raises: - ValueError: If the configured command prefix is empty or contains an empty value. - """ raw = settings.get("command_prefix") if raw is None: return [sys.executable, "-m", "lmms_eval"] @@ -720,17 +643,7 @@ def _lmms_eval_command( checkpoint: str, output_path: Path, ) -> tuple[list[str], dict[str, str], float | None]: - """ - Builds an lmms-eval command, environment, and optional timeout for a realized checkpoint. - - Parameters: - settings (Mapping[str, Any]): Downstream evaluation settings. - checkpoint (str): Path to the realized checkpoint. - output_path (Path): Directory for lmms-eval output. - - Returns: - tuple[list[str], dict[str, str], float | None]: The command arguments, environment variables, and timeout in seconds. - """ + """Build a deterministic lmms-eval CLI invocation for one realized checkpoint.""" tasks = _join_cli_values(settings.get("tasks"), path="downstream_evaluation.config.tasks") argv = [ @@ -780,19 +693,10 @@ def _lmms_eval_command( if settings.get("cache_dir") is not None: env.setdefault("LMMS_EVAL_HOME", str(settings["cache_dir"])) timeout = settings.get("timeout_seconds", settings.get("timeout")) - return argv, env, (float(timeout) if timeout is not None else _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS) + return argv, env, (float(timeout) if timeout is not None else None) def _metric_key(value: Any) -> str: - """ - Normalize a metric name component for use in metric keys. - - Parameters: - value (Any): The value to convert into a normalized metric name component. - - Returns: - str: The stripped string representation with spaces, commas, and slashes replaced by underscores. - """ return ( str(value) .strip() @@ -804,15 +708,6 @@ def _metric_key(value: Any) -> str: def _flatten_lmms_eval_metrics(payload: Mapping[str, Any]) -> dict[str, float]: - """ - Flatten finite numeric task metrics from an lmms-eval result payload. - - Parameters: - payload (Mapping[str, Any]): Result payload containing task metrics under the ``results`` key. - - Returns: - dict[str, float]: Metric names mapped to finite numeric values, or an empty dictionary when no valid results are present. - """ results = payload.get("results") if not isinstance(results, Mapping): return {} @@ -831,18 +726,6 @@ def _flatten_lmms_eval_metrics(payload: Mapping[str, Any]) -> dict[str, float]: def _lmms_eval_result_payload(output_path: Path) -> tuple[dict[str, Any], Path]: - """ - Finds the newest valid lmms-eval result payload under an output directory. - - Parameters: - output_path (Path): Directory containing lmms-eval output files. - - Returns: - tuple[dict[str, Any], Path]: The result payload and path of the newest JSON file containing a `results` mapping. - - Raises: - FileNotFoundError: If no valid result JSON file is found. - """ candidates = [] for path in sorted(output_path.rglob("*.json")): try: @@ -860,16 +743,6 @@ def _lmms_eval_result_payload(output_path: Path) -> tuple[dict[str, Any], Path]: def _write_lmms_eval_streams( output_path: Path, result: subprocess.CompletedProcess[str] ) -> dict[str, str]: - """ - Persist non-empty lmms-eval subprocess output streams and return their artifact paths. - - Parameters: - output_path (Path): Directory where stream files are written. - result (subprocess.CompletedProcess[str]): Completed subprocess result containing captured output. - - Returns: - dict[str, str]: Mapping of stream path keys to the paths of written output files. - """ stream_paths = {} for stream_name, text in (("stdout", result.stdout), ("stderr", result.stderr)): if not text: @@ -881,16 +754,6 @@ def _write_lmms_eval_streams( def _lmms_eval_output_tail(result: subprocess.CompletedProcess[str], *, max_lines: int = 20) -> str: - """ - Format the most recent subprocess output lines from stderr and stdout. - - Parameters: - result (subprocess.CompletedProcess[str]): Completed process containing captured output. - max_lines (int): Maximum number of lines to include from each stream. - - Returns: - str: Formatted stderr and stdout output tails. - """ sections = [] for stream_name, text in (("stderr", result.stderr), ("stdout", result.stdout)): lines = (text or "").strip().splitlines() @@ -906,23 +769,6 @@ def _downstream_evaluation( source, execution_identity: str, ) -> dict[str, Any]: - """ - Run downstream lmms-eval benchmarking for a materialized checkpoint. - - Parameters: - config (dict[str, Any]): Campaign configuration used to determine execution paths. - node (CompiledPostMIPNode): Post-MIP node containing lmms-eval settings. - source: Checkpoint artifact to evaluate. - execution_identity (str): Identity of the current node execution. - - Returns: - dict[str, Any]: Paths to the evaluation summary, raw result, command record, and captured streams, together with numeric metrics. - - Raises: - ValueError: If the source is not a checkpoint artifact. - RuntimeError: If lmms-eval fails or produces no numeric task metrics. - FileNotFoundError: If no valid lmms-eval result file is produced. - """ if source.artifact_kind is not ArtifactKind.CHECKPOINT: raise ValueError("downstream_evaluation requires materialized checkpoint artifacts") settings = dict(node.config.get("config") or {}) @@ -950,25 +796,15 @@ def _downstream_evaluation( ) # Campaign config controls the executable and arguments, but subprocess receives # an argv list directly; no shell parsing is involved. - try: - result = subprocess.run( - argv, - cwd=str(output), - env=env, - capture_output=True, - text=True, - timeout=timeout, - check=False, - ) - except subprocess.TimeoutExpired as timeout_error: - timeout_result = subprocess.CompletedProcess( - args=argv, - returncode=-1, - stdout=timeout_error.stdout.decode("utf-8", errors="replace") if timeout_error.stdout else "", - stderr=timeout_error.stderr.decode("utf-8", errors="replace") if timeout_error.stderr else "", - ) - _write_lmms_eval_streams(output, timeout_result) - raise + result = subprocess.run( + argv, + cwd=str(output), + env=env, + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) stream_paths = _write_lmms_eval_streams(output, result) if result.returncode: tail = _lmms_eval_output_tail(result) @@ -1055,22 +891,6 @@ def _run_candidate( input_revision_id: str, execution_identity: str, ) -> dict[str, Any]: - """ - Execute a candidate according to the node type and return its execution result. - - Parameters: - config (dict[str, Any]): Runtime configuration for the candidate execution. - node (CompiledPostMIPNode): Compiled node defining the execution type and model source. - ledger (CandidateLedger): Ledger containing the input candidate revision. - input_revision_id (str): Identifier of the candidate revision to execute. - execution_identity (str): Identifier for the current node execution. - - Returns: - dict[str, Any]: A successful result containing the input and source revision identifiers, architecture identifier, and executor-specific metadata. - - Raises: - ValueError: If the node type is not a supported candidate executor. - """ source = ledger.source_revision(input_revision_id, node.model_source) if node.node_type == "materialize": result = _materialize(config, node, ledger, input_revision_id, source, execution_identity) @@ -1114,18 +934,6 @@ def _distributed_shard(config: dict[str, Any], node: CompiledPostMIPNode) -> Ite def run_post_mip_node_shard( config: dict[str, Any], stage_id: str, *, shard_index: int = 0, shard_count: int = 1 ) -> Path: - """ - Execute the assigned candidate revisions for a post-MIP node shard and persist the results. - - Parameters: - config (dict[str, Any]): Post-MIP configuration. - stage_id (str): Identifier of the compiled node to execute. - shard_index (int): Zero-based index of this shard. - shard_count (int): Total number of shards distributing the candidate revisions. - - Returns: - Path: Path to the shard result artifact. - """ node = _compiled_node(config, stage_id) ledger = _ledger(config) ledger.ingest_mip(_puzzle_dir(config)) @@ -1176,7 +984,7 @@ def run_post_mip_node_shard( timeout_field = "timeout_seconds" elif not isinstance(error, subprocess.TimeoutExpired): timeout_field = "readiness_timeout" - default_timeout = _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS if node.node_type == "downstream_evaluation" else ( + default_timeout = 3600 if node.node_type == "downstream_evaluation" else ( 600 if timeout_field == "benchmark_timeout" else 1200 ) row["timeout_seconds"] = float( diff --git a/modelopt/torch/puzzletron/stages/pipeline.py b/modelopt/torch/puzzletron/stages/pipeline.py index 8a40ad5e52c..26e0c059416 100644 --- a/modelopt/torch/puzzletron/stages/pipeline.py +++ b/modelopt/torch/puzzletron/stages/pipeline.py @@ -207,13 +207,7 @@ def _vllm_stats_is_explicit(config: dict[str, Any]) -> bool: def _calculate_static_workload_stats(config: dict[str, Any], hydra_cfg: Any) -> None: - """ - Append an analytical memory profile for each configured MIP workload. - - Parameters: - config (dict[str, Any]): Pipeline configuration containing optional MIP workloads. - hydra_cfg (Any): Base subblock-statistics configuration to customize for each workload. - """ + """Append one analytical memory profile for every configured MIP workload.""" from ..subblock_stats.calc_subblock_stats import launch_calc_subblock_stats workloads = dict((config.get("mip") or {}).get("workloads") or {}) @@ -249,15 +243,6 @@ def _calculate_static_workload_stats(config: dict[str, Any], hydra_cfg: Any) -> def _scenario_hidden_width(puzzle_dir: Path) -> int | None: - """ - Read the hidden width from a scenario manifest. - - Parameters: - puzzle_dir (Path): Directory containing the scenario manifest. - - Returns: - int | None: The manifest's hidden width, or `None` when the manifest or value is absent. - """ manifest_path = puzzle_dir / "scenario_manifest.json" if not manifest_path.is_file(): return None @@ -273,20 +258,6 @@ def _has_runtime_measurement( measurement: Any, allow_missing_workload_id: bool = False, ) -> bool: - """ - Determine whether a statistics file contains a compatible runtime measurement. - - Parameters: - path (Path): Statistics file to inspect. - hidden_width (int): Model hidden width expected by the measurement. - measurement (Any): Runtime measurement configuration to match. - allow_missing_workload_id (bool): Whether entries without a workload identifier may match. - - Returns: - bool: `True` if a compatible runtime measurement is present, `False` otherwise. - """ - from ..subblock_stats.calc_subblock_stats import _runtime_reuse_key, _runtime_stats_identity - try: payload = json.loads(path.read_text()) except (OSError, ValueError): @@ -294,27 +265,34 @@ def _has_runtime_measurement( if not isinstance(payload, list): return False expected_backend = (measurement.runtime_stats or {}).get("backend") - requested_key = _runtime_reuse_key( - width=hidden_width, - batch_size=measurement.batch_size, - prefill_seq_len=measurement.prefill_seq_len, - generation_seq_len=measurement.generation_seq_len, - runtime_stats_config=measurement.runtime_stats or {}, - ) for entry in payload: if not isinstance(entry, dict): continue args = entry.get("args") or {} if not isinstance(args, dict) or args.get("runtime_stats") is not True: continue - persisted_key = _runtime_stats_identity( - args, - fallback_workload_id=measurement.measurement_id if allow_missing_workload_id else None, - ) - if persisted_key is None: + if int(args.get("n_embd", -1)) != int(hidden_width): continue - if persisted_key == requested_key: - return True + if args.get("weights_dtype") != "torch.bfloat16": + continue + if int(args.get("batch_size", -1)) != int(measurement.batch_size): + continue + if int(args.get("prefill_seq_len", -1)) != int(measurement.prefill_seq_len): + continue + if int(args.get("generation_seq_len", -1)) != int(measurement.generation_seq_len): + continue + if int(args.get("max_num_seqs", -1)) != int(measurement.max_num_seqs): + continue + if args.get("runtime_granularity", "subblock") != measurement.granularity: + continue + if expected_backend is not None and args.get("runtime_backend") != expected_backend: + continue + workload_id = args.get("workload_id") + if workload_id is None and not allow_missing_workload_id: + continue + if workload_id is not None and workload_id != measurement.measurement_id: + continue + return True return False @@ -325,18 +303,6 @@ def _runtime_measurement_candidate_paths( stats_path: Path, measurement: Any, ) -> list[tuple[Path, bool]]: - """ - Builds candidate paths for locating reusable runtime measurement statistics. - - Parameters: - config (dict[str, Any]): Configuration containing the statistics filename. - puzzle_dir (Path): Directory associated with the current scenario. - stats_path (Path): Primary statistics file path. - measurement (Any): Measurement configuration that may specify a relative statistics path. - - Returns: - list[tuple[Path, bool]]: Candidate statistics paths paired with a flag indicating whether each path came from a configured relative path. - """ stats_name = str( (config.get("vllm_stats") or {}).get("subblock_stats_filename", stats_path.name) ) @@ -363,22 +329,6 @@ def _runtime_reuse_source_path( hidden_width: int, measurement: Any, ) -> Path: - """ - Finds a reusable vLLM measurement file matching the requested hidden width and workload. - - Parameters: - config (dict[str, Any]): Runtime configuration used to resolve candidate measurement paths. - puzzle_dir (Path): Experiment directory containing scenario-specific measurement files. - stats_path (Path): Configured statistics file path. - hidden_width (int): Hidden width required for the reusable measurement. - measurement (Any): Workload measurement whose identity must match. - - Returns: - Path: The first candidate measurement file containing a compatible runtime measurement. - - Raises: - RuntimeError: If no candidate contains a matching reusable measurement. - """ candidates = _runtime_measurement_candidate_paths( config=config, puzzle_dir=puzzle_dir, @@ -405,7 +355,7 @@ def _refresh_scenario_runtime_workload_stats( hydra_cfg: Any, stats_path: Path, ) -> None: - """Refresh scenario-specific runtime statistics using measurements for the local hidden width.""" + """Refresh width-scenario runtime rows with the local parameter inventory identity.""" from ..subblock_stats.calc_subblock_stats import launch_calc_subblock_stats puzzle_dir = _puzzle_dir(config, hydra_cfg) @@ -445,13 +395,7 @@ def _refresh_scenario_runtime_workload_stats( def _write_runtime_subblock_library(path: Path, block_configs: tuple[Any, ...]) -> None: - """ - Write runtime subblock configurations to a JSON library file. - - Parameters: - path (Path): Destination path for the library file. - block_configs (tuple[Any, ...]): Block configurations to serialize. - """ + """Write the legacy subblock-library input without assembling a replacement library.""" rows = [] for block_config in block_configs: row = { @@ -1180,21 +1124,6 @@ def bypass_overfit_stage(config: dict[str, Any], manifest: StageManifest): def build_library_stage(config: dict[str, Any], manifest: StageManifest): - """ - Build the replacement and candidate libraries and record their associated statistics. - - The stage validates and shares the resolved scoring parent, optionally refreshes runtime - statistics, calculates static workload statistics, and publishes the resulting artifact - paths and execution metadata. - - Parameters: - config (dict[str, Any]): Pipeline configuration. - manifest (StageManifest): Manifest used to record stage completion and outputs. - - Returns: - StageManifest: Updated manifest containing the generated library paths, statistics - metadata, and scoring-parent information. - """ hydra_cfg = load_runtime_hydra_config(config) puzzle_dir = _puzzle_dir(config, hydra_cfg) candidate_library_path = puzzle_dir / "candidate_library.json" diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py index 127f4b2d8c4..38580f156b7 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py @@ -276,36 +276,24 @@ def _runtime_reuse_key_from_args( *, fallback_workload_id: str | None = None, ) -> tuple | None: - """ - Builds the identity used to match reusable runtime measurements. - - Parameters: - args (Mapping): Persisted calculation arguments containing runtime and workload settings. - fallback_workload_id (str | None): Workload identifier to use when `args` does not provide one. - - Returns: - tuple | None: Runtime-reuse identity, or `None` when runtime statistics are unavailable or the dtype is not bfloat16. - """ + """Return the exact runtime-reuse identity represented by persisted args.""" if not args.get("runtime_stats") or args.get("weights_dtype") != str(torch.bfloat16): return None - required_fields = ["n_embd", "batch_size", "prefill_seq_len", "generation_seq_len"] - if any(args.get(field) is None for field in required_fields): - return None workload_id = args.get("workload_id") if workload_id is None: workload_id = fallback_workload_id return ( int(args["n_embd"]), int(args["batch_size"]), - int(args["prefill_seq_len"]), - int(args["generation_seq_len"]), + int(args.get("prefill_seq_len")), + int(args.get("generation_seq_len")), args.get("max_num_seqs"), args.get("runtime_granularity", "subblock"), args.get("runtime_backend"), - args.get("num_iters", 30), - args.get("num_warmup_iters", 10), - max(2, int(args.get("repeat_block_n_times", 10))), + args.get("num_iters"), + args.get("num_warmup_iters"), + args.get("repeat_block_n_times"), _freeze_stats_args(args.get("vllm_args")), workload_id, ) @@ -319,15 +307,7 @@ def _runtime_reuse_key( generation_seq_len: int, runtime_stats_config: Mapping, ) -> tuple: - """ - Builds an exact identity for matching reusable runtime measurements. - - Parameters: - runtime_stats_config (Mapping): Runtime settings that determine measurement compatibility. - - Returns: - tuple: Identity containing model dimensions, runtime settings, vLLM arguments, and workload identity. - """ + """Return the exact runtime-reuse identity requested by this calculation.""" return ( int(width), @@ -352,21 +332,7 @@ def _reuse_runtime_stats( source_path: str, fallback_workload_id: str | None = None, ) -> dict: - """ - Reuse measured runtime statistics from a compatible source entry in refreshed statistics. - - Parameters: - target (dict): Statistics entry to update with reusable runtime data. - source (dict): Statistics entry containing the measured runtime data. - source_path (str): Path identifying the source statistics entry. - fallback_workload_id (str | None): Workload identifier to use when the source does not provide one. - - Returns: - dict: The updated target statistics entry. - - Raises: - KeyError: If a target subblock has no matching runtime statistics in the source entry. - """ + """Overlay immutable measured latency onto refreshed static statistics.""" # Synthetic vLLM timings are layer-independent: collection benchmarks the # set of unique subblock configs, while a post-scoring replacement library @@ -860,33 +826,6 @@ def calculate_subblock_stats( runtime_selection_identity: str | None = None, parameter_inventory: Mapping | None = None, ) -> dict: - """ - Compute parameter, memory, additive-metric, and optional runtime statistics for subblock configurations. - - Parameters: - calc_subblock_stats_config (DictConfig): Runtime measurement and calculation settings. - teacher_dir (Path): Directory containing the teacher model or checkpoint. - model_config (PretrainedConfig): Model configuration used for metric calculations. - descriptor (Type[ModelDescriptor]): Model descriptor defining architecture-specific behavior. - master_puzzle_dir (Path): Puzzle directory used for runtime measurement caches. - subblock_configs (list[immutabledict[str, SubblockConfig]]): Subblock configurations and their parent layer indices. - batch_size (int): Number of sequences in the workload. - prefill_seq_len (int): Input sequence length used for prefill calculations. - generation_seq_len (int): Number of generated tokens used for decode calculations. - n_embd (int): Model hidden size. - n_head (int): Number of attention heads. - vocab_size (int): Model vocabulary size. - runtime_stats_enabled (bool): Whether to measure runtime statistics. - use_cuda_graph (bool): Whether to use CUDA graphs during runtime measurement. - weights_dtype (torch.dtype): Data type used for model weights. - activations_dtype (torch.dtype): Data type used for activations. - kv_cache_dtype (torch.dtype): Data type used for the key-value cache. - runtime_selection_identity (str | None): Identity of the runtime subblock selection. - parameter_inventory (Mapping | None): Precomputed parameter inventory to use for parameter counts. - - Returns: - dict: Statistics for the requested workload, including calculation arguments, non-block statistics, and per-subblock metrics. - """ runtime_granularity = "subblock" runtime_stats_config = ( calc_subblock_stats_config.get("runtime_stats", {}) if runtime_stats_enabled else {} @@ -1307,26 +1246,12 @@ def _subblock_stats_already_complete( prefill_seq_len: int = 2048, generation_seq_len: int = 2048, ) -> bool: - """ - Determine whether existing statistics cover all requested configurations. - - Parameters: - existing_stats (list): Previously calculated statistics entries. - subblock_configs (list): Subblock configurations required for coverage. - batch_sizes (Iterable[int]): Batch sizes to verify. - data_types (list): Weight, activation, and KV-cache dtype combinations. - model_hidden_sizes (Iterable[int]): Model widths to verify. - runtime_stats_enabled (bool): Whether runtime measurements are required. - runtime_granularity (str): Required runtime measurement granularity. - runtime_max_num_seqs (int | None): Required maximum number of runtime sequences. - runtime_workload_id (str | None): Required runtime workload identity. - runtime_selection_identity (str | None): Required runtime subblock-selection identity. - parameter_inventory_identities (Mapping[int, str] | None): Inventory identity required for each model width. - prefill_seq_len (int): Required prefill sequence length. - generation_seq_len (int): Required generation sequence length. - - Returns: - bool: True if every requested configuration and required measurement is present, False otherwise. + """Whether ``existing_stats`` already covers every configuration this run would compute. + + When runtime benchmarking is enabled, the bf16 entries (the only ones for + which runtime is ever measured) must additionally already carry runtime + measurements **at the requested granularity** — switching subblock<->block must trigger a + recompute rather than silently reusing the other granularity's numbers. """ by_signature = {_arg_signature(entry["args"]): entry for entry in existing_stats} required_subblock_keys = { @@ -1429,27 +1354,6 @@ def calculate_subblock_stats_for_puzzle_dir( # from attach_helper import debugging_setup # debugging_setup() # You can optionally pass a name to identify the job (e.g. `debugging_setup(name="my_script")`) # ==== END === Setup for attach-helper ==== - """ - Compute and persist subblock statistics for all requested batch sizes, data types, and model widths. - - Parameters: - calc_subblock_stats_config (DictConfig): Configuration for statistics calculation and optional runtime measurement. - master_puzzle_dir (Path | str): Puzzle directory containing subblock configurations and output files. - teacher_dir (Path | str): Teacher checkpoint directory used for model metadata and parameter inventories. - descriptor (Type[ModelDescriptor]): Model descriptor defining architecture-specific behavior. - model_hidden_sizes (ListConfig): Hidden sizes to evaluate; the teacher hidden size is always included. - ffn_hidden_sizes (ListConfig): Additional FFN sizes to include in the subblock configurations. - batch_sizes (Iterable[int]): Batch sizes to evaluate. - prefill_seq_len (int): Number of prompt tokens used for runtime measurements. - generation_seq_len (int): Number of generated tokens used for runtime measurements. - runtime_stats_enabled (bool): Whether to compute or reuse runtime statistics. - merge_with_existing_stats (bool): Whether to update an existing incomplete statistics file. - subblock_stats_filename (str): Name of the JSON file used to persist statistics. - - Raises: - FileNotFoundError: If a configured runtime manifest or reusable runtime statistics file cannot be found. - ValueError: If runtime settings or reusable runtime statistics do not cover the requested configurations. - """ if isinstance(batch_sizes, str): batch_sizes = [ int(batch_size) for batch_size in batch_sizes.strip("[]").replace(" ", "").split(",") diff --git a/puzzletron_setup/bundle.py b/puzzletron_setup/bundle.py index 13bad12f5ef..3ee3e8f0936 100644 --- a/puzzletron_setup/bundle.py +++ b/puzzletron_setup/bundle.py @@ -329,19 +329,6 @@ def _post_mip_flows( global_kd_mesh: Mapping[str, Any], default_serving_topology: Mapping[str, Any], ) -> dict[str, Any]: - """ - Prepare post-MIP flow configurations with mesh settings, serving defaults, and smoke-run limits. - - Parameters: - state (Mapping[str, Any]): Campaign state containing post-MIP flow definitions. - smoke (bool): Whether to apply reduced settings for a smoke run. - common_mesh (Mapping[str, Any]): Mesh used by evaluation and materialization nodes. - global_kd_mesh (Mapping[str, Any]): Mesh used by global knowledge-distillation nodes. - default_serving_topology (Mapping[str, Any]): Default topology for serving-based nodes. - - Returns: - dict[str, Any]: The normalized post-MIP flow configurations. - """ flows = deepcopy(_mapping(_answers(state, "post_mip").get("flows"))) for flow in flows.values(): for node in _mapping(flow.get("nodes")).values(): @@ -808,21 +795,6 @@ def _dynamic_stage_entries( *, pool_source_evaluations: bool, ) -> dict[str, Any]: - """ - Builds scheduler entries for dynamic post-MIP stages. - - Parameters: - experiment (Mapping[str, Any]): Experiment configuration containing post-MIP flows. - workers (Mapping[str, Any]): Worker limits for pooled and sharded stages. - gpus_per_node (int): Number of GPUs assigned to each node. - common (Mapping[str, Any]): Parallel configuration for evaluation stages. - single_gpu (Mapping[str, Any]): Parallel configuration for materialization stages. - cpu_partition (str | None): CPU partition to assign to CPU stages. - pool_source_evaluations (bool): Whether source evaluations should use pooled workers. - - Returns: - dict[str, Any]: Scheduler entries keyed by post-MIP flow and node identifiers. - """ entries = {} candidate_limits = _post_mip_candidate_limits(experiment) for flow_id, flow in _mapping(_mapping(experiment.get("post_mip")).get("flows")).items(): diff --git a/puzzletron_setup/v2/validation.py b/puzzletron_setup/v2/validation.py index 787c2989ec7..d0cf216d693 100644 --- a/puzzletron_setup/v2/validation.py +++ b/puzzletron_setup/v2/validation.py @@ -188,12 +188,7 @@ def _dataset_subset_issues(state: WizardState) -> list[ValidationIssue]: def validate_state(state: WizardState) -> tuple[ValidationIssue, ...]: - """ - Validate wizard state and identify issues that prevent canonical compilation. - - Returns: - tuple[ValidationIssue, ...]: Validation issues sorted by configuration section and path. - """ + """Return actionable authoring issues before canonical compilation.""" issues: list[ValidationIssue] = [] required = ( "model.source", diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index 52d752d0bb5..ded6bbbff4c 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -3767,17 +3767,6 @@ def _post_mip_strategy(node: NodeDraft) -> str: def post_mip_section(session: WizardSession, resolver: DefaultsResolver, context: dict) -> bool: - """ - Configure post-MIP execution flows for each MIP run, using recommended or custom nodes. - - Parameters: - session (WizardSession): Wizard session used to read state and collect configuration. - resolver (DefaultsResolver): Resolver for stage resource defaults. - context (dict): Model and pruning context required to configure serving and evaluation nodes. - - Returns: - bool: `True` when post-MIP flows are configured, `False` when the section is exited through back navigation. - """ mip = _mapping_copy(session.state.collection("mip_config")) runs = _mapping_copy(mip.get("runs")) sequence = int(session.state.get_field("data.sequence_length", 4096)) @@ -4204,22 +4193,7 @@ def _serving_setting_prompt( pruning: Mapping[str, Any], stage_id: str, ) -> Any: - """ - Collect AIPerf serving workload settings and the vLLM serving topology. - - Parameters: - session (WizardSession): Wizard session used to collect and validate responses. - prefix (str): State key prefix for the serving settings. - defaults (Mapping[str, Any]): Default workload and topology values. - inventory (Any): Model inventory used to validate the topology. - pruning (Mapping[str, Any]): Pruning configuration relevant to topology validation. - stage_id (str): Pruning stage associated with the serving configuration. - - Returns: - Any: A mapping containing input and output sequence lengths, concurrency values, - request count, model selection mode, and topology, or the `BACK` sentinel when - the user navigates to the previous prompt. - """ + """Ask the complete AIPerf workload and serving-only parallel setting.""" values = {} for name, label, default in ( ("input_tokens", "Serving input sequence length (ISL):", defaults["input_tokens"]), @@ -4296,30 +4270,9 @@ def _downstream_evaluation_setting_prompt( pruning: Mapping[str, Any], stage_id: str, ) -> Any: - """ - Collect lmms-eval tasks, execution settings, model arguments, and vLLM topology. - - Parameters: - session (WizardSession): Wizard session used to prompt for settings. - prefix (str): State-key prefix for the prompted values. - defaults (Mapping[str, Any]): Existing values used as prompt defaults. - inventory (Any): Model inventory used to validate the vLLM topology. - pruning (Mapping[str, Any]): Pruning configuration relevant to topology validation. - stage_id (str): Identifier of the stage using the evaluation settings. - - Returns: - Any: A mapping containing lmms-eval tasks, sample and batch limits, timeout, model arguments, logging settings, and vLLM topology, or `BACK` if prompting is cancelled. - """ + """Ask lmms-eval task settings and the vLLM topology used to run them.""" def validate_tasks(value: str) -> bool | str: - """Validate a comma-separated list of lmms-eval tasks. - - Parameters: - value (str): Comma-separated task names. - - Returns: - bool | str: `True` if at least one task is provided, otherwise an error message. - """ tasks = [item.strip() for item in value.split(",") if item.strip()] return True if tasks else "Enter at least one lmms-eval task." @@ -4393,16 +4346,7 @@ def _configure_dynamic_resources( *, ask: bool, ) -> Any: - """ - Configure independent resource assignments for all nodes in a post-MIP flow. - - Parameters: - flow_id (str): Identifier of the flow whose nodes are configured. - ask (bool): Whether to prompt for resource and batch customizations. - - Returns: - True when configuration completes, or `BACK` when navigation is requested. - """ + """Attach an independent resource/batch card to every node in one flow.""" registry = ResourceProfileRegistry.from_dict( session.state.collection("parallel_profiles") or {} ) diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index 4b092569c34..2750be52f6f 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -640,18 +640,7 @@ def _ask_aiperf_config( runtime: Mapping[str, Any], defaults: Mapping[str, Any] | None = None, ) -> dict[str, Any]: - """ - Configure an AIPerf serving node's parallel topology and workload settings. - - Parameters: - detailed (bool): Whether to prompt for workload and timeout values. - moe (bool): Whether to configure expert parallelism for a mixture-of-experts model. - runtime (Mapping[str, Any]): Runtime defaults for input length, output length, and concurrency. - defaults (Mapping[str, Any] | None): Previously saved configuration values. - - Returns: - dict[str, Any]: The configured AIPerf topology, workload, and timeout settings. - """ + """Ask for one AIPerf node's independent Serving topology and workload.""" defaults = dict(defaults or {}) topology_defaults = dict(defaults.get("topology") or {}) checkpoint = prompts.checkpoint() @@ -751,26 +740,12 @@ def _ask_downstream_evaluation_config( moe: bool, defaults: Mapping[str, Any] | None = None, ) -> dict[str, Any]: - """ - Collect lmms-eval tasks, sampling settings, vLLM topology, and evaluation timeout. - - Parameters: - detailed (bool): Whether to prompt for the per-candidate timeout. - moe (bool): Whether to allow configuring expert parallelism. - defaults (Mapping[str, Any] | None): Previously saved settings used as prompt defaults. - - Returns: - dict[str, Any]: The configured downstream evaluation settings. - """ + """Ask for lmms-eval task and vLLM settings.""" defaults = defaults or {} - default_tasks = defaults.get("tasks", "ifeval,gsm8k") - if isinstance(default_tasks, list): - default_tasks = ",".join(default_tasks) tasks = prompts.text( "lmms-eval tasks (comma-separated):", - default=str(default_tasks), - validate=lambda value: bool(str(value).strip()) or "Enter at least one task.", + default=str(defaults.get("tasks", "ifeval,gsm8k")), ) limit = prompts.integer( "lmms-eval sample limit:", @@ -876,21 +851,6 @@ def _default_flow( objective: Mapping[str, Any] | None = None, include_initial_filter: bool = True, ) -> dict[str, Any]: - """ - Build the standard post-MIP evaluation and selection flow. - - Parameters: - run_id (str): Identifier of the MIP run. - run (Mapping[str, Any]): MIP run configuration. - runtime (Mapping[str, Any]): Runtime settings for serving evaluation. - data (Mapping[str, Any]): Dataset settings, including sequence length. - prefix (str): Prefix applied to generated node identifiers. - objective (Mapping[str, Any] | None): Objective used to configure ranking; the run's first objective is used when omitted. - include_initial_filter (bool): Whether to include the initial MIP-score filter. - - Returns: - dict[str, Any]: Flow configuration containing the source metadata and ordered post-MIP nodes. - """ def node_id(name: str) -> str: return f"{prefix}{name}" @@ -1050,20 +1010,6 @@ def _custom_flow( detailed: bool, moe: bool, ) -> dict[str, Any]: - """ - Build a custom post-MIP evaluation flow through interactive configuration. - - Parameters: - run_id (str): Identifier of the MIP run supplying candidate models. - runtime (Mapping[str, Any]): Runtime settings used by serving evaluations. - data (Mapping[str, Any]): Dataset settings used by evaluation nodes. - used_ids (set[str]): Node IDs already in use; newly configured IDs are added. - detailed (bool): Whether to collect detailed evaluation settings. - moe (bool): Whether to enable mixture-of-experts configuration options. - - Returns: - dict[str, Any]: A flow definition containing the MIP source and configured nodes. - """ nodes: OrderedDict[str, Any] = OrderedDict() available_metrics = ["mip.score"] transformer_nodes = [] @@ -1131,8 +1077,7 @@ def _custom_flow( detailed=detailed, moe=moe, ) - for task_name in node["config"].get("tasks", []): - available_metrics.append(f"{node_id}.{task_name}.strict-match") + available_metrics.append(f"{node_id}.gsm8k.exact_match") elif node_type == "global_kd": node["config"] = {"max_steps": prompts.integer("Global KD steps:", default=128)} elif node_type == "ptq": @@ -1281,20 +1226,6 @@ def _resource_rows( gpus_per_node: int, workers: Mapping[str, int], ) -> list[dict[str, Any]]: - """ - Calculate resource requirements for each campaign execution stage. - - Parameters: - state (AnswerState): Campaign configuration containing post-MIP flows and execution details. - common (Mapping[str, int]): Parallel mesh dimensions shared by common stages. - bypass (Mapping[str, int]): Parallel mesh dimensions for bypass processing. - global_kd (Mapping[str, int]): Parallel mesh dimensions for global knowledge distillation. - gpus_per_node (int): Number of GPUs available on each node. - workers (Mapping[str, int]): Worker limits for pool and sharded stages. - - Returns: - list[dict[str, Any]]: Resource rows containing each stage's name, instance count, GPUs per instance, and required node count. - """ from .bundle import _post_mip_candidate_limits, _serving_parallel rows = [] diff --git a/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py b/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py index 87f21abdbba..4e9a6e39c6d 100644 --- a/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py +++ b/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py @@ -1261,10 +1261,6 @@ def test_width_scenario_runtime_stats_reuse_root_measurement(tmp_path, monkeypat "generation_seq_len": 1024, "max_num_seqs": 1, "n_embd": 2688, - "num_iters": 30, - "num_warmup_iters": 10, - "repeat_block_n_times": 10, - "vllm_args": [], "workload_id": "serving-default", }, "subblocks": [], @@ -1412,6 +1408,7 @@ def test_runtime_stats_resume_signature_includes_workload_id(): **kwargs, runtime_workload_id="different-workload", ) + assert hydra_cfg.calc_subblock_stats.merge_with_existing_stats is False def test_sparse_runtime_selection_is_unique_and_layer_independent(): From 8ec08ef0f3caf36aca11e467b3e5e92912d6494d Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Fri, 7 Aug 2026 12:12:20 -0700 Subject: [PATCH 08/16] undid changes not related to evaluation Signed-off-by: Grzegorz Karch --- .../distributed_eval/run_coordinator.sh | 5 +- .../distributed_eval/run_depth_coordinator.sh | 5 +- .../distributed_eval/run_depth_pool.sh | 7 +- .../distributed_eval/run_replacement_pool.sh | 7 +- .../torch/puzzletron/benchmarks/aiperf.py | 44 +-- .../puzzletron/orchestration/adapters/pool.py | 4 +- .../orchestration/executors/slurm.py | 27 -- .../puzzletron/orchestration/task_launcher.py | 29 -- modelopt/torch/puzzletron/stages/pipeline.py | 159 +---------- .../subblock_stats/calc_subblock_stats.py | 127 ++------- .../test_aiperf_context_capacity.py | 44 --- .../test_orchestration_executors.py | 51 ---- .../test_orchestration_task_topology.py | 50 ---- .../puzzletron/test_sparse_runtime_stats.py | 266 ------------------ 14 files changed, 35 insertions(+), 790 deletions(-) diff --git a/examples/puzzletron/distributed_eval/run_coordinator.sh b/examples/puzzletron/distributed_eval/run_coordinator.sh index 62ea0fc9ed1..ac22b6774fd 100755 --- a/examples/puzzletron/distributed_eval/run_coordinator.sh +++ b/examples/puzzletron/distributed_eval/run_coordinator.sh @@ -3,8 +3,7 @@ set -Eeuo pipefail : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}" : "${CONFIG_PATH:?set CONFIG_PATH}" -: "${WORKER_WORLD_SIZE:=${WORLD_SIZE:-}}" -: "${WORKER_WORLD_SIZE:?set WORKER_WORLD_SIZE to one torchrun worker-group world size}" +: "${WORLD_SIZE:?set WORLD_SIZE to one torchrun worker-group world size}" : "${SOLUTIONS_PATH:?set SOLUTIONS_PATH}" : "${OUTPUT_DIR:?set OUTPUT_DIR}" @@ -43,7 +42,7 @@ if [[ ! -f "${CAMPAIGN_DIR}/manifest.json" ]]; then "${PYTHON_BIN}" -m modelopt.torch.puzzletron.distributed_eval.cli init \ --campaign-dir "${CAMPAIGN_DIR}" \ --config "${CONFIG_PATH}" \ - --world-size "${WORKER_WORLD_SIZE}" \ + --world-size "${WORLD_SIZE}" \ --evaluator-revision "${EVALUATOR_REVISION}" \ "${override_args[@]}" fi diff --git a/examples/puzzletron/distributed_eval/run_depth_coordinator.sh b/examples/puzzletron/distributed_eval/run_depth_coordinator.sh index dc37dee7b36..dbc7682f931 100755 --- a/examples/puzzletron/distributed_eval/run_depth_coordinator.sh +++ b/examples/puzzletron/distributed_eval/run_depth_coordinator.sh @@ -3,8 +3,7 @@ set -Eeuo pipefail : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}" : "${CONFIG_PATH:?set CONFIG_PATH}" -: "${WORKER_WORLD_SIZE:=${WORLD_SIZE:-}}" -: "${WORKER_WORLD_SIZE:?set WORKER_WORLD_SIZE to one torchrun worker-group world size}" +: "${WORLD_SIZE:?set WORLD_SIZE to one torchrun worker-group world size}" PYTHON_BIN="${PYTHON_BIN:-python}" @@ -19,7 +18,7 @@ if [[ ! -f "${CAMPAIGN_DIR}/manifest.json" ]]; then "${PYTHON_BIN}" -m modelopt.torch.puzzletron.distributed_eval.cli init \ --campaign-dir "${CAMPAIGN_DIR}" \ --config "${CONFIG_PATH}" \ - --world-size "${WORKER_WORLD_SIZE}" \ + --world-size "${WORLD_SIZE}" \ --stage depth \ --evaluator-revision "${EVALUATOR_REVISION:-puzzletron-depth-v1}" \ "${override_args[@]}" diff --git a/examples/puzzletron/distributed_eval/run_depth_pool.sh b/examples/puzzletron/distributed_eval/run_depth_pool.sh index c38044f87e3..3b8b663c9b4 100644 --- a/examples/puzzletron/distributed_eval/run_depth_pool.sh +++ b/examples/puzzletron/distributed_eval/run_depth_pool.sh @@ -6,8 +6,7 @@ set -Eeuo pipefail : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}" : "${CONFIG_PATH:?set CONFIG_PATH}" -: "${WORKER_WORLD_SIZE:=${WORLD_SIZE:-}}" -: "${WORKER_WORLD_SIZE:?set WORKER_WORLD_SIZE to one worker-group world size}" +: "${WORLD_SIZE:?set WORLD_SIZE to one worker-group world size}" : "${WORKER_COUNT:?set WORKER_COUNT to the number of worker groups}" : "${PUZZLETRON_GROUP_INDEX:=${PUZZLETRON_TASK_INDEX:-${SLURM_PROCID:-}}}" : "${PUZZLETRON_GROUP_INDEX:?run this script as one orchestrator worker-group task}" @@ -56,7 +55,7 @@ if [[ "${GROUP_INDEX}" == "0" && ! -f "${MANIFEST_PATH}" ]]; then -m modelopt.torch.puzzletron.distributed_eval.cli init \ --campaign-dir "${CAMPAIGN_DIR}" \ --config "${CONFIG_PATH}" \ - --world-size "${WORKER_WORLD_SIZE}" \ + --world-size "${WORLD_SIZE}" \ --stage depth \ --evaluator-revision "${EVALUATOR_REVISION:-puzzletron-depth-v1}" \ "${override_args[@]}" @@ -75,7 +74,7 @@ done # independent worker groups may share a node. export NNODES=1 export NODE_RANK=0 -export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORKER_WORLD_SIZE}}" +export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORLD_SIZE}}" export WORKER_GROUP_INDEX="${GROUP_INDEX}" export WORKER_ID="${WORKER_PREFIX}${GROUP_INDEX}" export WORKER_HOST="${WORKER_HOST:-$(hostname -f)}" diff --git a/examples/puzzletron/distributed_eval/run_replacement_pool.sh b/examples/puzzletron/distributed_eval/run_replacement_pool.sh index 46d5336a940..283a9ef11a3 100755 --- a/examples/puzzletron/distributed_eval/run_replacement_pool.sh +++ b/examples/puzzletron/distributed_eval/run_replacement_pool.sh @@ -6,8 +6,7 @@ set -Eeuo pipefail : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}" : "${CONFIG_PATH:?set CONFIG_PATH}" -: "${WORKER_WORLD_SIZE:=${WORLD_SIZE:-}}" -: "${WORKER_WORLD_SIZE:?set WORKER_WORLD_SIZE to one worker-group world size}" +: "${WORLD_SIZE:?set WORLD_SIZE to one worker-group world size}" : "${WORKER_COUNT:?set WORKER_COUNT to the number of worker groups}" : "${PUZZLETRON_GROUP_INDEX:=${PUZZLETRON_TASK_INDEX:-${SLURM_PROCID:-}}}" : "${PUZZLETRON_GROUP_INDEX:?run this script as one orchestrator worker-group task}" @@ -55,7 +54,7 @@ if [[ "${GROUP_INDEX}" == "0" && ! -f "${MANIFEST_PATH}" ]]; then -m modelopt.torch.puzzletron.distributed_eval.cli init \ --campaign-dir "${CAMPAIGN_DIR}" \ --config "${CONFIG_PATH}" \ - --world-size "${WORKER_WORLD_SIZE}" \ + --world-size "${WORLD_SIZE}" \ --stage replace_block \ --evaluator-revision "${EVALUATOR_REVISION:-puzzletron-distributed-replace-block-v1}" \ "${override_args[@]}" @@ -72,7 +71,7 @@ done export NNODES=1 export NODE_RANK=0 -export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORKER_WORLD_SIZE}}" +export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORLD_SIZE}}" export WORKER_GROUP_INDEX="${GROUP_INDEX}" export WORKER_ID="${WORKER_PREFIX}${GROUP_INDEX}" export WORKER_HOST="${WORKER_HOST:-$(hostname -f)}" diff --git a/modelopt/torch/puzzletron/benchmarks/aiperf.py b/modelopt/torch/puzzletron/benchmarks/aiperf.py index 5824247f8be..872822777bd 100644 --- a/modelopt/torch/puzzletron/benchmarks/aiperf.py +++ b/modelopt/torch/puzzletron/benchmarks/aiperf.py @@ -143,38 +143,6 @@ def _topology_vllm_args(topology: dict[str, Any]) -> list[str]: return args -def _extra_vllm_args(topology: dict[str, Any]) -> tuple[str, ...]: - """Return caller-provided vLLM args without treating a string as characters.""" - - raw = topology.get("extra_vllm_args", ()) - if raw is None: - return () - if isinstance(raw, str): - return (raw,) - return tuple(str(arg) for arg in raw) - - -def _has_vllm_option(args: Iterable[str], *options: str) -> bool: - """Return whether an option is present as ``--flag`` or ``--flag=value``.""" - - option_set = set(options) - return any(str(arg).split("=", 1)[0] in option_set for arg in args) - - -def _server_vllm_args( - checkpoint_dir: Path, topology: dict[str, Any], concurrency_values: Iterable[int] -) -> list[str]: - """Build stable vLLM server args derived from topology and benchmark demand.""" - - args = _topology_vllm_args(topology) - args.extend(_descriptor_vllm_args(checkpoint_dir)) - extra_args = _extra_vllm_args(topology) - if not _has_vllm_option(extra_args, "--max-num-seqs", "--max_num_seqs"): - args.extend(("--max-num-seqs", str(max(concurrency_values)))) - args.extend(extra_args) - return args - - def _exact_length_extra_inputs( extra_inputs: dict[str, Any] | None, output_tokens: int ) -> dict[str, Any]: @@ -437,8 +405,6 @@ def run_aiperf_sweep( model_name = f"puzzletron-{architecture_id[:16]}" tokenizer_dir = _short_tokenizer_alias(checkpoint_dir, artifact_dir) server_log = artifact_dir / "vllm_server.log" - server_max_model_len = _server_max_model_len(input_tokens, output_tokens, topology) - server_vllm_args = _server_vllm_args(checkpoint_dir, topology, concurrency_values) server_cmd = [ "vllm", "serve", @@ -450,10 +416,12 @@ def run_aiperf_sweep( "--served-model-name", model_name, "--max-model-len", - str(server_max_model_len), + str(_server_max_model_len(input_tokens, output_tokens, topology)), "--trust-remote-code", ] - server_cmd.extend(server_vllm_args) + server_cmd.extend(_topology_vllm_args(topology)) + server_cmd.extend(_descriptor_vllm_args(checkpoint_dir)) + server_cmd.extend(str(arg) for arg in topology.get("extra_vllm_args", ())) executable = _resolve_executable(executable) env = _clean_subprocess_environment( gpu_ids, @@ -495,10 +463,6 @@ def run_aiperf_sweep( "endpoint_type": endpoint_type, "extra_inputs": _exact_length_extra_inputs(extra_inputs, output_tokens), "use_server_token_count": use_server_token_count, - "server": { - "max_model_len": server_max_model_len, - "vllm_args": tuple(server_vllm_args), - }, "revisions": revisions, }, prefix="aiperf_result", diff --git a/modelopt/torch/puzzletron/orchestration/adapters/pool.py b/modelopt/torch/puzzletron/orchestration/adapters/pool.py index 7abee0a5263..73a489fe9d3 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/pool.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/pool.py @@ -242,8 +242,8 @@ def command( "CAMPAIGN_DIR": str(campaign_dir), "CONFIG_PATH": plan.experiment_config_path, "PUZZLE_DIR": str(replacement_puzzle_dir), + "WORLD_SIZE": str(node.gpus_per_instance), "NPROC_PER_NODE": str(node.gpus_per_instance), - "WORKER_WORLD_SIZE": str(node.gpus_per_instance), "WORKER_COUNT": str(worker_count), } if node.stage_id == "depth_importance": @@ -317,7 +317,7 @@ def command( "CAMPAIGN_DIR": str(campaign_dir), "CONFIG_PATH": plan.experiment_config_path, "PUZZLE_DIR": str(replacement_puzzle_dir), - "WORKER_WORLD_SIZE": str(node.gpus_per_instance), + "WORLD_SIZE": str(node.gpus_per_instance), "WORKER_ID": str(item.metadata.get("worker_id", 0)), "WORKER_COUNT": str(node.instances), } diff --git a/modelopt/torch/puzzletron/orchestration/executors/slurm.py b/modelopt/torch/puzzletron/orchestration/executors/slurm.py index 047c2a0ce9f..d33b865c327 100644 --- a/modelopt/torch/puzzletron/orchestration/executors/slurm.py +++ b/modelopt/torch/puzzletron/orchestration/executors/slurm.py @@ -90,31 +90,6 @@ def render_hook_lines(commands: Sequence[str]) -> str: return "\n".join(lines) -def _render_host_container_env(repository: str) -> str: - """Render host-side container runtime defaults for Pyxis/Enroot.""" - - cache_root = Path(repository) / ".cache" / "enroot" - lines = [ - 'if [[ -z "${ENROOT_CACHE_PATH:-}" ]]; then', - f" export ENROOT_CACHE_PATH={shlex.quote(str(cache_root / 'cache'))}", - "fi", - 'if [[ -z "${ENROOT_DATA_PATH:-}" ]]; then', - f" export ENROOT_DATA_PATH={shlex.quote(str(cache_root / 'data'))}", - "fi", - 'if [[ -z "${ENROOT_TEMP_PATH:-}" ]]; then', - ' export ENROOT_TEMP_PATH="/tmp/puzzletron-enroot-${USER:-$(id -u)}/tmp"', - "fi", - 'if [[ -z "${ENROOT_RUNTIME_PATH:-}" ]]; then', - ' export ENROOT_RUNTIME_PATH="/tmp/puzzletron-enroot-${USER:-$(id -u)}/runtime"', - "fi", - ( - 'mkdir -p "$ENROOT_CACHE_PATH" "$ENROOT_DATA_PATH" ' - '"$ENROOT_TEMP_PATH" "$ENROOT_RUNTIME_PATH"' - ), - ] - return "\n".join(lines) - - def render_sbatch_script( *, attempt: AttemptSpec, @@ -203,8 +178,6 @@ def render_sbatch_script( header_lines.append(f"#SBATCH --gpus-per-node={step_gpus_per_node}") header_lines.append(f"#SBATCH --output={log_path}") script = "\n".join(header_lines) + "\n" - if container_image: - script += _render_host_container_env(repository) + "\n" prologue_parts = [ "set -Eeuo pipefail", postrun_trap, diff --git a/modelopt/torch/puzzletron/orchestration/task_launcher.py b/modelopt/torch/puzzletron/orchestration/task_launcher.py index 42606b6180b..4d4b21c414a 100644 --- a/modelopt/torch/puzzletron/orchestration/task_launcher.py +++ b/modelopt/torch/puzzletron/orchestration/task_launcher.py @@ -29,17 +29,9 @@ TASK_IDENTITY_ENV_KEYS = frozenset( { "CUDA_VISIBLE_DEVICES", - "GROUP_RANK", - "GROUP_WORLD_SIZE", - "LOCAL_RANK", - "LOCAL_WORLD_SIZE", - "MASTER_ADDR", - "MASTER_PORT", "SLURM_LOCALID", "SLURM_NTASKS", "SLURM_PROCID", - "RANK", - "WORLD_SIZE", "PUZZLETRON_GROUP_INDEX", "PUZZLETRON_GROUP_RANK", "PUZZLETRON_GROUP_SIZE", @@ -154,25 +146,6 @@ def build_task_command( ) -def _direct_distributed_env(binding: TaskBinding) -> dict[str, str]: - """Return torch.distributed env for direct payloads that initialize env://.""" - - master_addr = "127.0.0.1" if binding.group_size == 1 else binding.master_addr - return { - "RANK": str(binding.group_rank), - "WORLD_SIZE": str(binding.group_size), - # The launcher slices CUDA_VISIBLE_DEVICES per task, so every direct - # payload has a local single-process view even when multiple tasks share - # a physical host. - "LOCAL_RANK": "0", - "LOCAL_WORLD_SIZE": "1", - "GROUP_RANK": str(binding.group_index), - "GROUP_WORLD_SIZE": str(binding.group_size), - "MASTER_ADDR": master_addr, - "MASTER_PORT": str(binding.master_port), - } - - def _required_index(env: Mapping[str, str], primary: str, fallback: str) -> int: value = env.get(primary, env.get(fallback)) if value is None: @@ -263,8 +236,6 @@ def main(argv: Sequence[str] | None = None) -> int: PUZZLETRON_MASTER_PORT=str(binding.master_port), PUZZLETRON_RENDEZVOUS_ID=binding.rendezvous_id, ) - if TaskLauncher(args.launcher) is TaskLauncher.DIRECT: - env.update(_direct_distributed_env(binding)) print( "puzzletron binding " f"host={binding.hostname} task={binding.task_index} " diff --git a/modelopt/torch/puzzletron/stages/pipeline.py b/modelopt/torch/puzzletron/stages/pipeline.py index 26e0c059416..6c54623fd77 100644 --- a/modelopt/torch/puzzletron/stages/pipeline.py +++ b/modelopt/torch/puzzletron/stages/pipeline.py @@ -33,7 +33,6 @@ from ..rpc_eval import EvaluationCache, EvaluationRequest, EvaluationResult from ..scoring_parent import ensure_scoring_parent from ..subblock_stats.measurements import apply_vllm_measurement, normalize_vllm_measurements -from ..tools.hydra_utils import clone_hydra_config from .common import complete_stage, experiment_dir, stage_manifest_path if TYPE_CHECKING: @@ -218,10 +217,10 @@ def _calculate_static_workload_stats(config: dict[str, Any], hydra_cfg: Any) -> "osl": int(hydra_cfg.calc_subblock_stats.generation_seq_len), "batch_size": int(hydra_cfg.calc_subblock_stats.batch_sizes[0]), } - } + } for raw_workload in workloads.values(): workload = dict(raw_workload or {}) - selected = clone_hydra_config(hydra_cfg) + selected = OmegaConf.create(OmegaConf.to_container(hydra_cfg, resolve=True)) OmegaConf.set_struct(selected, False) stats_cfg = selected.calc_subblock_stats concurrency = int(workload.get("concurrency", workload.get("batch_size", 1))) @@ -242,158 +241,6 @@ def _calculate_static_workload_stats(config: dict[str, Any], hydra_cfg: Any) -> launch_calc_subblock_stats(selected) -def _scenario_hidden_width(puzzle_dir: Path) -> int | None: - manifest_path = puzzle_dir / "scenario_manifest.json" - if not manifest_path.is_file(): - return None - manifest = json.loads(manifest_path.read_text()) - hidden_width = manifest.get("hidden_width") - return int(hidden_width) if hidden_width is not None else None - - -def _has_runtime_measurement( - path: Path, - *, - hidden_width: int, - measurement: Any, - allow_missing_workload_id: bool = False, -) -> bool: - try: - payload = json.loads(path.read_text()) - except (OSError, ValueError): - return False - if not isinstance(payload, list): - return False - expected_backend = (measurement.runtime_stats or {}).get("backend") - for entry in payload: - if not isinstance(entry, dict): - continue - args = entry.get("args") or {} - if not isinstance(args, dict) or args.get("runtime_stats") is not True: - continue - if int(args.get("n_embd", -1)) != int(hidden_width): - continue - if args.get("weights_dtype") != "torch.bfloat16": - continue - if int(args.get("batch_size", -1)) != int(measurement.batch_size): - continue - if int(args.get("prefill_seq_len", -1)) != int(measurement.prefill_seq_len): - continue - if int(args.get("generation_seq_len", -1)) != int(measurement.generation_seq_len): - continue - if int(args.get("max_num_seqs", -1)) != int(measurement.max_num_seqs): - continue - if args.get("runtime_granularity", "subblock") != measurement.granularity: - continue - if expected_backend is not None and args.get("runtime_backend") != expected_backend: - continue - workload_id = args.get("workload_id") - if workload_id is None and not allow_missing_workload_id: - continue - if workload_id is not None and workload_id != measurement.measurement_id: - continue - return True - return False - - -def _runtime_measurement_candidate_paths( - *, - config: dict[str, Any], - puzzle_dir: Path, - stats_path: Path, - measurement: Any, -) -> list[tuple[Path, bool]]: - stats_name = str( - (config.get("vllm_stats") or {}).get("subblock_stats_filename", stats_path.name) - ) - candidates: list[tuple[Path, bool]] = [(stats_path, False)] - root_dir = puzzle_dir - if ( - puzzle_dir.name == "depth-00" - and puzzle_dir.parent.name.startswith("width-") - and puzzle_dir.parent.parent.name == "scenarios" - ): - root_dir = puzzle_dir.parent.parent.parent - candidates.append((root_dir / stats_name, False)) - relative_stats_path = getattr(measurement, "relative_stats_path", None) - if relative_stats_path is not None: - candidates.append((root_dir / Path(relative_stats_path), True)) - return candidates - - -def _runtime_reuse_source_path( - config: dict[str, Any], - *, - puzzle_dir: Path, - stats_path: Path, - hidden_width: int, - measurement: Any, -) -> Path: - candidates = _runtime_measurement_candidate_paths( - config=config, - puzzle_dir=puzzle_dir, - stats_path=stats_path, - measurement=measurement, - ) - for candidate, allow_missing_workload_id in dict.fromkeys(candidates): - if _has_runtime_measurement( - candidate, - hidden_width=hidden_width, - measurement=measurement, - allow_missing_workload_id=allow_missing_workload_id, - ): - return candidate - raise RuntimeError( - "build-library cannot refresh width-scenario runtime stats because no " - f"reusable vLLM measurement for width {hidden_width}, workload " - f"{measurement.measurement_id!r} was found in {[str(path) for path, _ in candidates]}" - ) - - -def _refresh_scenario_runtime_workload_stats( - config: dict[str, Any], - hydra_cfg: Any, - stats_path: Path, -) -> None: - """Refresh width-scenario runtime rows with the local parameter inventory identity.""" - from ..subblock_stats.calc_subblock_stats import launch_calc_subblock_stats - - puzzle_dir = _puzzle_dir(config, hydra_cfg) - hidden_width = _scenario_hidden_width(puzzle_dir) - if hidden_width is None: - return - measurements = normalize_vllm_measurements(config) - for measurement in measurements.values(): - if measurement.model_hidden_sizes and hidden_width not in measurement.model_hidden_sizes: - continue - source_path = _runtime_reuse_source_path( - config, - puzzle_dir=puzzle_dir, - stats_path=stats_path, - hidden_width=hidden_width, - measurement=measurement, - ) - selected = clone_hydra_config(hydra_cfg) - stats_cfg = selected.calc_subblock_stats - stats_cfg.model_hidden_sizes = [hidden_width] - stats_cfg.batch_sizes = [measurement.batch_size] - stats_cfg.prefill_seq_len = measurement.prefill_seq_len - stats_cfg.generation_seq_len = measurement.generation_seq_len - if stats_cfg.get("runtime_stats") is None: - stats_cfg.runtime_stats = {} - runtime_cfg = stats_cfg.runtime_stats - for key, value in dict(measurement.runtime_stats).items(): - runtime_cfg[key] = _plain(value) - runtime_cfg.enabled = True - runtime_cfg.reuse_stats_path = str(source_path) - runtime_cfg.workload_id = measurement.measurement_id - runtime_cfg.reuse_workload_id_if_missing = measurement.measurement_id - runtime_cfg.max_num_seqs = measurement.max_num_seqs - runtime_cfg.granularity = measurement.granularity - stats_cfg.merge_with_existing_stats = True - launch_calc_subblock_stats(selected) - - def _write_runtime_subblock_library(path: Path, block_configs: tuple[Any, ...]) -> None: """Write the legacy subblock-library input without assembling a replacement library.""" rows = [] @@ -1153,8 +1000,6 @@ def build_library_stage(config: dict[str, Any], manifest: StageManifest): ) launch_build_replacement_library(hydra_cfg) - if _vllm_stats_is_explicit(config): - _refresh_scenario_runtime_workload_stats(config, hydra_cfg, stats_path) _calculate_static_workload_stats(config, hydra_cfg) # Keep this stage-local so handler registration and config preflight do not # load build-library implementation dependencies before the stage executes. diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py index 38580f156b7..8f44bc69e85 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py @@ -259,7 +259,6 @@ def _runtime_measurement_fields( "latency_difference_negative", ) _REUSABLE_RUNTIME_ARG_FIELDS = ( - "workload_id", "runtime_granularity", "runtime_backend", "num_iters", @@ -271,67 +270,7 @@ def _runtime_measurement_fields( ) -def _runtime_reuse_key_from_args( - args: Mapping, - *, - fallback_workload_id: str | None = None, -) -> tuple | None: - """Return the exact runtime-reuse identity represented by persisted args.""" - - if not args.get("runtime_stats") or args.get("weights_dtype") != str(torch.bfloat16): - return None - workload_id = args.get("workload_id") - if workload_id is None: - workload_id = fallback_workload_id - return ( - int(args["n_embd"]), - int(args["batch_size"]), - int(args.get("prefill_seq_len")), - int(args.get("generation_seq_len")), - args.get("max_num_seqs"), - args.get("runtime_granularity", "subblock"), - args.get("runtime_backend"), - args.get("num_iters"), - args.get("num_warmup_iters"), - args.get("repeat_block_n_times"), - _freeze_stats_args(args.get("vllm_args")), - workload_id, - ) - - -def _runtime_reuse_key( - *, - width: int, - batch_size: int, - prefill_seq_len: int, - generation_seq_len: int, - runtime_stats_config: Mapping, -) -> tuple: - """Return the exact runtime-reuse identity requested by this calculation.""" - - return ( - int(width), - int(batch_size), - int(prefill_seq_len), - int(generation_seq_len), - runtime_stats_config.get("max_num_seqs"), - runtime_stats_config.get("granularity", "subblock"), - runtime_stats_config.get("backend"), - runtime_stats_config.get("num_iters", 30), - runtime_stats_config.get("num_warmup_iters", 10), - max(2, int(runtime_stats_config.get("repeat_block_n_times", 10))), - _freeze_stats_args([str(arg) for arg in runtime_stats_config.get("vllm_args", [])]), - runtime_stats_config.get("workload_id"), - ) - - -def _reuse_runtime_stats( - target: dict, - source: dict, - *, - source_path: str, - fallback_workload_id: str | None = None, -) -> dict: +def _reuse_runtime_stats(target: dict, source: dict, *, source_path: str) -> dict: """Overlay immutable measured latency onto refreshed static statistics.""" # Synthetic vLLM timings are layer-independent: collection benchmarks the @@ -363,10 +302,7 @@ def _reuse_runtime_stats( target_args["runtime_stats"] = True target_args["runtime_reuse_source"] = str(source_path) for field in _REUSABLE_RUNTIME_ARG_FIELDS: - value = source_args.get(field) - if field == "workload_id" and value is None: - value = target_args.get("workload_id") or fallback_workload_id - target_args[field] = value + target_args[field] = source_args.get(field) source_non_block = source.get("non_block", {}) target_non_block = target.setdefault("non_block", {}) @@ -876,7 +812,6 @@ def calculate_subblock_stats( max_num_seqs=( runtime_stats_config.get("max_num_seqs") if runtime_stats_enabled else None ), - workload_id=runtime_stats_config.get("workload_id") if runtime_stats_enabled else None, repeat_block_n_times=( max(2, int(runtime_stats_config.get("repeat_block_n_times", 10))) if runtime_stats_enabled @@ -1225,7 +1160,6 @@ def _arg_signature(args: dict) -> tuple: runtime_stats, args.get("runtime_granularity") if runtime_stats else None, args.get("max_num_seqs") if runtime_stats else None, - args.get("workload_id") if runtime_stats else None, args.get("runtime_selection_identity"), args.get("parameter_inventory_identity"), ) @@ -1240,7 +1174,6 @@ def _subblock_stats_already_complete( runtime_stats_enabled: bool, runtime_granularity: str = "subblock", runtime_max_num_seqs: int | None = None, - runtime_workload_id: str | None = None, runtime_selection_identity: str | None = None, parameter_inventory_identities: Mapping[int, str] | None = None, prefill_seq_len: int = 2048, @@ -1288,7 +1221,6 @@ def _entry_subblock_keys(entry: dict) -> set[tuple[SubblockConfig, int]]: runtime_expected, runtime_granularity if runtime_expected else None, runtime_max_num_seqs if runtime_expected else None, - runtime_workload_id if runtime_expected else None, ( runtime_selection_identity if runtime_expected @@ -1358,8 +1290,6 @@ def calculate_subblock_stats_for_puzzle_dir( batch_sizes = [ int(batch_size) for batch_size in batch_sizes.strip("[]").replace(" ", "").split(",") ] - else: - batch_sizes = list(batch_sizes) master_puzzle_dir = Path(master_puzzle_dir) teacher_dir = ( @@ -1403,7 +1333,7 @@ def calculate_subblock_stats_for_puzzle_dir( teacher_hidden_size = int(lm_config.hidden_size) model_hidden_sizes = _unique_hidden_sizes(model_hidden_sizes, teacher_hidden_size) runtime_reuse_path = runtime_stats_config.get("reuse_stats_path") - runtime_reuse_by_key: dict[tuple, dict] = {} + runtime_reuse_by_width: dict[int, dict] = {} if runtime_stats_enabled and runtime_reuse_path: runtime_reuse_path = Path(str(runtime_reuse_path)) if not runtime_reuse_path.is_file(): @@ -1415,30 +1345,17 @@ def calculate_subblock_stats_for_puzzle_dir( f"Reusable runtime stats file does not exist: {runtime_reuse_path}" ) reusable_entries = json.loads(runtime_reuse_path.read_text()) - fallback_workload_id = runtime_stats_config.get("reuse_workload_id_if_missing") - for entry in reusable_entries: - key = _runtime_reuse_key_from_args( - entry.get("args", {}), - fallback_workload_id=fallback_workload_id, - ) - if key is not None: - runtime_reuse_by_key[key] = entry - requested_runtime_keys = { - _runtime_reuse_key( - width=int(width), - batch_size=int(batch_size), - prefill_seq_len=int(prefill_seq_len), - generation_seq_len=int(generation_seq_len), - runtime_stats_config=runtime_stats_config, - ) - for width in model_hidden_sizes - for batch_size in batch_sizes + runtime_reuse_by_width = { + int(entry["args"]["n_embd"]): entry + for entry in reusable_entries + if entry.get("args", {}).get("runtime_stats") + and entry.get("args", {}).get("weights_dtype") == str(torch.bfloat16) } - missing_runtime_keys = requested_runtime_keys - set(runtime_reuse_by_key) - if missing_runtime_keys: + missing_runtime_widths = set(model_hidden_sizes) - set(runtime_reuse_by_width) + if missing_runtime_widths: raise ValueError( - f"Reusable runtime stats {runtime_reuse_path} are missing requested " - f"runtime identities {sorted(missing_runtime_keys)}" + f"Reusable runtime stats {runtime_reuse_path} are missing widths " + f"{sorted(missing_runtime_widths)}" ) runtime_selection_identity = "reuse-" + hashlib.sha256( runtime_reuse_path.read_bytes() @@ -1493,9 +1410,6 @@ def calculate_subblock_stats_for_puzzle_dir( runtime_max_num_seqs=calc_subblock_stats_config.get("runtime_stats", {}).get( "max_num_seqs" ), - runtime_workload_id=calc_subblock_stats_config.get("runtime_stats", {}).get( - "workload_id" - ), runtime_selection_identity=runtime_selection_identity, parameter_inventory_identities=parameter_inventory_identities, prefill_seq_len=prefill_seq_len, @@ -1541,17 +1455,11 @@ def calculate_subblock_stats_for_puzzle_dir( curr_runtime_stats_enabled = ( runtime_stats_enabled if weights_dtype == torch.bfloat16 else False ) - reused_runtime_stats = None - if curr_runtime_stats_enabled: - reused_runtime_stats = runtime_reuse_by_key.get( - _runtime_reuse_key( - width=int(model_hidden_size), - batch_size=int(batch_size), - prefill_seq_len=int(prefill_seq_len), - generation_seq_len=int(generation_seq_len), - runtime_stats_config=runtime_stats_config, - ) - ) + reused_runtime_stats = ( + runtime_reuse_by_width.get(int(model_hidden_size)) + if curr_runtime_stats_enabled + else None + ) curr_subblock_stats = calculate_subblock_stats( calc_subblock_stats_config, @@ -1581,7 +1489,6 @@ def calculate_subblock_stats_for_puzzle_dir( curr_subblock_stats, reused_runtime_stats, source_path=str(runtime_reuse_path), - fallback_workload_id=runtime_stats_config.get("workload_id"), ) curr_subblock_stats["args"]["runtime_selection_identity"] = ( runtime_selection_identity diff --git a/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py b/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py index ee127c27999..616b4d24d9b 100644 --- a/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py +++ b/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py @@ -27,7 +27,6 @@ _prepare_vllm_checkpoint, _profile_command, _server_max_model_len, - _server_vllm_args, _topology_vllm_args, ) @@ -156,49 +155,6 @@ def test_vllm_topology_args_enable_dp_and_expert_parallel_only_when_requested(): assert "--expert-parallel-size" not in ep_args -def test_server_vllm_args_default_max_num_seqs_follows_concurrency(monkeypatch, tmp_path): - monkeypatch.setattr( - "modelopt.torch.puzzletron.benchmarks.aiperf._descriptor_vllm_args", - lambda checkpoint_dir: ["--descriptor-arg"], - ) - - args = _server_vllm_args( - tmp_path, - { - "tensor_parallel_size": 1, - "pipeline_parallel_size": 1, - "data_parallel_size": 1, - "gpu_group_size": 1, - }, - (1, 4), - ) - - assert args[args.index("--max-num-seqs") + 1] == "4" - assert "--descriptor-arg" in args - - -def test_server_vllm_args_respects_explicit_max_num_seqs(monkeypatch, tmp_path): - monkeypatch.setattr( - "modelopt.torch.puzzletron.benchmarks.aiperf._descriptor_vllm_args", - lambda checkpoint_dir: [], - ) - - args = _server_vllm_args( - tmp_path, - { - "tensor_parallel_size": 1, - "pipeline_parallel_size": 1, - "data_parallel_size": 1, - "gpu_group_size": 1, - "extra_vllm_args": ["--max_num_seqs=8"], - }, - (1,), - ) - - assert "--max-num-seqs" not in args - assert "--max_num_seqs=8" in args - - def test_profile_command_maps_each_workload_answer_to_aiperf_cli(tmp_path): command = _profile_command( executable=Path("/opt/aiperf/bin/aiperf"), diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index 9b628faefd6..c2bfab06d8f 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -194,47 +194,6 @@ def test_render_sbatch_script_omits_gpu_requests_for_cpu_stage(): assert "--gpu-bind" not in srun -def test_render_sbatch_script_sets_writable_enroot_paths_before_srun(): - runner = RunnerEnvironment( - kind="slurm", - contract=ExecutionContract( - repository="/repo", - venv="/repo/.venv", - container="/images/pytorch.sqsh", - container_mounts="/repo:/repo", - ), - slurm=SlurmRunnerConfig(account="acct", partition_cpu="cpu"), - ) - attempt = AttemptSpec( - attempt_id="a1", - work_id="build_library:0", - stage_id="build_library", - command=CommandSpec(argv=("python", "worker.py"), log_path="/tmp/out.log"), - allocation_nodes=1, - allocation_gpus=0, - metadata={"gpus_per_node": 0, "partition": "cpu"}, - task_topology=TaskTopology(task_count=1, gpus_per_task=0), - ) - - script = render_sbatch_script( - attempt=attempt, - runner=runner, - partition="cpu", - account="acct", - time_limit="4:00:00", - qos=None, - job_name="pt-build", - ) - - enroot_cache = "export ENROOT_CACHE_PATH=/repo/.cache/enroot/cache" - enroot_runtime = ( - 'export ENROOT_RUNTIME_PATH="/tmp/puzzletron-enroot-${USER:-$(id -u)}/runtime"' - ) - assert enroot_cache in script - assert enroot_runtime in script - assert script.index(enroot_cache) < script.index("srun ") - - def test_vllm_aggregation_uses_slurm_execution_contract(tmp_path: Path, monkeypatch): """Controller-side merges must run in the same container/venv as workers.""" @@ -529,8 +488,6 @@ def test_depth_pool_uses_one_four_node_gang_allocation(tmp_path: Path): assert attempt.allocation_nodes == 4 assert attempt.allocation_gpus == 32 assert attempt.metadata["kill_on_bad_exit"] is True - assert attempt.command.env["WORKER_WORLD_SIZE"] == "8" - assert "WORLD_SIZE" not in attempt.command.env assert attempt.command.argv[-1].endswith("run_depth_pool.sh") script = render_sbatch_script( @@ -546,8 +503,6 @@ def test_depth_pool_uses_one_four_node_gang_allocation(tmp_path: Path): assert "#SBATCH --ntasks=4" in script assert "#SBATCH --gpus-per-node=8" in script assert "--kill-on-bad-exit=1" in script - assert "export WORKER_WORLD_SIZE=8" in script - assert "export WORLD_SIZE=8" not in script assert "run_depth_pool.sh" in script @@ -596,8 +551,6 @@ def test_depth_pool_packs_four_two_gpu_workers_per_node(tmp_path: Path): assert attempt.allocation_gpus == 16 assert attempt.task_topology.task_count == 8 assert attempt.task_topology.gpus_per_task == 2 - assert attempt.command.env["WORKER_WORLD_SIZE"] == "2" - assert "WORLD_SIZE" not in attempt.command.env script = render_sbatch_script( attempt=attempt, @@ -712,8 +665,6 @@ def test_replacement_pool_uses_one_four_node_gang_allocation(tmp_path: Path): assert attempt.allocation_gpus == 32 assert attempt.metadata["kill_on_bad_exit"] is True assert attempt.metadata["partition"] == "batch" - assert attempt.command.env["WORKER_WORLD_SIZE"] == "8" - assert "WORLD_SIZE" not in attempt.command.env assert attempt.command.argv[-1].endswith("run_replacement_pool.sh") @@ -773,8 +724,6 @@ def test_replacement_pool_splits_workers_across_embedding_widths(tmp_path: Path) assert [attempt.allocation_gpus for attempt in attempts] == [16, 16] assert [attempt.task_topology.task_count for attempt in attempts] == [4, 4] assert [attempt.task_topology.gpus_per_task for attempt in attempts] == [4, 4] - assert [attempt.command.env["WORKER_WORLD_SIZE"] for attempt in attempts] == ["4", "4"] - assert all("WORLD_SIZE" not in attempt.command.env for attempt in attempts) assert [attempt.command.env["WORKER_COUNT"] for attempt in attempts] == ["4", "4"] assert [ attempt.command.env["FINALIZE_EXPECTED_COMPLETIONS"] for attempt in attempts diff --git a/tests/unit/torch/puzzletron/test_orchestration_task_topology.py b/tests/unit/torch/puzzletron/test_orchestration_task_topology.py index fe15e657068..22791749112 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_task_topology.py +++ b/tests/unit/torch/puzzletron/test_orchestration_task_topology.py @@ -156,56 +156,6 @@ def fake_execvpe(executable, command, env) -> None: assert result == 0 assert captured["env"]["CUDA_VISIBLE_DEVICES"] == expected assert captured["env"]["PUZZLETRON_TASK_LAUNCHER"] == "direct" - assert captured["env"]["RANK"] == "0" - assert captured["env"]["WORLD_SIZE"] == "1" - assert captured["env"]["LOCAL_RANK"] == "0" - assert captured["env"]["LOCAL_WORLD_SIZE"] == "1" - assert captured["env"]["MASTER_ADDR"] == "127.0.0.1" - assert captured["env"]["MASTER_PORT"].isdigit() - - -def test_cpu_direct_task_launcher_sets_single_rank_env(monkeypatch) -> None: - captured: dict[str, object] = {} - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") - monkeypatch.setenv("PUZZLETRON_TASK_INDEX", "0") - monkeypatch.setenv("PUZZLETRON_LOCAL_TASK_INDEX", "0") - monkeypatch.setenv("PUZZLETRON_TASK_HOSTS", "cpu-a") - - def fake_execvpe(executable, command, env) -> None: - captured.update(executable=executable, command=command, env=env) - - monkeypatch.setattr(task_launcher.os, "execvpe", fake_execvpe) - - result = task_launcher.main( - [ - "--attempt-id", - "attempt-a", - "--nodes", - "1", - "--gpus-per-node", - "0", - "--task-count", - "1", - "--gpus-per-task", - "0", - "--tasks-per-group", - "1", - "--launcher", - "direct", - "--", - "python", - "worker.py", - ] - ) - - assert result == 0 - assert captured["env"]["CUDA_VISIBLE_DEVICES"] == "" - assert captured["env"]["RANK"] == "0" - assert captured["env"]["WORLD_SIZE"] == "1" - assert captured["env"]["LOCAL_RANK"] == "0" - assert captured["env"]["LOCAL_WORLD_SIZE"] == "1" - assert captured["env"]["MASTER_ADDR"] == "127.0.0.1" - assert captured["env"]["MASTER_PORT"].isdigit() def _task_binding(*, group_size: int) -> task_launcher.TaskBinding: diff --git a/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py b/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py index 4e9a6e39c6d..94c160c95a5 100644 --- a/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py +++ b/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py @@ -22,7 +22,6 @@ from types import SimpleNamespace import pytest -import torch from immutabledict import immutabledict from omegaconf import OmegaConf @@ -67,7 +66,6 @@ _reuse_runtime_stats, _runtime_measurement_fields, _select_runtime_subblock_configs, - _subblock_stats_already_complete, _unique_hidden_sizes, _validate_sparse_runtime_settings, ) @@ -189,7 +187,6 @@ def test_runtime_stats_can_be_reused_while_static_metrics_are_refreshed(): "runtime_stats": True, "runtime_granularity": "subblock", "runtime_backend": "vllm", - "workload_id": "serving-default", }, "subblocks": [ { @@ -214,7 +211,6 @@ def test_runtime_stats_can_be_reused_while_static_metrics_are_refreshed(): assert result["args"]["runtime_stats"] is True assert result["args"]["runtime_reuse_source"] == "existing.json" - assert result["args"]["workload_id"] == "serving-default" assert result["subblocks"][0]["runtime_ms"] == 7.0 assert result["subblocks"][1]["runtime_ms"] == 7.0 assert result["subblocks"][0]["num_params"] == 123 @@ -225,42 +221,6 @@ def test_runtime_stats_can_be_reused_while_static_metrics_are_refreshed(): assert result["non_block"]["runtime_ms"] == 1.0 -def test_runtime_stats_reuse_preserves_target_workload_when_source_is_legacy(): - config = FFNConfig(intermediate_size=16) - target = { - "args": {"runtime_stats": False, "workload_id": None}, - "subblocks": [{"subblock_config": config.to_dict(), "parent_layer_index": 0}], - "non_block": {}, - } - source = { - "args": { - "runtime_stats": True, - "runtime_granularity": "subblock", - "runtime_backend": "vllm", - "workload_id": None, - }, - "subblocks": [ - { - "subblock_config": config.to_dict(), - "parent_layer_index": 0, - "runtime_ms": 7.0, - } - ], - "non_block": {}, - } - - result = _reuse_runtime_stats( - target, - source, - source_path="measurement.json", - fallback_workload_id="serving-default", - ) - - assert result["args"]["runtime_stats"] is True - assert result["args"]["workload_id"] == "serving-default" - assert result["subblocks"][0]["runtime_ms"] == 7.0 - - def test_runtime_block_candidates_load_converted_teacher_without_replacement_library(tmp_path): attention = AttentionConfig(num_query_heads=8, num_kv_heads=2) teacher_block = { @@ -1185,232 +1145,6 @@ def raise_candidate_library_import_error(*_args, **_kwargs): assert not (tmp_path / "manifests" / "build_library.json").exists() -def test_static_workload_stats_accept_instantiated_activation_pass_objects(monkeypatch): - class DummyMixin: - pass - - hydra_cfg = OmegaConf.create( - { - "calc_subblock_stats": { - "batch_sizes": [1], - "prefill_seq_len": 64, - "generation_seq_len": 16, - "runtime_stats": {"enabled": True}, - "merge_with_existing_stats": False, - }, - "pruning": { - "activation_passes": [ - { - "name": "attention_grouped", - "pruning_mixin": DummyMixin, - } - ] - }, - }, - flags={"allow_objects": True}, - ) - captured = [] - monkeypatch.setattr( - "modelopt.torch.puzzletron.subblock_stats.calc_subblock_stats.launch_calc_subblock_stats", - lambda cfg: captured.append(cfg), - ) - - pipeline_stages._calculate_static_workload_stats( - { - "mip": { - "workloads": { - "short": {"batch_size": 2, "isl": 32, "osl": 8}, - } - } - }, - hydra_cfg, - ) - - assert len(captured) == 1 - stats_cfg = captured[0].calc_subblock_stats - assert stats_cfg.batch_sizes == [2] - assert stats_cfg.prefill_seq_len == 32 - assert stats_cfg.generation_seq_len == 8 - assert stats_cfg.runtime_stats.enabled is False - assert stats_cfg.merge_with_existing_stats is True - assert ( - captured[0].pruning.activation_passes[0].pruning_mixin - is DummyMixin - ) - assert hydra_cfg.calc_subblock_stats.batch_sizes == [1] - assert hydra_cfg.calc_subblock_stats.runtime_stats.enabled is True - - -def test_width_scenario_runtime_stats_reuse_root_measurement(tmp_path, monkeypatch): - scenario = tmp_path / "scenarios" / "width-2688" / "depth-00" - scenario.mkdir(parents=True) - (scenario / "scenario_manifest.json").write_text( - json.dumps({"status": "complete", "hidden_width": 2688}) - ) - (tmp_path / "subblock_stats.json").write_text( - json.dumps( - [ - { - "args": { - "runtime_stats": True, - "runtime_granularity": "subblock", - "runtime_backend": "vllm", - "weights_dtype": "torch.bfloat16", - "batch_size": 1, - "prefill_seq_len": 4096, - "generation_seq_len": 1024, - "max_num_seqs": 1, - "n_embd": 2688, - "workload_id": "serving-default", - }, - "subblocks": [], - } - ] - ) - ) - hydra_cfg = OmegaConf.create( - { - "puzzle_dir": str(scenario), - "calc_subblock_stats": { - "model_hidden_sizes": [2688, 2560], - "batch_sizes": [1], - "prefill_seq_len": 4096, - "generation_seq_len": 1024, - "subblock_stats_filename": "subblock_stats.json", - "merge_with_existing_stats": True, - "runtime_stats": { - "enabled": True, - "backend": "vllm", - "granularity": "subblock", - "max_num_seqs": 1, - "num_iters": 30, - "num_warmup_iters": 10, - "repeat_block_n_times": 4, - "topology": {"gpu_group_size": 1}, - }, - }, - } - ) - config = { - "puzzle_dir": str(scenario), - "vllm_stats": { - "enabled": True, - "subblock_stats_filename": "subblock_stats.json", - "measurements": { - "serving-default": { - "prefill_seq_len": 4096, - "generation_seq_len": 1024, - "batch_size": 1, - "max_num_seqs": 1, - "granularity": "subblock", - "runtime_stats": { - "backend": "vllm", - "granularity": "subblock", - "max_num_seqs": 1, - "num_iters": 30, - "num_warmup_iters": 10, - "repeat_block_n_times": 4, - "topology": {"gpu_group_size": 1}, - }, - } - }, - }, - } - calls = [] - - def launch(cfg): - calls.append(cfg) - assert list(cfg.calc_subblock_stats.model_hidden_sizes) == [2688] - assert cfg.calc_subblock_stats.runtime_stats.enabled is True - assert cfg.calc_subblock_stats.runtime_stats.workload_id == "serving-default" - assert ( - cfg.calc_subblock_stats.runtime_stats.reuse_workload_id_if_missing - == "serving-default" - ) - assert cfg.calc_subblock_stats.runtime_stats.reuse_stats_path == str( - tmp_path / "subblock_stats.json" - ) - - monkeypatch.setattr( - "modelopt.torch.puzzletron.subblock_stats.calc_subblock_stats.launch_calc_subblock_stats", - launch, - ) - - pipeline_stages._refresh_scenario_runtime_workload_stats( - config, - hydra_cfg, - scenario / "subblock_stats.json", - ) - - assert len(calls) == 1 - - -def test_runtime_stats_resume_signature_includes_workload_id(): - config = FFNConfig(intermediate_size=16) - runtime_fields = { - "runtime_ms": 1.0, - "prefill_runtime_ms": 1.0, - "decode_runtime_ms": 1.0, - "decode_runtime_ms_per_token": 1.0, - "weight_memory_mib": 1.0, - "kv_cache_bytes_per_token": 1.0, - "state_cache_bytes_per_sequence": 1.0, - "prefill_flops": 1.0, - "decode_flops": 1.0, - } - existing = { - "args": { - "batch_size": 1, - "prefill_seq_len": 4096, - "generation_seq_len": 1024, - "weights_dtype": "torch.bfloat16", - "activations_dtype": "torch.bfloat16", - "kv_cache_dtype": "torch.bfloat16", - "n_embd": 2688, - "runtime_stats": True, - "runtime_granularity": "subblock", - "max_num_seqs": 1, - "workload_id": "serving-default", - "runtime_selection_identity": "reuse-root-aggregate", - "parameter_inventory_identity": "scenario-inventory", - }, - "subblocks": [ - { - "subblock_config": config.to_dict(), - "parent_layer_index": 0, - **runtime_fields, - "additive_metric_provenance": { - field: "test" for field in runtime_fields - }, - } - ], - } - kwargs = dict( - existing_stats=[existing], - subblock_configs=[_indexed(config, 0)], - batch_sizes=[1], - data_types=[(torch.bfloat16, torch.bfloat16, torch.bfloat16)], - model_hidden_sizes=[2688], - runtime_stats_enabled=True, - runtime_granularity="subblock", - runtime_max_num_seqs=1, - runtime_selection_identity="reuse-root-aggregate", - parameter_inventory_identities={2688: "scenario-inventory"}, - prefill_seq_len=4096, - generation_seq_len=1024, - ) - - assert _subblock_stats_already_complete( - **kwargs, - runtime_workload_id="serving-default", - ) - assert not _subblock_stats_already_complete( - **kwargs, - runtime_workload_id="different-workload", - ) - assert hydra_cfg.calc_subblock_stats.merge_with_existing_stats is False - - def test_sparse_runtime_selection_is_unique_and_layer_independent(): teacher_ffn = FFNConfig(intermediate_size=16) reduced_ffn = FFNConfig(intermediate_size=8) From 63c9f29b2e60dfc3abd1f034b67aaeff78a9b61e Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Fri, 7 Aug 2026 12:46:43 -0700 Subject: [PATCH 09/16] fix: validate lmms-eval task completion Signed-off-by: Grzegorz Karch --- modelopt/torch/puzzletron/post_mip/runner.py | 129 +++++++++++++++-- .../torch/puzzletron/test_post_mip_runner.py | 132 +++++++++++++++++- 2 files changed, 252 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 1bbe463d0b7..5af05c631bd 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -568,6 +568,14 @@ def _join_cli_values(value: Any, *, path: str) -> str: return ",".join(values) +def _configured_lmms_eval_tasks(settings: Mapping[str, Any]) -> tuple[str, ...]: + tasks = _join_cli_values(settings.get("tasks"), path="downstream_evaluation.config.tasks") + values = tuple(task.strip() for task in tasks.split(",")) + if not values or any(not task for task in values): + raise ValueError("downstream_evaluation.config.tasks must contain non-empty task names") + return values + + def _model_arg_string(values: Mapping[str, Any]) -> str: parts = [] for key, value in values.items(): @@ -645,7 +653,7 @@ def _lmms_eval_command( ) -> tuple[list[str], dict[str, str], float | None]: """Build a deterministic lmms-eval CLI invocation for one realized checkpoint.""" - tasks = _join_cli_values(settings.get("tasks"), path="downstream_evaluation.config.tasks") + tasks = ",".join(_configured_lmms_eval_tasks(settings)) argv = [ *_command_prefix(settings), "--model", @@ -707,6 +715,18 @@ def _metric_key(value: Any) -> str: ) +def _numeric_metrics(task_payload: Mapping[str, Any]) -> dict[str, float]: + metrics = {} + for metric_name, value in task_payload.items(): + if ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and math.isfinite(value) + ): + metrics[str(metric_name)] = float(value) + return metrics + + def _flatten_lmms_eval_metrics(payload: Mapping[str, Any]) -> dict[str, float]: results = payload.get("results") if not isinstance(results, Mapping): @@ -715,16 +735,105 @@ def _flatten_lmms_eval_metrics(payload: Mapping[str, Any]) -> dict[str, float]: for task_name, task_payload in results.items(): if not isinstance(task_payload, Mapping): continue - for metric_name, value in task_payload.items(): - if ( - isinstance(value, (int, float)) - and not isinstance(value, bool) - and math.isfinite(value) - ): - metrics[f"{_metric_key(task_name)}.{_metric_key(metric_name)}"] = float(value) + for metric_name, value in _numeric_metrics(task_payload).items(): + metrics[f"{_metric_key(task_name)}.{_metric_key(metric_name)}"] = value return metrics +def _resolved_lmms_eval_tasks( + payload: Mapping[str, Any], configured_tasks: Sequence[str] +) -> tuple[str, ...]: + group_subtasks = payload.get("group_subtasks") + if not isinstance(group_subtasks, Mapping): + group_subtasks = {} + + def expand(task: str, seen: frozenset[str]) -> tuple[str, ...]: + raw_subtasks = group_subtasks.get(task) + if ( + isinstance(raw_subtasks, Sequence) + and not isinstance(raw_subtasks, str) + and raw_subtasks + and task not in seen + ): + expanded = [] + for raw_subtask in raw_subtasks: + expanded.extend(expand(str(raw_subtask), seen | {task})) + return tuple(dict.fromkeys(expanded)) + return (task,) + + resolved = [] + for task in configured_tasks: + resolved.extend(expand(task, frozenset())) + return tuple(dict.fromkeys(resolved)) + + +def _sample_count(payload: Mapping[str, Any], task: str) -> float | None: + samples = payload.get("n-samples", payload.get("n_samples")) + if not isinstance(samples, Mapping): + return None + value = samples.get(task) + if isinstance(value, Mapping): + if "effective" in value: + value = value["effective"] + elif "original" in value: + value = value["original"] + else: + return None + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return None + return float(value) + + +def _validate_lmms_eval_completion( + payload: Mapping[str, Any], configured_tasks: Sequence[str] +) -> dict[str, float]: + results = payload.get("results") + if not isinstance(results, Mapping): + raise RuntimeError("lmms-eval result is missing the results mapping") + + expected_tasks = _resolved_lmms_eval_tasks(payload, configured_tasks) + missing_results = [task for task in expected_tasks if task not in results] + if missing_results: + raise RuntimeError( + "lmms-eval result is missing configured task results: " + f"{sorted(missing_results)}" + ) + + missing_metrics = [ + task + for task in expected_tasks + if not isinstance(results[task], Mapping) or not _numeric_metrics(results[task]) + ] + if missing_metrics: + raise RuntimeError( + "lmms-eval result has no numeric metrics for configured tasks: " + f"{sorted(missing_metrics)}" + ) + + sample_counts = {} + missing_samples = [] + zero_samples = [] + for task in expected_tasks: + sample_count = _sample_count(payload, task) + if sample_count is None: + missing_samples.append(task) + elif sample_count <= 0: + zero_samples.append(task) + else: + sample_counts[task] = sample_count + if missing_samples: + raise RuntimeError( + "lmms-eval result is missing sample counts for configured tasks: " + f"{sorted(missing_samples)}" + ) + if zero_samples: + raise RuntimeError( + "lmms-eval result has zero effective samples for configured tasks: " + f"{sorted(zero_samples)}" + ) + return sample_counts + + def _lmms_eval_result_payload(output_path: Path) -> tuple[dict[str, Any], Path]: candidates = [] for path in sorted(output_path.rglob("*.json")): @@ -817,6 +926,9 @@ def _downstream_evaluation( except FileNotFoundError as error: tail = _lmms_eval_output_tail(result) raise FileNotFoundError(str(error) + (f": {tail}" if tail else "")) from error + sample_counts = _validate_lmms_eval_completion( + payload, _configured_lmms_eval_tasks(settings) + ) metrics = _flatten_lmms_eval_metrics(payload) if not metrics: raise RuntimeError(f"lmms-eval result has no numeric task metrics: {result_path}") @@ -828,6 +940,7 @@ def _downstream_evaluation( "checkpoint": source.artifact["checkpoint"], "metrics": metrics, "result_path": str(result_path), + "sample_counts": sample_counts, }, ) return { diff --git a/tests/unit/torch/puzzletron/test_post_mip_runner.py b/tests/unit/torch/puzzletron/test_post_mip_runner.py index 432e9a89f1c..7c413da3f4d 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -240,7 +240,12 @@ def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): "results": { "ifeval": {"prompt_level_strict_acc,none": 0.5}, "gsm8k": {"exact_match,strict-match": 0.75}, - } + }, + "group_subtasks": {"ifeval": [], "gsm8k": []}, + "n-samples": { + "ifeval": {"original": 541, "effective": 4}, + "gsm8k": {"original": 1319, "effective": 4}, + }, } ) ) @@ -280,6 +285,131 @@ def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): } assert Path(result["result_path"]).is_file() assert Path(result["raw_result_path"]).name == "results.json" + summary = json.loads(Path(result["result_path"]).read_text()) + assert summary["sample_counts"] == {"gsm8k": 4.0, "ifeval": 4.0} + + +def test_lmms_eval_completion_validates_resolved_task_expansion(): + sample_counts = runner._validate_lmms_eval_completion( + { + "results": { + "arc_challenge": {"acc,none": 0.25}, + "hellaswag": {"acc_norm,none": 0.5}, + }, + "group_subtasks": { + "leaderboard": ["arc_challenge", "hellaswag"], + "arc_challenge": [], + "hellaswag": [], + }, + "n-samples": { + "arc_challenge": {"original": 1172, "effective": 8}, + "hellaswag": {"original": 10042, "effective": 8}, + }, + }, + ("leaderboard",), + ) + + assert sample_counts == {"arc_challenge": 8.0, "hellaswag": 8.0} + + +def test_downstream_evaluation_rejects_missing_configured_task(monkeypatch, tmp_path): + def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): + del env, capture_output, text, timeout, check + output = Path(cwd) + (output / "results.json").write_text( + json.dumps( + { + "results": {"ifeval": {"prompt_level_strict_acc,none": 0.5}}, + "group_subtasks": {"ifeval": []}, + "n-samples": {"ifeval": {"original": 541, "effective": 4}}, + } + ) + ) + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(runner.subprocess, "run", fake_run) + node = SimpleNamespace( + node_id="lmms_eval", + flow_id="runtime", + stage_id="post.runtime.lmms_eval", + config={ + "config": { + "command_prefix": ["python", "-m", "lmms_eval"], + "tasks": ["ifeval", "gsm8k"], + "topology": {"gpu_group_size": 1}, + } + }, + ) + source = SimpleNamespace( + architecture_id="architecture", + artifact_kind=ArtifactKind.CHECKPOINT, + artifact={"checkpoint": str(tmp_path / "checkpoint")}, + ) + + try: + runner._downstream_evaluation( + {"puzzle_dir": str(tmp_path)}, node, source, "execution" + ) + except RuntimeError as error: + message = str(error) + else: + raise AssertionError("expected incomplete lmms-eval result to fail") + + assert "missing configured task results" in message + assert "gsm8k" in message + + +def test_downstream_evaluation_rejects_zero_sample_task(monkeypatch, tmp_path): + def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): + del env, capture_output, text, timeout, check + output = Path(cwd) + (output / "results.json").write_text( + json.dumps( + { + "results": { + "ifeval": {"prompt_level_strict_acc,none": 0.5}, + "gsm8k": {"exact_match,strict-match": 0.75}, + }, + "group_subtasks": {"ifeval": [], "gsm8k": []}, + "n-samples": { + "ifeval": {"original": 541, "effective": 4}, + "gsm8k": {"original": 1319, "effective": 0}, + }, + } + ) + ) + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(runner.subprocess, "run", fake_run) + node = SimpleNamespace( + node_id="lmms_eval", + flow_id="runtime", + stage_id="post.runtime.lmms_eval", + config={ + "config": { + "command_prefix": ["python", "-m", "lmms_eval"], + "tasks": ["ifeval", "gsm8k"], + "topology": {"gpu_group_size": 1}, + } + }, + ) + source = SimpleNamespace( + architecture_id="architecture", + artifact_kind=ArtifactKind.CHECKPOINT, + artifact={"checkpoint": str(tmp_path / "checkpoint")}, + ) + + try: + runner._downstream_evaluation( + {"puzzle_dir": str(tmp_path)}, node, source, "execution" + ) + except RuntimeError as error: + message = str(error) + else: + raise AssertionError("expected zero-sample lmms-eval result to fail") + + assert "zero effective samples" in message + assert "gsm8k" in message def test_downstream_evaluation_reports_lmms_eval_output_when_results_are_missing( From 7eda1409b1895218740a91f4747a2e17623c4b78 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Fri, 7 Aug 2026 12:52:34 -0700 Subject: [PATCH 10/16] fix: align lmms metric suggestions Signed-off-by: Grzegorz Karch --- puzzletron_setup/wizard.py | 22 ++++++++++++++++++- .../torch/puzzletron/test_setup_bundle.py | 18 +++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index 2750be52f6f..90dba4cd26f 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -39,6 +39,10 @@ _MESH_KEYS = ("tp", "cp", "pp", "dp_shard", "dp_replicate", "ep") _DEFAULT_MIP_SOLUTION_COUNT = 3 _DEFAULT_HOMOGENEOUS_SOLUTIONS_PER_SCENARIO = 8 +_DOWNSTREAM_EVALUATION_METRICS_BY_TASK = { + "gsm8k": ("exact_match_strict-match",), + "ifeval": ("prompt_level_strict_acc_none",), +} def _default(state: AnswerState, section: str, key: str, fallback: Any) -> Any: @@ -841,6 +845,20 @@ def _ask_downstream_evaluation_config( } +def _downstream_evaluation_metric_suggestions(node_id: str, config: Mapping[str, Any]) -> list[str]: + """Return filter metric names produced by the downstream-evaluation runner.""" + + suggestions = [] + tasks = config.get("tasks") or () + if isinstance(tasks, str): + tasks = [item.strip() for item in tasks.split(",") if item.strip()] + for task in tasks: + task_name = str(task).strip() + for metric in _DOWNSTREAM_EVALUATION_METRICS_BY_TASK.get(task_name, ()): + suggestions.append(f"{node_id}.{task_name}.{metric}") + return suggestions + + def _default_flow( run_id: str, run: Mapping[str, Any], @@ -1077,7 +1095,9 @@ def _custom_flow( detailed=detailed, moe=moe, ) - available_metrics.append(f"{node_id}.gsm8k.exact_match") + available_metrics.extend( + _downstream_evaluation_metric_suggestions(node_id, node["config"]) + ) elif node_type == "global_kd": node["config"] = {"max_steps": prompts.integer("Global KD steps:", default=128)} elif node_type == "ptq": diff --git a/tests/unit/torch/puzzletron/test_setup_bundle.py b/tests/unit/torch/puzzletron/test_setup_bundle.py index 51f2862282e..5f4f264ea65 100644 --- a/tests/unit/torch/puzzletron/test_setup_bundle.py +++ b/tests/unit/torch/puzzletron/test_setup_bundle.py @@ -34,6 +34,7 @@ _ask_mesh, _ask_mip, _default_flow, + _downstream_evaluation_metric_suggestions, _resource_rows, ) @@ -549,6 +550,23 @@ def test_render_execution_uses_common_mesh_for_post_mip_evaluation_only() -> Non assert execution["post.run.materialized"]["instances"] == 1 +def test_downstream_evaluation_metric_suggestions_match_runner_keys() -> None: + assert _downstream_evaluation_metric_suggestions( + "lmms_eval", + {"tasks": ["ifeval", "gsm8k", "custom_task"]}, + ) == [ + "lmms_eval.ifeval.prompt_level_strict_acc_none", + "lmms_eval.gsm8k.exact_match_strict-match", + ] + assert _downstream_evaluation_metric_suggestions( + "lmms_eval", + {"tasks": "gsm8k,ifeval"}, + ) == [ + "lmms_eval.gsm8k.exact_match_strict-match", + "lmms_eval.ifeval.prompt_level_strict_acc_none", + ] + + def test_render_execution_uses_vllm_mesh_for_post_mip_downstream_evaluation() -> None: state = { "answers": { From db5feb519392ce405afd70073be67aba767da305 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Fri, 7 Aug 2026 15:58:29 -0700 Subject: [PATCH 11/16] fix: guard lmms-eval reserved args Signed-off-by: Grzegorz Karch --- modelopt/torch/puzzletron/post_mip/runner.py | 123 +++++++++++++++++- .../torch/puzzletron/test_post_mip_runner.py | 51 ++++++++ 2 files changed, 171 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 5af05c631bd..cedf2140d14 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -24,6 +24,7 @@ import json import math import os +import shlex import subprocess import sys import traceback @@ -538,6 +539,34 @@ def _aiperf( "reasoning_parser", } ) +_LMMS_EVAL_RESERVED_TOPOLOGY_MODEL_ARG_FIELDS = frozenset( + { + "tensor_parallel_size", + "pipeline_parallel_size", + "data_parallel_size", + "prefill_context_parallel_size", + "decode_context_parallel_size", + "enable_expert_parallel", + "distributed_executor_backend", + "expert_parallel_size", + "gpu_group_size", + "tp", + "pp", + "dp", + "prefill_cp", + "decode_cp", + "ep", + } +) +_LMMS_EVAL_RESERVED_EXTRA_ARG_FLAGS = frozenset( + { + "--model_args", + "--model-args", + "--output_path", + "--output-path", + "--tasks", + } +) def _as_cli_bool(value: bool) -> str: @@ -568,6 +597,63 @@ def _join_cli_values(value: Any, *, path: str) -> str: return ",".join(values) +def _lmms_eval_model_arg_keys(value: str) -> tuple[str, ...]: + keys: list[str] = [] + start = 0 + depth = 0 + quote: str | None = None + escaped = False + + def append(segment: str) -> None: + key, separator, _ = segment.strip().partition("=") + if separator and key.strip(): + keys.append(key.strip()) + + for index, char in enumerate(value): + if escaped: + escaped = False + continue + if quote: + if char == "\\": + escaped = True + elif char == quote: + quote = None + continue + if char in {"'", '"'}: + quote = char + elif char in "([{": + depth += 1 + elif char in ")]}" and depth: + depth -= 1 + elif char == "," and depth == 0: + append(value[start:index]) + start = index + 1 + append(value[start:]) + return tuple(keys) + + +def _lmms_eval_reserved_model_arg_fields(checkpoint_arg: str) -> frozenset[str]: + return frozenset( + key + for key in ( + str(checkpoint_arg).strip(), + *_LMMS_EVAL_RESERVED_TOPOLOGY_MODEL_ARG_FIELDS, + ) + if key + ) + + +def _reject_reserved_lmms_eval_model_args( + keys: Sequence[Any], reserved_fields: frozenset[str] +) -> None: + reserved = sorted({str(key).strip() for key in keys} & reserved_fields) + if reserved: + raise ValueError( + "downstream_evaluation.config.model_args must not set reserved " + f"lmms-eval model arguments: {', '.join(reserved)}" + ) + + def _configured_lmms_eval_tasks(settings: Mapping[str, Any]) -> tuple[str, ...]: tasks = _join_cli_values(settings.get("tasks"), path="downstream_evaluation.config.tasks") values = tuple(task.strip() for task in tasks.split(",")) @@ -601,6 +687,7 @@ def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> checkpoint_arg = str(settings.get("checkpoint_arg", "model")) topology = dict(settings.get("topology") or {}) canonical_topology = normalize_vllm_topology(topology) if topology else {} + reserved_fields = _lmms_eval_reserved_model_arg_fields(checkpoint_arg) derived = { checkpoint_arg: checkpoint, } @@ -621,14 +708,17 @@ def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> derived[key] = settings[key] if isinstance(raw, str): + _reject_reserved_lmms_eval_model_args( + _lmms_eval_model_arg_keys(raw), reserved_fields + ) prefix = raw.strip().strip(",") suffix = _model_arg_string(derived) return ",".join(part for part in (prefix, suffix) if part) if raw is not None and not isinstance(raw, Mapping): raise TypeError("downstream_evaluation.config.model_args must be a mapping or string") + _reject_reserved_lmms_eval_model_args(tuple((raw or {}).keys()), reserved_fields) merged = dict(raw or {}) - for key, value in derived.items(): - merged.setdefault(key, value) + merged.update(derived) return _model_arg_string(merged) @@ -645,6 +735,33 @@ def _command_prefix(settings: Mapping[str, Any]) -> list[str]: return values +def _lmms_eval_extra_args(settings: Mapping[str, Any]) -> list[str]: + raw = settings.get("extra_args") + if raw is None: + return [] + if isinstance(raw, str): + values = shlex.split(raw) + elif isinstance(raw, Sequence): + values = [str(item) for item in raw] + else: + raise TypeError("downstream_evaluation.config.extra_args must be a string or sequence") + if any(not value for value in values): + raise ValueError("downstream_evaluation.config.extra_args must not contain empty values") + reserved = sorted( + { + value.split("=", 1)[0] + for value in values + if value.split("=", 1)[0] in _LMMS_EVAL_RESERVED_EXTRA_ARG_FLAGS + } + ) + if reserved: + raise ValueError( + "downstream_evaluation.config.extra_args must not set reserved " + f"lmms-eval flags: {', '.join(reserved)}" + ) + return values + + def _lmms_eval_command( settings: Mapping[str, Any], *, @@ -692,7 +809,7 @@ def _lmms_eval_command( ) if bool(settings.get("log_samples", False)): argv.append("--log_samples") - argv.extend(str(item) for item in settings.get("extra_args") or ()) + argv.extend(_lmms_eval_extra_args(settings)) env = os.environ.copy() for key, value in dict(settings.get("env") or {}).items(): diff --git a/tests/unit/torch/puzzletron/test_post_mip_runner.py b/tests/unit/torch/puzzletron/test_post_mip_runner.py index 7c413da3f4d..74b12d0fb7f 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -226,6 +226,57 @@ def test_lmms_eval_command_maps_checkpoint_and_vllm_topology(tmp_path): assert timeout == 123 +def test_lmms_eval_command_rejects_reserved_model_args(tmp_path): + cases = ( + ({"model": "/ckpts/wrong"}, "model"), + ("dtype=bfloat16,tensor_parallel_size=1", "tensor_parallel_size"), + ) + for model_args, expected in cases: + try: + runner._lmms_eval_command( + { + "tasks": ["ifeval"], + "topology": {"gpu_group_size": 1}, + "model_args": model_args, + }, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + except ValueError as error: + message = str(error) + else: + raise AssertionError("expected reserved lmms-eval model_args to fail") + + assert "reserved lmms-eval model arguments" in message + assert expected in message + + +def test_lmms_eval_command_rejects_reserved_extra_args(tmp_path): + cases = ( + (["--tasks", "gsm8k"], "--tasks"), + ("--output_path /tmp/other", "--output_path"), + (["--model_args=model=/ckpts/wrong"], "--model_args"), + ) + for extra_args, expected in cases: + try: + runner._lmms_eval_command( + { + "tasks": ["ifeval"], + "topology": {"gpu_group_size": 1}, + "extra_args": extra_args, + }, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + except ValueError as error: + message = str(error) + else: + raise AssertionError("expected reserved lmms-eval extra_args to fail") + + assert "reserved lmms-eval flags" in message + assert expected in message + + def test_downstream_evaluation_runs_lmms_eval_and_flattens_metrics(monkeypatch, tmp_path): captured = {} From bdb20f77674c8cb7c898b0dc3c90ed6113fe1a46 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Sat, 8 Aug 2026 10:37:23 -0700 Subject: [PATCH 12/16] fix: clean up lmms-eval process group Signed-off-by: Grzegorz Karch --- modelopt/torch/puzzletron/post_mip/runner.py | 70 ++++++- .../torch/puzzletron/test_post_mip_runner.py | 176 ++++++++++++++++-- 2 files changed, 230 insertions(+), 16 deletions(-) diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index cedf2140d14..4de690a60b8 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -25,6 +25,7 @@ import math import os import shlex +import signal import subprocess import sys import traceback @@ -567,6 +568,7 @@ def _aiperf( "--tasks", } ) +_LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS = 10.0 def _as_cli_bool(value: bool) -> str: @@ -989,6 +991,69 @@ def _lmms_eval_output_tail(result: subprocess.CompletedProcess[str], *, max_line return "\n".join(sections) +def _signal_lmms_eval_process_group( + process: subprocess.Popen[str], signal_number: int +) -> None: + try: + if os.name == "posix": + os.killpg(process.pid, signal_number) + else: + process.send_signal(signal_number) + except ProcessLookupError: + pass + + +def _lmms_eval_process_group_exists(process: subprocess.Popen[str]) -> bool: + if os.name != "posix": + return process.poll() is None + try: + os.killpg(process.pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def _run_lmms_eval_process( + argv: list[str], + *, + cwd: str, + env: Mapping[str, str], + timeout: float | None, +) -> subprocess.CompletedProcess[str]: + process = subprocess.Popen( + argv, + cwd=cwd, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + start_new_session=os.name == "posix", + ) + try: + stdout, stderr = process.communicate(timeout=timeout) + except subprocess.TimeoutExpired as error: + _signal_lmms_eval_process_group(process, signal.SIGTERM) + try: + stdout, stderr = process.communicate( + timeout=_LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS + ) + except subprocess.TimeoutExpired: + _signal_lmms_eval_process_group(process, signal.SIGKILL) + stdout, stderr = process.communicate() + else: + if _lmms_eval_process_group_exists(process): + _signal_lmms_eval_process_group(process, signal.SIGKILL) + raise subprocess.TimeoutExpired( + argv, + timeout, + output=stdout if stdout is not None else error.output, + stderr=stderr if stderr is not None else error.stderr, + ) from error + return subprocess.CompletedProcess(argv, process.returncode, stdout, stderr) + + def _downstream_evaluation( config: dict[str, Any], node: CompiledPostMIPNode, @@ -1022,14 +1087,11 @@ def _downstream_evaluation( ) # Campaign config controls the executable and arguments, but subprocess receives # an argv list directly; no shell parsing is involved. - result = subprocess.run( + result = _run_lmms_eval_process( argv, cwd=str(output), env=env, - capture_output=True, - text=True, timeout=timeout, - check=False, ) stream_paths = _write_lmms_eval_streams(output, result) if result.returncode: diff --git a/tests/unit/torch/puzzletron/test_post_mip_runner.py b/tests/unit/torch/puzzletron/test_post_mip_runner.py index 74b12d0fb7f..7b650284a51 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -17,6 +17,7 @@ # SPDX-License-Identifier: Apache-2.0 import json +import signal import subprocess from pathlib import Path from types import SimpleNamespace @@ -277,11 +278,162 @@ def test_lmms_eval_command_rejects_reserved_extra_args(tmp_path): assert expected in message +def test_lmms_eval_timeout_terminates_process_group(monkeypatch, tmp_path): + created = [] + signals = [] + + class FakeProcess: + pid = 1234 + returncode = None + + def __init__(self): + self.communicate_timeouts = [] + + def communicate(self, timeout=None): + self.communicate_timeouts.append(timeout) + if len(self.communicate_timeouts) == 1: + raise subprocess.TimeoutExpired( + ["python", "-m", "lmms_eval"], + timeout, + output="partial stdout", + stderr="partial stderr", + ) + self.returncode = -signal.SIGTERM + return "partial stdout", "partial stderr" + + def fake_popen(argv, **kwargs): + process = FakeProcess() + created.append((argv, kwargs, process)) + return process + + def fake_killpg(pid, signal_number): + if signal_number == 0: + raise ProcessLookupError + signals.append((pid, signal_number)) + + monkeypatch.setattr(runner.subprocess, "Popen", fake_popen) + monkeypatch.setattr(runner.os, "killpg", fake_killpg) + + try: + runner._run_lmms_eval_process( + ["python", "-m", "lmms_eval"], + cwd=str(tmp_path), + env={}, + timeout=7.0, + ) + except subprocess.TimeoutExpired as error: + assert error.timeout == 7.0 + assert error.output == "partial stdout" + assert error.stderr == "partial stderr" + else: + raise AssertionError("expected lmms-eval timeout to be raised") + + argv, kwargs, process = created[0] + assert argv == ["python", "-m", "lmms_eval"] + assert kwargs["start_new_session"] is True + assert process.communicate_timeouts == [ + 7.0, + runner._LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS, + ] + assert signals == [(1234, signal.SIGTERM)] + + +def test_lmms_eval_timeout_kills_remaining_process_group(monkeypatch, tmp_path): + signals = [] + + class FakeProcess: + pid = 3456 + returncode = None + + def communicate(self, timeout=None): + if timeout == 7.0: + raise subprocess.TimeoutExpired( + ["python", "-m", "lmms_eval"], + timeout, + output="partial stdout", + stderr="partial stderr", + ) + self.returncode = -signal.SIGTERM + return "partial stdout", "partial stderr" + + def fake_killpg(pid, signal_number): + if signal_number != 0: + signals.append((pid, signal_number)) + + monkeypatch.setattr(runner.subprocess, "Popen", lambda *args, **kwargs: FakeProcess()) + monkeypatch.setattr(runner.os, "killpg", fake_killpg) + + try: + runner._run_lmms_eval_process( + ["python", "-m", "lmms_eval"], + cwd=str(tmp_path), + env={}, + timeout=7.0, + ) + except subprocess.TimeoutExpired: + pass + else: + raise AssertionError("expected lmms-eval timeout to be raised") + + assert signals == [(3456, signal.SIGTERM), (3456, signal.SIGKILL)] + + +def test_lmms_eval_timeout_kills_stubborn_process_group(monkeypatch, tmp_path): + signals = [] + + class FakeProcess: + pid = 5678 + returncode = None + + def __init__(self): + self.communicate_timeouts = [] + + def communicate(self, timeout=None): + self.communicate_timeouts.append(timeout) + if len(self.communicate_timeouts) < 3: + raise subprocess.TimeoutExpired( + ["python", "-m", "lmms_eval"], + timeout, + output="partial stdout", + stderr="partial stderr", + ) + self.returncode = -signal.SIGKILL + return "partial stdout", "partial stderr" + + process = FakeProcess() + monkeypatch.setattr(runner.subprocess, "Popen", lambda *args, **kwargs: process) + monkeypatch.setattr( + runner.os, + "killpg", + lambda pid, signal_number: signals.append((pid, signal_number)), + ) + + try: + runner._run_lmms_eval_process( + ["python", "-m", "lmms_eval"], + cwd=str(tmp_path), + env={}, + timeout=7.0, + ) + except subprocess.TimeoutExpired as error: + assert error.output == "partial stdout" + assert error.stderr == "partial stderr" + else: + raise AssertionError("expected lmms-eval timeout to be raised") + + assert process.communicate_timeouts == [ + 7.0, + runner._LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS, + None, + ] + assert signals == [(5678, signal.SIGTERM), (5678, signal.SIGKILL)] + + def test_downstream_evaluation_runs_lmms_eval_and_flattens_metrics(monkeypatch, tmp_path): captured = {} - def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): - del env, capture_output, text, timeout, check + def fake_run(argv, *, cwd, env, timeout): + del env, timeout captured["argv"] = argv output = Path(cwd) / "nested" output.mkdir(parents=True) @@ -302,7 +454,7 @@ def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): ) return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") - monkeypatch.setattr(runner.subprocess, "run", fake_run) + monkeypatch.setattr(runner, "_run_lmms_eval_process", fake_run) node = SimpleNamespace( node_id="lmms_eval", flow_id="runtime", @@ -364,8 +516,8 @@ def test_lmms_eval_completion_validates_resolved_task_expansion(): def test_downstream_evaluation_rejects_missing_configured_task(monkeypatch, tmp_path): - def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): - del env, capture_output, text, timeout, check + def fake_run(argv, *, cwd, env, timeout): + del env, timeout output = Path(cwd) (output / "results.json").write_text( json.dumps( @@ -378,7 +530,7 @@ def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): ) return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") - monkeypatch.setattr(runner.subprocess, "run", fake_run) + monkeypatch.setattr(runner, "_run_lmms_eval_process", fake_run) node = SimpleNamespace( node_id="lmms_eval", flow_id="runtime", @@ -411,8 +563,8 @@ def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): def test_downstream_evaluation_rejects_zero_sample_task(monkeypatch, tmp_path): - def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): - del env, capture_output, text, timeout, check + def fake_run(argv, *, cwd, env, timeout): + del env, timeout output = Path(cwd) (output / "results.json").write_text( json.dumps( @@ -431,7 +583,7 @@ def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): ) return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") - monkeypatch.setattr(runner.subprocess, "run", fake_run) + monkeypatch.setattr(runner, "_run_lmms_eval_process", fake_run) node = SimpleNamespace( node_id="lmms_eval", flow_id="runtime", @@ -466,8 +618,8 @@ def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): def test_downstream_evaluation_reports_lmms_eval_output_when_results_are_missing( monkeypatch, tmp_path ): - def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): - del cwd, env, capture_output, text, timeout, check + def fake_run(argv, *, cwd, env, timeout): + del cwd, env, timeout return subprocess.CompletedProcess( argv, 0, @@ -475,7 +627,7 @@ def fake_run(argv, *, cwd, env, capture_output, text, timeout, check): stderr="", ) - monkeypatch.setattr(runner.subprocess, "run", fake_run) + monkeypatch.setattr(runner, "_run_lmms_eval_process", fake_run) node = SimpleNamespace( node_id="lmms_eval", flow_id="runtime", From 28aed98bfe879824e8046732179e06345f9e719e Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Sat, 8 Aug 2026 10:50:39 -0700 Subject: [PATCH 13/16] docs: pin lmms-eval example version Signed-off-by: Grzegorz Karch --- examples/puzzletron/README.md | 17 +++++++++++++- examples/puzzletron/ci_environment.json | 1 + .../nano_30b_a3b_bf16/runs/lmms_eval.yaml | 1 + examples/puzzletron/docs/post_mip_pipeline.md | 22 +++++++++++++++++-- examples/puzzletron/requirements.txt | 1 + noxfile.py | 2 ++ 6 files changed, 41 insertions(+), 3 deletions(-) diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 89b986cde48..2047bc81a10 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -284,6 +284,7 @@ import os from packaging.version import Version import aiperf +import lmms_eval import modelopt import nemo_automodel import torch @@ -293,7 +294,14 @@ import vllm with open(os.environ["PUZZLETRON_CI_ENVIRONMENT"], encoding="utf-8") as stream: ci_environment = json.load(stream) -for package in ("torch", "vllm", "nemo-automodel", "aiperf", "nvidia-modelopt"): +for package in ( + "torch", + "vllm", + "nemo-automodel", + "aiperf", + "lmms-eval", + "nvidia-modelopt", +): print(package, metadata.version(package)) print("torch CUDA", torch.version.cuda) @@ -306,6 +314,7 @@ assert Version(metadata.version("torchvision")).release == Version( ci_environment["torchvision"] ).release assert transformers.__version__ == ci_environment["transformers"] +assert metadata.version("lmms-eval") == ci_environment["lmms_eval"] assert Version(metadata.version("nemo-automodel")).base_version == ( ci_environment["nemo_automodel"]["base_version"] ) @@ -546,6 +555,12 @@ After `mip`, prepare one deduplicated online-evaluation plan. Repeat `--profile-id` for every configured profile; aliases ensure that an identical architecture is evaluated once while remaining visible in every profile. +Downstream `lmms-eval` nodes require the same GPU worker environment as vLLM and +the Puzzletron example requirements. The reproducible example path is pinned to +`lmms-eval==0.7.2` by `examples/puzzletron/requirements.txt` and verified +against `ci_environment.json`; do not run the checked-in `lmms_eval.yaml` +example against an unpinned evaluator install. + ```bash python examples/puzzletron/run_profile_online_evaluation.py \ --puzzle-dir "$PUZZLE_DIR" --prepare \ diff --git a/examples/puzzletron/ci_environment.json b/examples/puzzletron/ci_environment.json index d7520a69109..1e7d0f5be5f 100644 --- a/examples/puzzletron/ci_environment.json +++ b/examples/puzzletron/ci_environment.json @@ -5,6 +5,7 @@ "torch": "2.11.0", "torchvision": "0.26.0", "transformers": "5.8.1", + "lmms_eval": "0.7.2", "nemo_automodel": { "base_version": "0.5.0", "repository": "https://github.com/Separius/Automodel.git", diff --git a/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml b/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml index 4bd8c6b2954..bd3a8c492b4 100644 --- a/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml +++ b/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml @@ -5,6 +5,7 @@ defaults: - _self_ # Opt-in downstream lmms-eval workflow for the realized runtime-075 candidate. +# Install examples/puzzletron/requirements.txt to use the pinned lmms-eval version. # Non-empty post_mip.flows replaces the legacy post-MIP tail in the v2 orchestrator. zero_shot_evaluation: enabled: false diff --git a/examples/puzzletron/docs/post_mip_pipeline.md b/examples/puzzletron/docs/post_mip_pipeline.md index 0b28d0b0e99..a6263f73ea1 100644 --- a/examples/puzzletron/docs/post_mip_pipeline.md +++ b/examples/puzzletron/docs/post_mip_pipeline.md @@ -105,9 +105,11 @@ metric lists or cases. Later filters reference metrics as `mip.` or - `evaluation`: evaluates either a config-only candidate or a checkpoint and publishes all result metrics. - `aiperf`: benchmarks a checkpoint and publishes all result metrics. +- `downstream_evaluation`: runs `lmms-eval` against a materialized checkpoint + and publishes task metrics. - `global_kd`: produces a new checkpoint revision. -- `ptq` and `downstream_evaluation`: reserved interfaces; configuring either - currently fails plan compilation with a clear not-implemented error. +- `ptq`: reserved interface; configuring it currently fails plan compilation + with a clear not-implemented error. Nodes that require checkpoints never materialize implicitly. Add a `materialize` node where the transition is needed. @@ -129,6 +131,22 @@ Selection still follows `input`; `model_source` only chooses the artifact operat on. This supports a long KD run selected using short-KD/PTQ results but restarted from the original candidate. +## Downstream evaluation + +`downstream_evaluation` shells out to `python -m lmms_eval` from the GPU worker +environment. Install `examples/puzzletron/requirements.txt` in that environment; +it pins the evaluator package used by the checked-in example: + +```bash +python -m pip install -r examples/puzzletron/requirements.txt +python -c 'import importlib.metadata as m; assert m.version("lmms-eval") == "0.7.2"' +``` + +The runner derives the realized checkpoint path, vLLM topology arguments, task +list, and output path from the campaign config. Use `model_args` only for +non-derived model options such as dtype or maximum model length, and `extra_args` +only for non-reserved `lmms-eval` flags. + ## Filters `top_k` accepts one integer or separate homogeneous/heterogeneous quotas. diff --git a/examples/puzzletron/requirements.txt b/examples/puzzletron/requirements.txt index 70280023afd..478fe2dfc78 100644 --- a/examples/puzzletron/requirements.txt +++ b/examples/puzzletron/requirements.txt @@ -1,4 +1,5 @@ aiohttp>=3.9,<4 +lmms-eval==0.7.2 math-verify ray # Transformers comes from nvidia-modelopt[hf]; CPU CI overlays the exact diff --git a/noxfile.py b/noxfile.py index 86c21502a8f..c55ac725047 100644 --- a/noxfile.py +++ b/noxfile.py @@ -110,6 +110,7 @@ def puzzletron_v2(session): "torch": PUZZLETRON_V2_CI_ENVIRONMENT["torch"], "torchvision": PUZZLETRON_V2_CI_ENVIRONMENT["torchvision"], "transformers": PUZZLETRON_V2_CI_ENVIRONMENT["transformers"], + "lmms-eval": PUZZLETRON_V2_CI_ENVIRONMENT["lmms_eval"], "nemo-automodel": PUZZLETRON_V2_AUTOMODEL_SOURCE["base_version"], } session.run( @@ -125,6 +126,7 @@ def puzzletron_v2(session): "'torch': Version(version('torch')).base_version, " "'torchvision': Version(version('torchvision')).base_version, " "'transformers': Version(version('transformers')).base_version, " + "'lmms-eval': Version(version('lmms-eval')).base_version, " "'nemo-automodel': Version(version('nemo-automodel')).base_version}; " "mismatches = {name: (actual[name], expected_version) " "for name, expected_version in expected.items() " From 5c8f36b440f7667531df5d2f1c9dfdcaf7e2e0f8 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Sat, 8 Aug 2026 13:51:37 -0700 Subject: [PATCH 14/16] fix: isolate lmms-eval dependency pin Signed-off-by: Grzegorz Karch --- examples/puzzletron/README.md | 27 +++++++++++++------ .../nano_30b_a3b_bf16/runs/lmms_eval.yaml | 6 ++++- examples/puzzletron/docs/post_mip_pipeline.md | 14 +++++++--- examples/puzzletron/requirements.txt | 1 - noxfile.py | 2 -- 5 files changed, 34 insertions(+), 16 deletions(-) diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 2047bc81a10..ea77fb43039 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -284,7 +284,6 @@ import os from packaging.version import Version import aiperf -import lmms_eval import modelopt import nemo_automodel import torch @@ -299,7 +298,6 @@ for package in ( "vllm", "nemo-automodel", "aiperf", - "lmms-eval", "nvidia-modelopt", ): print(package, metadata.version(package)) @@ -314,7 +312,6 @@ assert Version(metadata.version("torchvision")).release == Version( ci_environment["torchvision"] ).release assert transformers.__version__ == ci_environment["transformers"] -assert metadata.version("lmms-eval") == ci_environment["lmms_eval"] assert Version(metadata.version("nemo-automodel")).base_version == ( ci_environment["nemo_automodel"]["base_version"] ) @@ -555,11 +552,25 @@ After `mip`, prepare one deduplicated online-evaluation plan. Repeat `--profile-id` for every configured profile; aliases ensure that an identical architecture is evaluated once while remaining visible in every profile. -Downstream `lmms-eval` nodes require the same GPU worker environment as vLLM and -the Puzzletron example requirements. The reproducible example path is pinned to -`lmms-eval==0.7.2` by `examples/puzzletron/requirements.txt` and verified -against `ci_environment.json`; do not run the checked-in `lmms_eval.yaml` -example against an unpinned evaluator install. +Downstream `lmms-eval` nodes use a separate evaluator Python. The reproducible +example path is pinned to `lmms-eval==0.7.2` by +`examples/puzzletron/requirements-lmms-eval.txt` and recorded in +`ci_environment.json`. Keep this separate from the Puzzletron runtime +environment because `lmms-eval==0.7.2` pins `wandb==0.25.0`, while the pinned +AutoModel build requires a newer `wandb`. + +```bash +python3 -m venv /workspace/.venv-lmms-eval +source /workspace/.venv-lmms-eval/bin/activate +python -m pip install --upgrade pip "setuptools>=80,<81" wheel packaging +VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cu129 \ + python -m pip install --no-build-isolation -e "${VLLM_ROOT}" +python -m pip install -r "${MODEL_OPT_ROOT}/examples/puzzletron/requirements-lmms-eval.txt" +python -c 'import importlib.metadata as m; assert m.version("lmms-eval") == "0.7.2"' +deactivate + +export PUZZLETRON_LMMS_EVAL_PYTHON=/workspace/.venv-lmms-eval/bin/python +``` ```bash python examples/puzzletron/run_profile_online_evaluation.py \ diff --git a/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml b/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml index bd3a8c492b4..a5e8d8c77b4 100644 --- a/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml +++ b/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml @@ -5,7 +5,7 @@ defaults: - _self_ # Opt-in downstream lmms-eval workflow for the realized runtime-075 candidate. -# Install examples/puzzletron/requirements.txt to use the pinned lmms-eval version. +# Use PUZZLETRON_LMMS_EVAL_PYTHON to point at the isolated pinned evaluator env. # Non-empty post_mip.flows replaces the legacy post-MIP tail in the v2 orchestrator. zero_shot_evaluation: enabled: false @@ -39,6 +39,10 @@ post_mip: type: downstream_evaluation input: materialized config: + command_prefix: + - ${oc.env:PUZZLETRON_LMMS_EVAL_PYTHON} + - -m + - lmms_eval model: vllm checkpoint_arg: model tasks: diff --git a/examples/puzzletron/docs/post_mip_pipeline.md b/examples/puzzletron/docs/post_mip_pipeline.md index a6263f73ea1..fd13f55ab38 100644 --- a/examples/puzzletron/docs/post_mip_pipeline.md +++ b/examples/puzzletron/docs/post_mip_pipeline.md @@ -133,13 +133,19 @@ from the original candidate. ## Downstream evaluation -`downstream_evaluation` shells out to `python -m lmms_eval` from the GPU worker -environment. Install `examples/puzzletron/requirements.txt` in that environment; -it pins the evaluator package used by the checked-in example: +`downstream_evaluation` shells out to `python -m lmms_eval` through +`command_prefix`. Install the pinned evaluator into an isolated environment +rather than the Puzzletron runtime environment, because `lmms-eval==0.7.2` pins +`wandb==0.25.0` and the pinned AutoModel build requires a newer `wandb`: ```bash -python -m pip install -r examples/puzzletron/requirements.txt +python3 -m venv /workspace/.venv-lmms-eval +source /workspace/.venv-lmms-eval/bin/activate +python -m pip install -r examples/puzzletron/requirements-lmms-eval.txt python -c 'import importlib.metadata as m; assert m.version("lmms-eval") == "0.7.2"' +deactivate + +export PUZZLETRON_LMMS_EVAL_PYTHON=/workspace/.venv-lmms-eval/bin/python ``` The runner derives the realized checkpoint path, vLLM topology arguments, task diff --git a/examples/puzzletron/requirements.txt b/examples/puzzletron/requirements.txt index 478fe2dfc78..70280023afd 100644 --- a/examples/puzzletron/requirements.txt +++ b/examples/puzzletron/requirements.txt @@ -1,5 +1,4 @@ aiohttp>=3.9,<4 -lmms-eval==0.7.2 math-verify ray # Transformers comes from nvidia-modelopt[hf]; CPU CI overlays the exact diff --git a/noxfile.py b/noxfile.py index c55ac725047..86c21502a8f 100644 --- a/noxfile.py +++ b/noxfile.py @@ -110,7 +110,6 @@ def puzzletron_v2(session): "torch": PUZZLETRON_V2_CI_ENVIRONMENT["torch"], "torchvision": PUZZLETRON_V2_CI_ENVIRONMENT["torchvision"], "transformers": PUZZLETRON_V2_CI_ENVIRONMENT["transformers"], - "lmms-eval": PUZZLETRON_V2_CI_ENVIRONMENT["lmms_eval"], "nemo-automodel": PUZZLETRON_V2_AUTOMODEL_SOURCE["base_version"], } session.run( @@ -126,7 +125,6 @@ def puzzletron_v2(session): "'torch': Version(version('torch')).base_version, " "'torchvision': Version(version('torchvision')).base_version, " "'transformers': Version(version('transformers')).base_version, " - "'lmms-eval': Version(version('lmms-eval')).base_version, " "'nemo-automodel': Version(version('nemo-automodel')).base_version}; " "mismatches = {name: (actual[name], expected_version) " "for name, expected_version in expected.items() " From f15c992e5336981721e9706e96962e5f821747c1 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Tue, 11 Aug 2026 10:35:22 +0200 Subject: [PATCH 15/16] Fix lmms-eval timeout typing Signed-off-by: Johannes Rausch --- modelopt/torch/puzzletron/post_mip/runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 4de690a60b8..4de011235d4 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -1047,7 +1047,7 @@ def _run_lmms_eval_process( _signal_lmms_eval_process_group(process, signal.SIGKILL) raise subprocess.TimeoutExpired( argv, - timeout, + error.timeout, output=stdout if stdout is not None else error.output, stderr=stderr if stderr is not None else error.stderr, ) from error From ebfe4e7c15d49eceb46e315a7cebcd303bd3a73c Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Tue, 11 Aug 2026 15:12:47 +0200 Subject: [PATCH 16/16] Address downstream evaluation review findings Make lmms-eval command construction deterministic and bounded, and clarify that the runner never invokes a shell. Signed-off-by: Johannes Rausch --- examples/puzzletron/docs/post_mip_pipeline.md | 6 +- modelopt/torch/puzzletron/post_mip/runner.py | 61 ++++++++----------- .../torch/puzzletron/test_post_mip_runner.py | 27 +++++--- 3 files changed, 47 insertions(+), 47 deletions(-) diff --git a/examples/puzzletron/docs/post_mip_pipeline.md b/examples/puzzletron/docs/post_mip_pipeline.md index fd13f55ab38..bd26139acf6 100644 --- a/examples/puzzletron/docs/post_mip_pipeline.md +++ b/examples/puzzletron/docs/post_mip_pipeline.md @@ -133,8 +133,10 @@ from the original candidate. ## Downstream evaluation -`downstream_evaluation` shells out to `python -m lmms_eval` through -`command_prefix`. Install the pinned evaluator into an isolated environment +`downstream_evaluation` runs `python -m lmms_eval` as a subprocess through +`command_prefix`. The runner passes an argument list directly and does not invoke +a shell. Values in `command_prefix` and `extra_args` are arguments; shell syntax +is not interpreted. Install the pinned evaluator into an isolated environment rather than the Puzzletron runtime environment, because `lmms-eval==0.7.2` pins `wandb==0.25.0` and the pinned AutoModel build requires a newer `wandb`: diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 4de011235d4..3612edeac8f 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -568,6 +568,7 @@ def _aiperf( "--tasks", } ) +_DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS = 3600.0 _LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS = 10.0 @@ -700,19 +701,15 @@ def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> "pipeline_parallel_size": canonical_topology["pp"], "data_parallel_size": canonical_topology["dp"], "enable_expert_parallel": canonical_topology["enable_expert_parallel"], - "distributed_executor_backend": canonical_topology[ - "distributed_executor_backend" - ], + "distributed_executor_backend": canonical_topology["distributed_executor_backend"], } ) - for key in _LMMS_EVAL_MODEL_ARG_FIELDS: + for key in sorted(_LMMS_EVAL_MODEL_ARG_FIELDS): if key in settings: derived[key] = settings[key] if isinstance(raw, str): - _reject_reserved_lmms_eval_model_args( - _lmms_eval_model_arg_keys(raw), reserved_fields - ) + _reject_reserved_lmms_eval_model_args(_lmms_eval_model_arg_keys(raw), reserved_fields) prefix = raw.strip().strip(",") suffix = _model_arg_string(derived) return ",".join(part for part in (prefix, suffix) if part) @@ -820,28 +817,21 @@ def _lmms_eval_command( if settings.get("cache_dir") is not None: env.setdefault("LMMS_EVAL_HOME", str(settings["cache_dir"])) timeout = settings.get("timeout_seconds", settings.get("timeout")) - return argv, env, (float(timeout) if timeout is not None else None) + if timeout is None: + timeout = _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS + return argv, env, float(timeout) def _metric_key(value: Any) -> str: return ( - str(value) - .strip() - .replace(" ", "_") - .replace(",", "_") - .replace("/", "_") - .replace("\\", "_") + str(value).strip().replace(" ", "_").replace(",", "_").replace("/", "_").replace("\\", "_") ) def _numeric_metrics(task_payload: Mapping[str, Any]) -> dict[str, float]: metrics = {} for metric_name, value in task_payload.items(): - if ( - isinstance(value, (int, float)) - and not isinstance(value, bool) - and math.isfinite(value) - ): + if isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value): metrics[str(metric_name)] = float(value) return metrics @@ -914,8 +904,7 @@ def _validate_lmms_eval_completion( missing_results = [task for task in expected_tasks if task not in results] if missing_results: raise RuntimeError( - "lmms-eval result is missing configured task results: " - f"{sorted(missing_results)}" + f"lmms-eval result is missing configured task results: {sorted(missing_results)}" ) missing_metrics = [ @@ -991,9 +980,7 @@ def _lmms_eval_output_tail(result: subprocess.CompletedProcess[str], *, max_line return "\n".join(sections) -def _signal_lmms_eval_process_group( - process: subprocess.Popen[str], signal_number: int -) -> None: +def _signal_lmms_eval_process_group(process: subprocess.Popen[str], signal_number: int) -> None: try: if os.name == "posix": os.killpg(process.pid, signal_number) @@ -1036,12 +1023,15 @@ def _run_lmms_eval_process( except subprocess.TimeoutExpired as error: _signal_lmms_eval_process_group(process, signal.SIGTERM) try: - stdout, stderr = process.communicate( - timeout=_LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS - ) + stdout, stderr = process.communicate(timeout=_LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS) except subprocess.TimeoutExpired: _signal_lmms_eval_process_group(process, signal.SIGKILL) - stdout, stderr = process.communicate() + try: + stdout, stderr = process.communicate( + timeout=_LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS + ) + except subprocess.TimeoutExpired as kill_error: + stdout, stderr = kill_error.output, kill_error.stderr else: if _lmms_eval_process_group_exists(process): _signal_lmms_eval_process_group(process, signal.SIGKILL) @@ -1097,17 +1087,14 @@ def _downstream_evaluation( if result.returncode: tail = _lmms_eval_output_tail(result) raise RuntimeError( - f"lmms-eval failed with exit code {result.returncode}" - + (f": {tail}" if tail else "") + f"lmms-eval failed with exit code {result.returncode}" + (f": {tail}" if tail else "") ) try: payload, result_path = _lmms_eval_result_payload(output) except FileNotFoundError as error: tail = _lmms_eval_output_tail(result) raise FileNotFoundError(str(error) + (f": {tail}" if tail else "")) from error - sample_counts = _validate_lmms_eval_completion( - payload, _configured_lmms_eval_tasks(settings) - ) + sample_counts = _validate_lmms_eval_completion(payload, _configured_lmms_eval_tasks(settings)) metrics = _flatten_lmms_eval_metrics(payload) if not metrics: raise RuntimeError(f"lmms-eval result has no numeric task metrics: {result_path}") @@ -1276,8 +1263,12 @@ def run_post_mip_node_shard( timeout_field = "timeout_seconds" elif not isinstance(error, subprocess.TimeoutExpired): timeout_field = "readiness_timeout" - default_timeout = 3600 if node.node_type == "downstream_evaluation" else ( - 600 if timeout_field == "benchmark_timeout" else 1200 + default_timeout = ( + _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS + if node.node_type == "downstream_evaluation" + else 600 + if timeout_field == "benchmark_timeout" + else 1200 ) row["timeout_seconds"] = float( getattr(error, "timeout", None) diff --git a/tests/unit/torch/puzzletron/test_post_mip_runner.py b/tests/unit/torch/puzzletron/test_post_mip_runner.py index 7b650284a51..4ad3c649666 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -227,6 +227,19 @@ def test_lmms_eval_command_maps_checkpoint_and_vllm_topology(tmp_path): assert timeout == 123 +def test_lmms_eval_command_uses_bounded_default_timeout(tmp_path): + _, _, timeout = runner._lmms_eval_command( + { + "tasks": ["ifeval"], + "topology": {"gpu_group_size": 1}, + }, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + + assert timeout == runner._DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS + + def test_lmms_eval_command_rejects_reserved_model_args(tmp_path): cases = ( ({"model": "/ckpts/wrong"}, "model"), @@ -424,7 +437,7 @@ def communicate(self, timeout=None): assert process.communicate_timeouts == [ 7.0, runner._LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS, - None, + runner._LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS, ] assert signals == [(5678, signal.SIGTERM), (5678, signal.SIGKILL)] @@ -550,9 +563,7 @@ def fake_run(argv, *, cwd, env, timeout): ) try: - runner._downstream_evaluation( - {"puzzle_dir": str(tmp_path)}, node, source, "execution" - ) + runner._downstream_evaluation({"puzzle_dir": str(tmp_path)}, node, source, "execution") except RuntimeError as error: message = str(error) else: @@ -603,9 +614,7 @@ def fake_run(argv, *, cwd, env, timeout): ) try: - runner._downstream_evaluation( - {"puzzle_dir": str(tmp_path)}, node, source, "execution" - ) + runner._downstream_evaluation({"puzzle_dir": str(tmp_path)}, node, source, "execution") except RuntimeError as error: message = str(error) else: @@ -648,9 +657,7 @@ def fake_run(argv, *, cwd, env, timeout): ) try: - runner._downstream_evaluation( - {"puzzle_dir": str(tmp_path)}, node, source, "execution" - ) + runner._downstream_evaluation({"puzzle_dir": str(tmp_path)}, node, source, "execution") except FileNotFoundError as error: message = str(error) else: