[DRAFT] Make distributed Puzzletron campaigns reliable

examples/puzzletron/distributed_eval/run_coordinator.sh

@@ -18,7 +18,8 @@ set -Eeuo pipefail
 
 : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}"
 : "${CONFIG_PATH:?set CONFIG_PATH}"
-: "${WORLD_SIZE:?set WORLD_SIZE to one torchrun worker-group world size}"
+: "${PUZZLETRON_WORKER_WORLD_SIZE:=${WORLD_SIZE:-}}"
+: "${PUZZLETRON_WORKER_WORLD_SIZE:?set PUZZLETRON_WORKER_WORLD_SIZE to one worker-group world size}"
 : "${SOLUTIONS_PATH:?set SOLUTIONS_PATH}"
 : "${OUTPUT_DIR:?set OUTPUT_DIR}"
 
@@ -57,7 +58,7 @@ if [[ ! -f "${CAMPAIGN_DIR}/manifest.json" ]]; then
   "${PYTHON_BIN}" -m modelopt.torch.puzzletron.distributed_eval.cli init \
     --campaign-dir "${CAMPAIGN_DIR}" \
     --config "${CONFIG_PATH}" \
-    --world-size "${WORLD_SIZE}" \
+    --world-size "${PUZZLETRON_WORKER_WORLD_SIZE}" \
     --evaluator-revision "${EVALUATOR_REVISION}" \
     "${override_args[@]}"
 fi

examples/puzzletron/distributed_eval/run_depth_coordinator.sh

@@ -1,9 +1,25 @@
 #!/usr/bin/env bash
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 set -Eeuo pipefail
 
 : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}"
 : "${CONFIG_PATH:?set CONFIG_PATH}"
-: "${WORLD_SIZE:?set WORLD_SIZE to one torchrun worker-group world size}"
+: "${PUZZLETRON_WORKER_WORLD_SIZE:=${WORLD_SIZE:-}}"
+: "${PUZZLETRON_WORKER_WORLD_SIZE:?set PUZZLETRON_WORKER_WORLD_SIZE to one worker-group world size}"
 
 PYTHON_BIN="${PYTHON_BIN:-python}"
 
@@ -18,7 +34,7 @@ if [[ ! -f "${CAMPAIGN_DIR}/manifest.json" ]]; then
   "${PYTHON_BIN}" -m modelopt.torch.puzzletron.distributed_eval.cli init \
     --campaign-dir "${CAMPAIGN_DIR}" \
     --config "${CONFIG_PATH}" \
-    --world-size "${WORLD_SIZE}" \
+    --world-size "${PUZZLETRON_WORKER_WORLD_SIZE}" \
     --stage depth \
     --evaluator-revision "${EVALUATOR_REVISION:-puzzletron-depth-v1}" \
     "${override_args[@]}"

examples/puzzletron/distributed_eval/run_depth_pool.sh

@@ -18,7 +18,8 @@ set -Eeuo pipefail
 
 : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}"
 : "${CONFIG_PATH:?set CONFIG_PATH}"
-: "${WORLD_SIZE:?set WORLD_SIZE to one worker-group world size}"
+: "${PUZZLETRON_WORKER_WORLD_SIZE:=${WORLD_SIZE:-}}"
+: "${PUZZLETRON_WORKER_WORLD_SIZE:?set PUZZLETRON_WORKER_WORLD_SIZE to one worker-group world size}"
 : "${WORKER_COUNT:?set WORKER_COUNT to the number of worker groups}"
 : "${PUZZLETRON_GROUP_INDEX:=${PUZZLETRON_TASK_INDEX:-${SLURM_PROCID:-}}}"
 : "${PUZZLETRON_GROUP_INDEX:?run this script as one orchestrator worker-group task}"
