From d938fc6d90f3d4ec45c40a050bed91ab8ebd7278 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Wed, 29 Jul 2026 16:35:55 -0700 Subject: [PATCH 1/3] 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, 831 insertions(+), 22 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 63235712436..05944ed6288 100644 --- a/modelopt/torch/puzzletron/orchestration/compiler.py +++ b/modelopt/torch/puzzletron/orchestration/compiler.py @@ -56,6 +56,7 @@ "aiperf": ExecutionStrategy.SHARDED, } + def _mapping(value: Any) -> dict[str, Any]: return dict(value) if isinstance(value, Mapping) else {} @@ -84,7 +85,6 @@ def _mapping(value: Any) -> dict[str, Any]: "downstream_evaluation": { "kind": "evaluator", "accepts": {"checkpoint"}, - "implemented": False, }, } @@ -492,7 +492,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: @@ -501,7 +501,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 8a833ca4e6d..d9c45c143e6 100644 --- a/puzzletron_setup/v2/validation.py +++ b/puzzletron_setup/v2/validation.py @@ -285,7 +285,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 2b93e947935..af87270c860 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -3641,7 +3641,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", {}), @@ -3738,7 +3738,7 @@ def post_mip_section(session: WizardSession, resolver: DefaultsResolver, context "aiperf", "global_kd", ("PTQ — unavailable", "unavailable"), - ("Downstream evaluation — unavailable", "unavailable"), + "downstream_evaluation", ], default="evaluation", ) @@ -3812,6 +3812,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", @@ -4079,6 +4091,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, @@ -4131,7 +4220,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 52a9e9374f5..6938eeeee22 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -728,6 +728,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], @@ -911,7 +1019,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", ) @@ -958,9 +1066,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." @@ -1159,6 +1274,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 1a0c1c085cc..1d23353dd4e 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_compiler.py +++ b/tests/unit/torch/puzzletron/test_orchestration_compiler.py @@ -175,3 +175,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 fa2258d1f94..a56a109617b 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_state_validation.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_state_validation.py @@ -155,6 +155,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 e4824915afe29c5b22d7d9ec3839378de2fb62ed Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Wed, 29 Jul 2026 16:36:52 -0700 Subject: [PATCH 2/3] 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 495d4551b0e110a094b08fc4dad8c1f7b27ee6ea Mon Sep 17 00:00:00 2001 From: "coderabbitai[bot]" <136622811+coderabbitai[bot]@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:11:12 +0200 Subject: [PATCH 3/3] =?UTF-8?q?=F0=9F=93=9D=20Add=20docstrings=20to=20`gka?= =?UTF-8?q?rch/add=5Flmms=5Feval`=20(#2035)?= 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/2034#issuecomment-5124533292 The following files were modified: * `modelopt/torch/puzzletron/orchestration/adapters/post_mip.py` * `modelopt/torch/puzzletron/orchestration/compiler.py` * `modelopt/torch/puzzletron/orchestration/progress.py` * `modelopt/torch/puzzletron/post_mip/builtin.py` * `modelopt/torch/puzzletron/post_mip/reporting.py` * `modelopt/torch/puzzletron/post_mip/runner.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_orchestration_compiler.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`
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> --- .../orchestration/adapters/post_mip.py | 13 ++ .../puzzletron/orchestration/compiler.py | 9 + .../puzzletron/orchestration/progress.py | 11 ++ modelopt/torch/puzzletron/post_mip/builtin.py | 9 + .../torch/puzzletron/post_mip/reporting.py | 22 ++- modelopt/torch/puzzletron/post_mip/runner.py | 175 +++++++++++++++++- puzzletron_setup/bundle.py | 28 +++ puzzletron_setup/v2/validation.py | 10 +- puzzletron_setup/v2/wizard.py | 63 ++++++- puzzletron_setup/wizard.py | 70 ++++++- 10 files changed, 400 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index 90f1f745f9a..4f993841183 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: + """ + Create the work plan for a post-MIP node. + + Parameters: + plan (CampaignPlan): Campaign execution plan containing node configuration and artifacts. + node (StagePlanNode): Post-MIP node whose workers should be planned. + + Returns: + WorkPlan: A work plan with the required worker allocation and aggregation enabled. + + Raises: + RuntimeError: If an evaluation node has no candidate architectures available. + """ 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 05944ed6288..236dcdc6d6e 100644 --- a/modelopt/torch/puzzletron/orchestration/compiler.py +++ b/modelopt/torch/puzzletron/orchestration/compiler.py @@ -58,6 +58,15 @@ def _mapping(value: Any) -> dict[str, Any]: + """ + Convert a mapping to a dictionary. + + Parameters: + value (Any): Value to normalize. + + Returns: + dict[str, Any]: A dictionary copy of the value, or an empty dictionary if the value is not a mapping. + """ return dict(value) if isinstance(value, Mapping) else {} diff --git a/modelopt/torch/puzzletron/orchestration/progress.py b/modelopt/torch/puzzletron/orchestration/progress.py index a9c943357e6..05484b49b70 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): Post-MIP stage identifier containing the flow and node IDs. + config (Mapping[str, Any] | None): Configuration containing post-MIP flow definitions. + + Returns: + str | None: A candidate completion summary with failure or timeout counts, or `None` when progress data is unavailable. + """ parts = stage_id.split(".", 2) if len(parts) != 3: return None diff --git a/modelopt/torch/puzzletron/post_mip/builtin.py b/modelopt/torch/puzzletron/post_mip/builtin.py index 050352af445..30c70442598 100644 --- a/modelopt/torch/puzzletron/post_mip/builtin.py +++ b/modelopt/torch/puzzletron/post_mip/builtin.py @@ -125,4 +125,13 @@ class DownstreamEvaluationNode(PostMIPNode): @classmethod def render_report(cls, node, payload): + """Render the downstream evaluation report for a node. + + Parameters: + node: The node associated with the report. + payload: Report data containing the section identifier. + + Returns: + str: The rendered downstream evaluation report. + """ 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..20320c22775 100644 --- a/modelopt/torch/puzzletron/post_mip/reporting.py +++ b/modelopt/torch/puzzletron/post_mip/reporting.py @@ -319,7 +319,16 @@ 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 an HTML report of AIPerf candidate throughput, latency, and selection results. + + Parameters: + section_id (str): Identifier used to generate the report's chart element ID. + payload (Mapping[str, Any]): Report data containing observations and status information. + + Returns: + str: HTML containing a status summary, throughput chart, and candidate results table. + """ observations = list(payload.get("observations") or ()) rows = [] @@ -368,7 +377,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 a shared comparison of candidate short KD runs and their loss histories. + + Parameters: + section_id (str): Identifier used to construct the chart element IDs. + payload (Mapping[str, Any]): Report data containing candidate KD runs and status information. + + Returns: + str: HTML containing the comparison heading, status summary, loss chart containers, 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 3a1ee61afab..9479c84b687 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -470,6 +470,18 @@ def _evaluate( def _aiperf( config: dict[str, Any], node: CompiledPostMIPNode, source, execution_identity: str ) -> dict[str, Any]: + """ + Run an AIPerf benchmark for a candidate checkpoint across configured concurrency levels. + + Parameters: + config (dict[str, Any]): Runtime configuration used to determine the artifact output directory. + node (CompiledPostMIPNode): Node configuration containing benchmark settings. + source: Candidate source containing the checkpoint to benchmark. + execution_identity (str): Identity of the current node execution. + + Returns: + dict[str, Any]: A mapping containing aggregated benchmark metrics and raw artifact paths. + """ from ..benchmarks import run_aiperf_sweep settings = dict(node.config.get("config") or {}) @@ -541,10 +553,27 @@ def _aiperf( def _as_cli_bool(value: bool) -> str: + """Convert a boolean value to its CLI string representation. + + Parameters: + value (bool): The boolean value to convert. + + Returns: + str: "True" for a true value or "False" for a false value. + """ return "True" if value else "False" def _as_lmms_eval_arg(value: Any) -> str: + """ + Convert a supported value to its deterministic LMMS-Eval argument representation. + + Parameters: + value (Any): The value to convert. + + Returns: + str: A string representation suitable for an LMMS-Eval argument. + """ if isinstance(value, bool): return _as_cli_bool(value) if isinstance(value, (int, float)) and not isinstance(value, bool): @@ -555,6 +584,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 string. + + Parameters: + value (Any): String or sequence of values to join. + path (str): Configuration path used in validation error messages. + + Returns: + str: The stripped input string or comma-separated sequence values. + + Raises: + TypeError: If value is neither a string nor a sequence. + ValueError: If value or any sequence item is empty. + """ if isinstance(value, str): text = value.strip() if not text: @@ -569,6 +612,18 @@ def _join_cli_values(value: Any, *, path: str) -> str: def _model_arg_string(values: Mapping[str, Any]) -> str: + """ + Build a comma-separated lmms-eval model arguments string. + + Parameters: + values (Mapping[str, Any]): Model argument names and values to serialize. + + Returns: + str: The serialized model 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 +644,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 configured values into an LMMS-Eval model-arguments string. + + Parameters: + settings (Mapping[str, Any]): LMMS-Eval settings containing optional model arguments and topology. + checkpoint (str): Path to the checkpoint supplied to the evaluator. + + 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 +690,18 @@ def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> def _command_prefix(settings: Mapping[str, Any]) -> list[str]: + """ + Build 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 arguments, or the default Python module invocation. + + Raises: + ValueError: If the configured command prefix is empty. + """ raw = settings.get("command_prefix") if raw is None: return [sys.executable, "-m", "lmms_eval"] @@ -643,7 +720,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.""" + """ + Build the LMMS-Eval command, environment, and optional timeout for a checkpoint evaluation. + + Parameters: + settings (Mapping[str, Any]): Evaluation settings used to construct the invocation. + checkpoint (str): Path to the checkpoint to evaluate. + output_path (Path): Directory where LMMS-Eval writes its outputs. + + 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 +784,15 @@ def _lmms_eval_command( def _metric_key(value: Any) -> str: + """ + Normalize a metric name for use as a stable key. + + Parameters: + value (Any): Value to convert into a normalized metric key. + + Returns: + str: The stripped string representation with spaces, commas, slashes, and backslashes replaced by underscores. + """ return ( str(value) .strip() @@ -708,6 +804,15 @@ def _metric_key(value: Any) -> str: def _flatten_lmms_eval_metrics(payload: Mapping[str, Any]) -> dict[str, float]: + """ + Flatten finite numeric LMMS-Eval task metrics into normalized metric names. + + Parameters: + payload (Mapping[str, Any]): LMMS-Eval output containing task results. + + Returns: + dict[str, float]: Numeric metrics keyed by normalized task and metric names. + """ results = payload.get("results") if not isinstance(results, Mapping): return {} @@ -726,6 +831,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 results payload under a directory. + + Parameters: + output_path (Path): Directory containing LMMS-Eval output files. + + Returns: + tuple[dict[str, Any], Path]: The results payload and the path of its JSON file. + + Raises: + FileNotFoundError: If no JSON file containing a mapping-valued ``results`` field is found. + """ candidates = [] for path in sorted(output_path.rglob("*.json")): try: @@ -746,6 +863,22 @@ def _downstream_evaluation( source, execution_identity: str, ) -> dict[str, Any]: + """ + Run LMMS-Eval on a materialized checkpoint and collect its numeric task metrics. + + Parameters: + config (dict[str, Any]): Runtime configuration used to determine output locations. + node (CompiledPostMIPNode): Compiled downstream evaluation node configuration. + source: Candidate source containing the checkpoint to evaluate. + execution_identity (str): Identity of the current node execution. + + Returns: + dict[str, Any]: Evaluation metrics and paths to the summary, raw result, and recorded command. + + Raises: + ValueError: If the source does not contain a checkpoint artifact. + RuntimeError: If LMMS-Eval fails or produces no numeric task metrics. + """ if source.artifact_kind is not ArtifactKind.CHECKPOINT: raise ValueError("downstream_evaluation requires materialized checkpoint artifacts") settings = dict(node.config.get("config") or {}) @@ -815,7 +948,17 @@ def _post_mip_kd_settings( config: Mapping[str, Any], node_settings: Mapping[str, Any], ) -> dict[str, Any]: - """Resolve KD settings while preserving the transformer's checkpoint contract.""" + """ + Resolve global distillation settings with node-specific overrides. + + Parameters: + config (Mapping[str, Any]): Configuration containing optional global distillation settings. + node_settings (Mapping[str, Any]): Node-specific distillation settings that override global values. + + Returns: + dict[str, Any]: The merged distillation settings, with disabled string values for + ``save_consolidated`` normalized to ``True``. + """ settings = copy.deepcopy(dict(config.get("global_distillation") or {})) settings.update(copy.deepcopy(dict(node_settings))) @@ -863,6 +1006,22 @@ def _run_candidate( input_revision_id: str, execution_identity: str, ) -> dict[str, Any]: + """ + Execute a candidate through the configured post-MIP node and assemble its result. + + Parameters: + config (dict[str, Any]): Pipeline configuration. + node (CompiledPostMIPNode): Compiled node defining the candidate operation. + ledger (CandidateLedger): Ledger containing candidate revisions. + input_revision_id (str): Revision identifier to process. + execution_identity (str): Identity of the current node execution. + + Returns: + dict[str, Any]: Successful result containing input, source, architecture, and operation-specific fields. + + 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) @@ -906,6 +1065,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 a shard of a compiled post-MIP node and persist its per-candidate results. + + Parameters: + config (dict[str, Any]): Runtime configuration for the post-MIP execution. + 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 candidates. + + Returns: + Path: Path to the shard results JSON file. + """ node = _compiled_node(config, stage_id) ledger = _ledger(config) ledger.ingest_mip(_puzzle_dir(config)) diff --git a/puzzletron_setup/bundle.py b/puzzletron_setup/bundle.py index 3ee3e8f0936..c11bbba398e 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 for the selected execution budget. + + Parameters: + state (Mapping[str, Any]): Setup state containing post-MIP flow definitions. + smoke (bool): Whether to apply smoke-budget limits. + 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]: Deep-copied post-MIP flows with normalized node configurations and budget-specific limits. + """ 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]: + """ + Build execution stage entries for post-MIP flow nodes. + + 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 allocated per node. + common (Mapping[str, Any]): Parallelization settings for evaluation stages. + single_gpu (Mapping[str, Any]): Parallelization settings for materialization stages. + cpu_partition (str | None): Optional CPU partition for selector and materialization stages. + pool_source_evaluations (bool): Whether source evaluations use a persistent worker pool. + + Returns: + dict[str, Any]: Mapping of post-MIP stage identifiers to their execution configurations. + """ 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 d9c45c143e6..a0dc43e6680 100644 --- a/puzzletron_setup/v2/validation.py +++ b/puzzletron_setup/v2/validation.py @@ -186,7 +186,15 @@ def _dataset_subset_issues(state: WizardState) -> list[ValidationIssue]: def validate_state(state: WizardState) -> tuple[ValidationIssue, ...]: - """Return actionable authoring issues before canonical compilation.""" + """ + Validate persisted wizard state and return actionable cross-field authoring issues. + + Parameters: + state (WizardState): Persisted wizard state to validate. + + Returns: + tuple[ValidationIssue, ...]: Issues sorted by validation category and path. + """ issues: list[ValidationIssue] = [] required = ( "model.source", diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index af87270c860..b9fcabc7066 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -3597,6 +3597,17 @@ def _post_mip_strategy(node: NodeDraft) -> str: def post_mip_section(session: WizardSession, resolver: DefaultsResolver, context: dict) -> bool: + """ + Configure post-MIP flows for each configured MIP run. + + Parameters: + session (WizardSession): Interactive wizard session used to read and store configuration. + resolver (DefaultsResolver): Resolver for stage resource defaults. + context (dict): Wizard context containing the inspected model and its inventory. + + Returns: + bool: `True` when post-MIP flows are configured, `False` if the user backs out. + """ mip = _mapping_copy(session.state.collection("mip_config")) runs = _mapping_copy(mip.get("runs")) sequence = int(session.state.get_field("data.sequence_length", 4096)) @@ -4023,7 +4034,19 @@ def _serving_setting_prompt( pruning: Mapping[str, Any], stage_id: str, ) -> Any: - """Ask the complete AIPerf workload and serving-only parallel setting.""" + """ + Collect serving workload parameters and its vLLM parallel topology. + + Parameters: + prefix (str): State path prefix used for the interactive prompts. + defaults (Mapping[str, Any]): Default serving settings. + inventory (Any): Model inventory used to validate the topology. + pruning (Mapping[str, Any]): Pruning configuration used to validate the topology. + stage_id (str): Stage associated with the serving topology. + + Returns: + Any: The configured serving settings, or `BACK` if the user navigates backward. + """ values = {} for name, label, default in ( ("input_tokens", "Serving input sequence length (ISL):", defaults["input_tokens"]), @@ -4100,9 +4123,31 @@ 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.""" + """ + Prompt for lmms-eval tasks and the vLLM configuration used to run them. + + Parameters: + session (WizardSession): Interactive wizard session. + prefix (str): State path prefix for the prompted settings. + defaults (Mapping[str, Any]): Default downstream evaluation settings. + inventory (Any): Model inventory used to validate the vLLM topology. + pruning (Mapping[str, Any]): Pruning configuration used to validate the vLLM topology. + stage_id (str): Stage identifier associated with the vLLM topology. + + Returns: + dict: Configured lmms-eval settings, including tasks, sampling limits, timeout, vLLM topology, and model arguments; `BACK` if the user backs out. + """ def validate_tasks(value: str) -> bool | str: + """ + Validate a comma-separated list of lmms-eval tasks. + + Parameters: + value (str): The task list to validate. + + 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." @@ -4176,7 +4221,19 @@ def _configure_dynamic_resources( *, ask: bool, ) -> Any: - """Attach an independent resource/batch card to every node in one flow.""" + """ + Configure independent resources for each node in a post-MIP flow. + + Parameters: + session (WizardSession): Interactive session used to read state and collect resource settings. + editor (PostMIPFlowEditor): Flow editor containing the nodes to configure. + flow_id (str): Identifier of the post-MIP flow. + model (Any): Model information used to validate parallel profiles. + ask (bool): Whether to prompt for per-node resource and batch customization. + + Returns: + True when configuration completes, or BACK if the user exits during a prompt. + """ registry = ResourceProfileRegistry.from_dict( session.state.collection("parallel_profiles") or {} ) diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index 6938eeeee22..79e05c60372 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -635,7 +635,19 @@ 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 topology and workload. + + Parameters: + prompts (PromptSession): Interactive prompt session. + detailed (bool): Whether to collect detailed workload settings. + 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): Existing topology and workload values to use as defaults. + + Returns: + dict[str, Any]: AIPerf workload configuration containing serving topology, token counts, concurrency, and timeout settings. + """ defaults = dict(defaults or {}) topology_defaults = dict(defaults.get("topology") or {}) checkpoint = prompts.checkpoint() @@ -735,7 +747,18 @@ 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: + prompts (PromptSession): Interactive prompt session. + detailed (bool): Whether to prompt for the per-candidate timeout. + moe (bool): Whether to offer expert parallelism for the vLLM topology. + defaults (Mapping[str, Any] | None): Previously saved configuration values to use as prompt defaults. + + Returns: + dict[str, Any]: Downstream evaluation configuration containing tasks, sampling settings, vLLM topology, model arguments, and timeout. + """ defaults = defaults or {} tasks = prompts.text( @@ -846,6 +869,21 @@ def _default_flow( objective: Mapping[str, Any] | None = None, include_initial_filter: bool = True, ) -> dict[str, Any]: + """ + Build an opinionated post-MIP evaluation and selection flow. + + Parameters: + run_id (str): Identifier of the MIP run used as the flow source. + run (Mapping[str, Any]): MIP run configuration, including objectives and homogeneous candidate settings. + runtime (Mapping[str, Any]): Runtime workload settings for the serving evaluation. + data (Mapping[str, Any]): Dataset settings used to configure evaluation block size. + prefix (str): Prefix applied to generated node identifiers. + objective (Mapping[str, Any] | None): Objective to use for the flow; defaults to the first objective in the run. + include_initial_filter (bool): Whether to include the initial MIP-score filtering node. + + Returns: + dict[str, Any]: Flow definition containing its source specification and ordered orchestration nodes. + """ def node_id(name: str) -> str: return f"{prefix}{name}" @@ -1005,6 +1043,20 @@ def _custom_flow( detailed: bool, moe: bool, ) -> dict[str, Any]: + """ + Interactively build a post-MIP orchestration flow for a run. + + Parameters: + run_id (str): Identifier of the MIP run associated with the flow. + runtime (Mapping[str, Any]): Runtime settings used to configure serving evaluations. + data (Mapping[str, Any]): Data settings used to configure model evaluations. + used_ids (set[str]): Node identifiers already used by other flows. + detailed (bool): Whether to collect detailed configuration options. + moe (bool): Whether the selected model is a mixture-of-experts model. + + Returns: + dict[str, Any]: A flow definition containing the run source and configured post-MIP nodes. + """ nodes: OrderedDict[str, Any] = OrderedDict() available_metrics = ["mip.score"] transformer_nodes = [] @@ -1221,6 +1273,20 @@ def _resource_rows( gpus_per_node: int, workers: Mapping[str, int], ) -> list[dict[str, Any]]: + """ + Builds a resource plan for each campaign execution stage. + + Parameters: + state (AnswerState): Campaign state containing post-MIP flows and execution mode. + common (Mapping[str, int]): Common parallel mesh dimensions. + bypass (Mapping[str, int]): Bypass parallel mesh dimensions. + global_kd (Mapping[str, int]): Global knowledge-distillation mesh dimensions. + gpus_per_node (int): Number of GPUs available on each node. + workers (Mapping[str, int]): Worker limits for pool and sharded workloads. + + Returns: + list[dict[str, Any]]: Resource rows containing each stage name, instance count, GPUs per instance, and required node count. + """ from .bundle import _post_mip_candidate_limits, _serving_parallel rows = []