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/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}
diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py
index ba46a5b8178..4f993841183 100644
--- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py
+++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py
@@ -108,10 +108,23 @@ 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
- 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..236dcdc6d6e 100644
--- a/modelopt/torch/puzzletron/orchestration/compiler.py
+++ b/modelopt/torch/puzzletron/orchestration/compiler.py
@@ -56,7 +56,17 @@
"aiperf": ExecutionStrategy.SHARDED,
}
+
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 {}
@@ -84,7 +94,6 @@ def _mapping(value: Any) -> dict[str, Any]:
"downstream_evaluation": {
"kind": "evaluator",
"accepts": {"checkpoint"},
- "implemented": False,
},
}
@@ -492,7 +501,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 +510,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..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
@@ -434,6 +445,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..30c70442598 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,18 @@ class DownstreamEvaluationNode(PostMIPNode):
NodeKind.EVALUATOR,
frozenset({ArtifactKind.CHECKPOINT}),
distributed=True,
- implemented=False,
default_strategy="sharded",
)
+
+ @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 86e3c11f44e..20320c22775 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",
]
@@ -318,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 = []
@@ -356,8 +366,27 @@ 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."""
+ """
+ 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 819cbbce010..9479c84b687 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
@@ -468,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 {})
@@ -523,11 +537,428 @@ 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:
+ """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):
+ 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:
+ """
+ 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:
+ 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:
+ """
+ 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:
+ 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:
+ """
+ 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 {})
+ 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]:
+ """
+ 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"]
+ 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 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 = [
+ *_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:
+ """
+ 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()
+ .replace(" ", "_")
+ .replace(",", "_")
+ .replace("/", "_")
+ .replace("\\", "_")
+ )
+
+
+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 {}
+ 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]:
+ """
+ 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:
+ 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]:
+ """
+ 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 {})
+ 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],
) -> 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)))
@@ -575,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)
@@ -582,6 +1029,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:
@@ -616,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))
@@ -649,7 +1110,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 +1122,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..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():
@@ -346,6 +359,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(
@@ -788,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():
@@ -816,7 +851,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..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",
@@ -285,7 +293,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..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))
@@ -3641,7 +3652,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 +3749,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 +3823,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",
@@ -4011,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"]),
@@ -4079,6 +4114,105 @@ 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:
+ """
+ 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."
+
+ 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,
@@ -4087,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 {}
)
@@ -4131,7 +4277,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..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()
@@ -728,6 +740,125 @@ 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]:
+ """
+ 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(
+ "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],
@@ -738,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}"
@@ -897,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 = []
@@ -911,7 +1071,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 +1118,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."
@@ -1106,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 = []
@@ -1159,6 +1340,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(