Add Puzzletron downstream evaluation

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/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}

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:

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:

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

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)

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(
+        "<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."""
+    """
+    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 = []