@@ -69,7 +70,7 @@ if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" && ! -f "${MANIFEST_PATH
     -m modelopt.torch.puzzletron.distributed_eval.cli init \
     --campaign-dir "${CAMPAIGN_DIR}" \
     --config "${CONFIG_PATH}" \
-    --world-size "${WORLD_SIZE}" \
+    --world-size "${PUZZLETRON_WORKER_WORLD_SIZE}" \
     --stage depth \
     --evaluator-revision "${EVALUATOR_REVISION:-puzzletron-depth-v1}" \
     "${override_args[@]}"
@@ -86,7 +87,7 @@ done
 
 # Every scheduler task owns one GPU slice and starts one worker group. Multiple
 # independent worker groups may share a node.
-export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORLD_SIZE}}"
+export NPROC_PER_NODE="${NPROC_PER_NODE:-${PUZZLETRON_WORKER_WORLD_SIZE}}"
 export WORKER_GROUP_INDEX="${GROUP_INDEX}"
 export WORKER_ID="${WORKER_PREFIX}${GROUP_INDEX}"
 export WORKER_HOST="${WORKER_HOST:-$(hostname -f)}"
@@ -175,6 +176,13 @@ PY
   wait "${coordinator_pid}"
   coordinator_rc=$?
   set -e
+  if ((coordinator_rc == 0)); then
+    CUDA_VISIBLE_DEVICES="" "${PYTHON_BIN}" \
+      "${SCRIPT_DIR}/../finalize_depth_importance.py" \
+      --config "${CONFIG_PATH}" \
+      --puzzle-dir "${PUZZLE_DIR}" \
+      --output-dir "${OUTPUT_DIR}"
+  fi
   drain_workers
 fi
 

examples/puzzletron/distributed_eval/run_replacement_pool.sh

@@ -18,7 +18,8 @@ set -Eeuo pipefail
 
 : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}"
 : "${CONFIG_PATH:?set CONFIG_PATH}"
-: "${WORLD_SIZE:?set WORLD_SIZE to one worker-group world size}"
+: "${PUZZLETRON_WORKER_WORLD_SIZE:=${WORLD_SIZE:-}}"
+: "${PUZZLETRON_WORKER_WORLD_SIZE:?set PUZZLETRON_WORKER_WORLD_SIZE to one worker-group world size}"
 : "${WORKER_COUNT:?set WORKER_COUNT to the number of worker groups}"
 : "${PUZZLETRON_GROUP_INDEX:=${PUZZLETRON_TASK_INDEX:-${SLURM_PROCID:-}}}"
 : "${PUZZLETRON_GROUP_INDEX:?run this script as one orchestrator worker-group task}"
@@ -68,7 +69,7 @@ if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" && ! -f "${MANIFEST_PATH
     -m modelopt.torch.puzzletron.distributed_eval.cli init \
     --campaign-dir "${CAMPAIGN_DIR}" \
     --config "${CONFIG_PATH}" \
-    --world-size "${WORLD_SIZE}" \
+    --world-size "${PUZZLETRON_WORKER_WORLD_SIZE}" \
     --stage replace_block \
     --evaluator-revision "${EVALUATOR_REVISION:-puzzletron-distributed-replace-block-v1}" \
     "${override_args[@]}"
@@ -83,7 +84,7 @@ while [[ ! -f "${MANIFEST_PATH}" ]]; do
   sleep 1
 done
 
-export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORLD_SIZE}}"
+export NPROC_PER_NODE="${NPROC_PER_NODE:-${PUZZLETRON_WORKER_WORLD_SIZE}}"
 export WORKER_GROUP_INDEX="${GROUP_INDEX}"
 export WORKER_ID="${WORKER_PREFIX}${GROUP_INDEX}"
 export WORKER_HOST="${WORKER_HOST:-$(hostname -f)}"

