Add downstream eval to puzzletron v2

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/``.
 

examples/puzzletron/README.md

@@ -268,7 +268,13 @@ import vllm
 with open(os.environ["PUZZLETRON_CI_ENVIRONMENT"], encoding="utf-8") as stream:
     ci_environment = json.load(stream)
 
-for package in ("torch", "vllm", "nemo-automodel", "aiperf", "nvidia-modelopt"):
+for package in (
+    "torch",
+    "vllm",
+    "nemo-automodel",
+    "aiperf",
+    "nvidia-modelopt",
+):
     print(package, metadata.version(package))
 
 print("torch CUDA", torch.version.cuda)
@@ -501,6 +507,26 @@ After `mip`, prepare one deduplicated online-evaluation plan. Repeat
 `--profile-id` for every configured profile; aliases ensure that an identical
 architecture is evaluated once while remaining visible in every profile.
 
+Downstream `lmms-eval` nodes use a separate evaluator Python. The reproducible
+example path is pinned to `lmms-eval==0.7.2` by
+`examples/puzzletron/requirements-lmms-eval.txt` and recorded in
+`ci_environment.json`. Keep this separate from the Puzzletron runtime
+environment because `lmms-eval==0.7.2` pins `wandb==0.25.0`, while the pinned
+AutoModel build requires a newer `wandb`.
+
+```bash
+python3 -m venv /workspace/.venv-lmms-eval
+source /workspace/.venv-lmms-eval/bin/activate
+python -m pip install --upgrade pip "setuptools>=80,<81" wheel packaging
+VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cu129 \
+  python -m pip install --no-build-isolation -e "${VLLM_ROOT}"
+python -m pip install -r "${MODEL_OPT_ROOT}/examples/puzzletron/requirements-lmms-eval.txt"
+python -c 'import importlib.metadata as m; assert m.version("lmms-eval") == "0.7.2"'
+deactivate
+
+export PUZZLETRON_LMMS_EVAL_PYTHON=/workspace/.venv-lmms-eval/bin/python
+```
+
 ```bash
 python examples/puzzletron/run_profile_online_evaluation.py \
   --puzzle-dir "$PUZZLE_DIR" --prepare \

examples/puzzletron/ci_environment.json

@@ -5,6 +5,7 @@
   "torch": "2.11.0",
   "torchvision": "0.26.0",
   "transformers": "5.8.1",
+  "lmms_eval": "0.7.2",
   "nemo_automodel": {
     "base_version": "0.5.0",
     "repository": "https://github.com/Separius/Automodel.git",

examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml

@@ -0,0 +1,68 @@
+# @package _global_
+
+defaults:
+  - default
+  - _self_
+
+# Opt-in downstream lmms-eval workflow for the realized runtime-075 candidate.
+# Use PUZZLETRON_LMMS_EVAL_PYTHON to point at the isolated pinned evaluator env.
+# Non-empty post_mip.flows replaces the legacy post-MIP tail in the v2 orchestrator.
+zero_shot_evaluation:
+  enabled: false
+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:
+            command_prefix:
+              - ${oc.env:PUZZLETRON_LMMS_EVAL_PYTHON}
+              - -m
+              - lmms_eval
+            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}

examples/puzzletron/docs/post_mip_pipeline.md

@@ -105,9 +105,11 @@ metric lists or cases. Later filters reference metrics as `mip.<metric>` or
 - `evaluation`: evaluates either a config-only candidate or a checkpoint and
   publishes all result metrics.
 - `aiperf`: benchmarks a checkpoint and publishes all result metrics.
+- `downstream_evaluation`: runs `lmms-eval` against a materialized checkpoint
+  and publishes task metrics.
 - `global_kd`: produces a new checkpoint revision.
-- `ptq` and `downstream_evaluation`: reserved interfaces; configuring either
-  currently fails plan compilation with a clear not-implemented error.
+- `ptq`: reserved interface; configuring it currently fails plan compilation
+  with a clear not-implemented error.
 
 Nodes that require checkpoints never materialize implicitly. Add a `materialize`
 node where the transition is needed.
@@ -129,6 +131,28 @@ Selection still follows `input`; `model_source` only chooses the artifact operat
 on. This supports a long KD run selected using short-KD/PTQ results but restarted
 from the original candidate.
 
+## Downstream evaluation
+
+`downstream_evaluation` shells out to `python -m lmms_eval` through
+`command_prefix`. Install the pinned evaluator into an isolated environment
+rather than the Puzzletron runtime environment, because `lmms-eval==0.7.2` pins
+`wandb==0.25.0` and the pinned AutoModel build requires a newer `wandb`:
+
+```bash
+python3 -m venv /workspace/.venv-lmms-eval
+source /workspace/.venv-lmms-eval/bin/activate
+python -m pip install -r examples/puzzletron/requirements-lmms-eval.txt
+python -c 'import importlib.metadata as m; assert m.version("lmms-eval") == "0.7.2"'
+deactivate
+
+export PUZZLETRON_LMMS_EVAL_PYTHON=/workspace/.venv-lmms-eval/bin/python
+```
+
+The runner derives the realized checkpoint path, vLLM topology arguments, task
+list, and output path from the campaign config. Use `model_args` only for
+non-derived model options such as dtype or maximum model length, and `extra_args`
+only for non-reserved `lmms-eval` flags.
+
 ## Filters
 
 `top_k` accepts one integer or separate homogeneous/heterogeneous quotas.

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:

modelopt/torch/puzzletron/orchestration/compiler.py

@@ -57,6 +57,7 @@
     "aiperf": ExecutionStrategy.SHARDED,
 }
 
+
 def _mapping(value: Any) -> dict[str, Any]:
     return dict(value) if isinstance(value, Mapping) else {}
 
@@ -85,7 +86,6 @@ def _mapping(value: Any) -> dict[str, Any]:
     "downstream_evaluation": {
         "kind": "evaluator",
         "accepts": {"checkpoint"},
-        "implemented": False,
     },
 }
 
@@ -516,7 +516,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:
@@ -525,7 +525,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:

modelopt/torch/puzzletron/orchestration/controller.py

@@ -74,6 +74,29 @@ def create_executor(plan: CampaignPlan, *, local: bool = False) -> Executor:
     raise ValueError(f"Unsupported runner kind: {plan.runner.kind}")
 
 
+def _stage_dashboard_display_name(
+    config: Mapping[str, Any],
+    stage_id: str,
+    *,
+    granularity: str | None = None,
+) -> str:
+    if stage_id.startswith("post."):
+        parts = stage_id.split(".", 2)
+        if len(parts) == 3:
+            _prefix, flow_id, node_id = parts
+            node = (
+                (config.get("post_mip") or {})
+                .get("flows", {})
+                .get(flow_id, {})
+                .get("nodes", {})
+                .get(node_id, {})
+            )
+            node_type = str(node.get("type") or "") if isinstance(node, Mapping) else ""
+            if node_type == "downstream_evaluation":
+                return "Downstream Evaluation"
+    return stage_display_name(stage_id, granularity=granularity)
+
+
 @dataclass
 class DryRunSubmission:
     stage_id: str
@@ -748,7 +771,8 @@ def _stage_views(self) -> list[StageView]:
             views.append(
                 StageView(
                     stage_id=node.stage_id,
-                    display_name=stage_display_name(
+                    display_name=_stage_dashboard_display_name(
+                        self.plan.experiment_config,
                         node.stage_id,
                         granularity=str(granularity) if granularity is not None else None,
                     ),

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

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)

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(
+        "<h3>Candidate evaluation</h3>",
+        "<h3>Downstream evaluation</h3>",
+        1,
+    )
+
+
 def render_global_kd_report(section_id: str, payload: Mapping[str, Any]) -> str:
     """Render several candidate KD histories on shared, lineage-colored plots."""