examples/puzzletron/finalize_depth_importance.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Publish a terminal stage manifest for distributed depth importance."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+from modelopt.torch.puzzletron.manifest import stage_manifest_from_config, write_stage_manifest
+from modelopt.torch.puzzletron.pipeline_config import pipeline_config_from_path
+
+__all__ = ["finalize_depth_importance", "main"]
+
+
+def _validated_trajectory(path: Path, *, expected_removals: int) -> dict[str, Any]:
+    try:
+        payload = json.loads(path.read_text())
+    except (OSError, ValueError) as error:
+        raise RuntimeError(f"depth trajectory is unavailable or invalid: {path}") from error
+    selected = payload.get("selected")
+    if (
+        payload.get("status") != "complete"
+        or int(payload.get("max_removals", -1)) != expected_removals
+        or not isinstance(selected, list)
+        or len(selected) != expected_removals
+    ):
+        raise RuntimeError(
+            f"depth trajectory is incomplete: expected {expected_removals} removals in {path}"
+        )
+    return payload
+
+
+def finalize_depth_importance(
+    config_path: str | Path,
+    puzzle_dir: str | Path,
+    output_dir: str | Path,
+    *,
+    overrides: list[str] | None = None,
+) -> dict[str, Any]:
+    """Validate the durable trajectory and publish the canonical stage manifest."""
+
+    config = pipeline_config_from_path(config_path, overrides=overrides)
+    puzzle_dir = Path(puzzle_dir)
+    config["puzzle_dir"] = str(puzzle_dir)
+    config.setdefault("experiment", {})["dir"] = str(puzzle_dir)
+    depth = config.get("depth_importance") or {}
+    expected_removals = int(depth.get("max_removals", depth.get("max_subblocks_to_remove", 10)))
+    trajectory_path = Path(output_dir) / "trajectory.json"
+    trajectory = _validated_trajectory(
+        trajectory_path,
+        expected_removals=expected_removals,
+    )
+
+    outputs = {
+        "trajectory_path": str(trajectory_path),
+        "scenario_count": len(trajectory.get("scenarios") or ()),
+        "selected": trajectory["selected"],
+    }
+    manifest = stage_manifest_from_config("depth_importance", config)
+    manifest.complete(outputs=outputs)
+    write_stage_manifest(
+        puzzle_dir / "manifests" / "depth_importance.json",
+        manifest,
+    )
+    return outputs
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True)
+    parser.add_argument("--puzzle-dir", required=True)
+    parser.add_argument("--output-dir", required=True)
+    args = parser.parse_args()
+
+    overrides = [
+        override
+        for override in os.environ.get("DISTRIBUTED_EVAL_OVERRIDES", "").splitlines()
+        if override
+    ]
+    finalize_depth_importance(
+        args.config,
+        args.puzzle_dir,
+        args.output_dir,
+        overrides=overrides,
+    )
+
+
+if __name__ == "__main__":
+    main()

modelopt/torch/puzzletron/distillation/global_kd_recipe.py

@@ -302,14 +302,19 @@ def _align_dtensor_to_module_mesh(value, module):
         return value
     target_mesh = weight.device_mesh
     if not isinstance(value, DTensor):
+        target_names = tuple(getattr(target_mesh, "mesh_dim_names", ()) or ())
+        # FSDP-only heads accept rank-local activations and gather their sharded
+        # parameters internally. Treating distinct DP batches as replicated
+        # DTensors would give them incorrect global semantics.
+        if "tp" not in target_names:
+            return value
         # Hidden states captured from TP modules can be materialized as local
         # tensors even though the saved LM-head projection owns DTensor
         # parameters. The head's colwise TP contract consumes replicated hidden
         # states, so restore that layout without communication before dispatch.
-        target_names = tuple(getattr(target_mesh, "mesh_dim_names", ()) or ())
         if target_names != ("tp",):
             raise RuntimeError(
-                "Cannot infer replicated hidden-state placement for a non-TP-only "
+                "Cannot infer replicated hidden-state placement for a mixed "
                 f"LM-head mesh: names={target_names}"
             )
         if int(value.shape[-1]) != int(weight.shape[-1]):
@@ -355,6 +360,18 @@ def _project_teacher_hidden_on_reference_mesh(hidden, teacher_head, reference_lo
     projection = getattr(teacher_head, "_puzzletron_projection_forward", None)
     if not isinstance(weight, DTensor):
         return projection(hidden) if projection is not None else teacher_head(hidden)
+    mesh_names = tuple(getattr(weight.device_mesh, "mesh_dim_names", ()) or ())
+    if "tp" not in mesh_names:
+        if projection is None:
+            return teacher_head(hidden)
+        passthrough = teacher_head.forward
+        teacher_head.forward = projection
+        try:
+            # Invoke Module.__call__ so composable FSDP can all-gather the
+            # sharded LM-head parameters before the saved projection runs.
+            return teacher_head(hidden)
+        finally:
+            teacher_head.forward = passthrough
 
     # TP may shard the captured hidden width. Preserve the DTensor so operator
     # dispatch retains its global shape and placements; projecting local shards

modelopt/torch/puzzletron/orchestration/adapters/pool.py

@@ -291,7 +291,7 @@ def command(
                 "CAMPAIGN_DIR": str(campaign_dir),
                 "CONFIG_PATH": plan.experiment_config_path,
                 "PUZZLE_DIR": str(replacement_puzzle_dir),
-                "WORLD_SIZE": str(node.gpus_per_instance),
+                "PUZZLETRON_WORKER_WORLD_SIZE": str(node.gpus_per_instance),
                 "NPROC_PER_NODE": str(topology.gpus_per_task),
                 "WORKER_COUNT": str(worker_count),
             }
@@ -365,7 +365,7 @@ def command(
             "CAMPAIGN_DIR": str(campaign_dir),
             "CONFIG_PATH": plan.experiment_config_path,
             "PUZZLE_DIR": str(replacement_puzzle_dir),
-            "WORLD_SIZE": str(node.gpus_per_instance),
+            "PUZZLETRON_WORKER_WORLD_SIZE": str(node.gpus_per_instance),
             "WORKER_ID": str(item.metadata.get("worker_id", 0)),
             "WORKER_COUNT": str(node.instances),
         }

modelopt/torch/puzzletron/orchestration/controller.py

@@ -1225,6 +1225,11 @@ def _stage_views(self) -> list[StageView]:
                 for handle_id, entry in self._active.items()
                 if entry[1].startswith(f"{node.stage_id}:")
             ]
+            stage_logs = [
+                path
+                for _handle_id, (_handle, work_id, _attempt_id) in stage_entries
+                for path in active_logs.get(work_id, ())
+            ]
             active_states = [
                 self._last_states.get(handle_id, JobState.UNKNOWN)
                 for handle_id, _entry in stage_entries
@@ -1255,11 +1260,6 @@ def _stage_views(self) -> list[StageView]:
                     if any(state is JobState.RUNNING for state in active_states)
                     else "pending"
                 )
-                stage_logs = [
-                    path
-                    for _handle_id, (_handle, work_id, _attempt_id) in stage_entries
-                    for path in active_logs.get(work_id, ())
-                ]
                 progress = summarize_stage_artifacts(
                     self.plan.puzzle_dir,
                     node.stage_id,

puzzletron_setup/v2/bundle.py

@@ -334,6 +334,24 @@ def render_runner_v2(state: WizardState, budget: str) -> dict[str, Any]:
     return _render_runner_v2(resolve_campaign_config(state), budget)
 
 
+def _stage_instances(
+    stage_id: str,
+    strategy: str,
+    configured_instances: int,
+    experiment: Mapping[str, Any],
+    budget: str,
+) -> int:
+    instances = 1 if budget == "smoke" else configured_instances
+    embedding = _mapping(experiment.get("embedding_pruning"))
+    if (
+        stage_id == "replacement_scoring"
+        and strategy == "persistent_pool"
+        and embedding.get("enabled")
+    ):
+        instances = max(instances, len(embedding.get("widths") or ()))
+    return instances
+
+
 def _render_execution_v2(
     config: ResolvedCampaignConfig,
     budget: str,
@@ -347,7 +365,13 @@ def _render_execution_v2(
     for stage_id, resource in config.stage_resources.items():
         entry = {
             "strategy": resource.strategy,
-            "instances": min(resource.instances, 1) if budget == "smoke" else resource.instances,
+            "instances": _stage_instances(
+                stage_id,
+                resource.strategy,
+                resource.instances,
+                experiment,
+                budget,
+            ),
             "resource": resource.resource,
             "gpus_per_node": (
                 resource.gpus_per_node if resource.gpus_per_node is not None else default_gpus

puzzletron_setup/v2/wizard_common.py

@@ -244,11 +244,11 @@ def _depth_granularity_choices(inventory: Any) -> list[tuple[str, str]]:
 def _default_axis_values(axis: Any) -> list[int]:
     legal_values = tuple(int(value) for value in axis.values)
     teacher = int(axis.teacher_value)
-    half = min(
+    target = min(
         legal_values,
-        key=lambda value: (abs(value - teacher // 2), -value),
+        key=lambda value: (abs(value - teacher * 3 // 4), -value),
     )
-    return list(dict.fromkeys((teacher, half)))
+    return list(dict.fromkeys((teacher, target)))
 
 
 def _text_field(