diff --git a/examples/puzzletron/distributed_eval/run_coordinator.sh b/examples/puzzletron/distributed_eval/run_coordinator.sh index dc06ba6ca48..bdff0f60048 100755 --- a/examples/puzzletron/distributed_eval/run_coordinator.sh +++ b/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 diff --git a/examples/puzzletron/distributed_eval/run_depth_coordinator.sh b/examples/puzzletron/distributed_eval/run_depth_coordinator.sh index dbc7682f931..ed01a577e27 100755 --- a/examples/puzzletron/distributed_eval/run_depth_coordinator.sh +++ b/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[@]}" diff --git a/examples/puzzletron/distributed_eval/run_depth_pool.sh b/examples/puzzletron/distributed_eval/run_depth_pool.sh index 2ab40f3796b..557b18a049c 100644 --- a/examples/puzzletron/distributed_eval/run_depth_pool.sh +++ b/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 diff --git a/examples/puzzletron/distributed_eval/run_replacement_pool.sh b/examples/puzzletron/distributed_eval/run_replacement_pool.sh index 7f6b430c3cd..52b75f35799 100755 --- a/examples/puzzletron/distributed_eval/run_replacement_pool.sh +++ b/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)}" diff --git a/examples/puzzletron/finalize_depth_importance.py b/examples/puzzletron/finalize_depth_importance.py new file mode 100644 index 00000000000..219f23b21a5 --- /dev/null +++ b/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() diff --git a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py index 01ebc247eb8..3e5eb2dd9d4 100644 --- a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py +++ b/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 diff --git a/modelopt/torch/puzzletron/orchestration/adapters/pool.py b/modelopt/torch/puzzletron/orchestration/adapters/pool.py index c9f3c3d756d..7677bb24cdf 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/pool.py +++ b/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), } diff --git a/modelopt/torch/puzzletron/orchestration/controller.py b/modelopt/torch/puzzletron/orchestration/controller.py index 9973ed14a9d..5e04ef1df47 100644 --- a/modelopt/torch/puzzletron/orchestration/controller.py +++ b/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, diff --git a/puzzletron_setup/v2/bundle.py b/puzzletron_setup/v2/bundle.py index ef93b1e1760..b14a31ece31 100644 --- a/puzzletron_setup/v2/bundle.py +++ b/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 diff --git a/puzzletron_setup/v2/wizard_common.py b/puzzletron_setup/v2/wizard_common.py index 10ec36b22a7..fbe59d7a9b1 100644 --- a/puzzletron_setup/v2/wizard_common.py +++ b/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( diff --git a/tests/unit/torch/puzzletron/test_depth_finalizer.py b/tests/unit/torch/puzzletron/test_depth_finalizer.py new file mode 100644 index 00000000000..6f9cd36061e --- /dev/null +++ b/tests/unit/torch/puzzletron/test_depth_finalizer.py @@ -0,0 +1,90 @@ +# 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. + +"""Tests for distributed depth-importance finalization.""" + +import json + +import pytest + +import examples.puzzletron.finalize_depth_importance as depth_finalizer + + +def test_depth_finalizer_publishes_terminal_manifest(tmp_path, monkeypatch): + config_path = tmp_path / "experiment.yaml" + config_path.touch() + output_dir = tmp_path / "depth" / "iterative" + output_dir.mkdir(parents=True) + selected = [ + {"kind": "attention", "layer_idx": 1}, + {"kind": "mamba", "layer_idx": 2}, + ] + (output_dir / "trajectory.json").write_text( + json.dumps( + { + "status": "complete", + "max_removals": 2, + "selected": selected, + "scenarios": [{}, {}, {}], + } + ) + ) + config = { + "model": {"path": "tiny-model"}, + "depth_importance": {"enabled": True, "max_removals": 2}, + } + loaded_overrides = [] + + def load_config(path, *, overrides=None): + assert path == config_path + loaded_overrides.extend(overrides or ()) + return config + + monkeypatch.setattr(depth_finalizer, "pipeline_config_from_path", load_config) + + outputs = depth_finalizer.finalize_depth_importance( + config_path, + tmp_path, + output_dir, + overrides=["++depth_importance.micro_batch_size=1"], + ) + + manifest = json.loads((tmp_path / "manifests" / "depth_importance.json").read_text()) + assert loaded_overrides == ["++depth_importance.micro_batch_size=1"] + assert outputs == { + "trajectory_path": str(output_dir / "trajectory.json"), + "scenario_count": 3, + "selected": selected, + } + assert manifest["stage"] == "depth_importance" + assert manifest["status"] == "success" + assert manifest["outputs"] == outputs + + +def test_depth_finalizer_rejects_incomplete_trajectory(tmp_path, monkeypatch): + config_path = tmp_path / "experiment.yaml" + output_dir = tmp_path / "depth" + output_dir.mkdir() + (output_dir / "trajectory.json").write_text( + json.dumps({"status": "complete", "max_removals": 2, "selected": [{}]}) + ) + monkeypatch.setattr( + depth_finalizer, + "pipeline_config_from_path", + lambda _path, *, overrides=None: {"depth_importance": {"max_removals": 2}}, + ) + + with pytest.raises(RuntimeError, match="depth trajectory is incomplete"): + depth_finalizer.finalize_depth_importance(config_path, tmp_path, output_dir) diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index cf2ad591abb..ecae6e42d7a 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -1570,6 +1570,11 @@ def test_replacement_pool_splits_workers_across_embedding_widths(tmp_path: Path) assert [attempt.task_topology.task_count for attempt in attempts] == [4, 4] assert [attempt.task_topology.gpus_per_task for attempt in attempts] == [4, 4] assert [attempt.command.env["WORKER_COUNT"] for attempt in attempts] == ["4", "4"] + assert [attempt.command.env["PUZZLETRON_WORKER_WORLD_SIZE"] for attempt in attempts] == [ + "4", + "4", + ] + assert all("WORLD_SIZE" not in attempt.command.env for attempt in attempts) assert [attempt.command.env["FINALIZE_OVERRIDES"] for attempt in attempts] == [ "++replacement_scoring.automodel.lm_head_backend=streaming", "++replacement_scoring.automodel.lm_head_backend=streaming", diff --git a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py index dcc672b712b..b1a70576bbb 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py +++ b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py @@ -227,6 +227,30 @@ def _write_configs(tmp_path: Path): return experiment, runner, execution +def test_failed_active_stage_view_preserves_worker_logs(tmp_path): + experiment, runner_path, execution_path = _write_configs(tmp_path) + plan = compile_campaign_plan( + experiment_config_path=experiment, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + ) + executor = _FakeExecutor() + controller = CampaignController(plan, executor=executor) + handle = JobHandle( + backend="fake", + handle_id="active-worker", + attempt_id="attempt-1", + metadata={"log_paths": ("worker.log",)}, + ) + controller._active[handle.handle_id] = (handle, "vllm_stats:0", handle.attempt_id) + controller._failed_stages.add("vllm_stats") + + view = next(item for item in controller._stage_views() if item.stage_id == "vllm_stats") + + assert view.status == "failed" + assert view.log_paths == ("worker.log",) + + def _write_sanity_drain_configs(tmp_path: Path): run_dir = tmp_path / "run" experiment = tmp_path / "experiment.yaml" diff --git a/tests/unit/torch/puzzletron/test_setup_v2_quick.py b/tests/unit/torch/puzzletron/test_setup_v2_quick.py index f69af67ce42..6593894fa07 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_quick.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_quick.py @@ -45,6 +45,7 @@ data_section, infrastructure_section, ) +from puzzletron_setup.v2.wizard_common import _default_axis_values _QWEN_FAMILY_CONFIG = "examples/puzzletron/configs/families/qwen3_5/family.yaml" _NEMOTRON_FAMILY_CONFIG = "examples/puzzletron/configs/families/nemotron3/family.yaml" @@ -71,6 +72,18 @@ def test_only_cpu_slurm_integer_defaults_accept_null() -> None: # Public facade and guided-profile defaults +def test_default_axis_values_target_75_percent_of_teacher() -> None: + axis = AxisInventory( + axis_id="heads", + label="Heads", + teacher_value=16, + values=(16, 12, 8), + alignment=1, + ) + + assert _default_axis_values(axis) == [16, 12] + + def test_guided_profile_defaults_are_selected_by_model_family(tmp_path): families = {} for family, num_solutions in (("first", 2), ("second", 7)): diff --git a/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py b/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py index 8cf142fc5c9..1bf1d1a9c19 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py @@ -557,6 +557,25 @@ def test_execution_uses_resolved_stage_resource_and_parallel_profile(tmp_path: P } +def test_replacement_pool_reserves_one_instance_per_embedding_width(tmp_path: Path) -> None: + state = _campaign_state(tmp_path) + resources = deepcopy(state.collection("stage_resources")) + resources["replacement_scoring"] = { + "strategy": "persistent_pool", + "instances": 1, + "resource": "gpu", + "gpus_per_node": 8, + "profile_name": "model", + } + state.set_collection("stage_resources", resources) + + smoke = render_execution_v2(state, "smoke") + production = render_execution_v2(state, "production") + + assert smoke["execution"]["stages"]["replacement_scoring"]["instances"] == 2 + assert production["execution"]["stages"]["replacement_scoring"]["instances"] == 2 + + def test_execution_strips_model_runtime_fields_from_inline_parallel_mesh(tmp_path: Path) -> None: state = _campaign_state(tmp_path) resources = deepcopy(state.collection("stage_resources")) diff --git a/tests/unit/torch/puzzletron/test_training_flash_kld.py b/tests/unit/torch/puzzletron/test_training_flash_kld.py index 1de04815f54..d4793055643 100644 --- a/tests/unit/torch/puzzletron/test_training_flash_kld.py +++ b/tests/unit/torch/puzzletron/test_training_flash_kld.py @@ -363,3 +363,51 @@ def __init__(self): def test_tp_teacher_projection_redistributes_hidden_and_rewraps_logits(): spawn_multiprocess_job(size=2, job=_tp_teacher_projection_job, backend="gloo") + + +def _dp_shard_hidden_alignment_job(_rank: int, _size: int) -> None: + class HookedHead(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(4, 4)) + self.hook_active = False + + def forward(self, value): + assert self.hook_active + return value + 1 + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.lm_head = HookedHead() + + mesh = init_device_mesh("cpu", (2,), mesh_dim_names=("dp_shard_cp",)) + head = torch.nn.Linear(4, 8, bias=False) + head.weight = torch.nn.Parameter(distribute_tensor(head.weight.detach(), mesh, (Shard(0),))) + hidden = torch.randn(3, 4) + + assert _align_dtensor_to_module_mesh(hidden, head) is hidden + + teacher = Model() + teacher.lm_head.weight = torch.nn.Parameter( + distribute_tensor(teacher.lm_head.weight.detach(), mesh, (Shard(0),)) + ) + assert _install_distillation_head_passthrough([teacher]) == 1 + teacher.lm_head.register_forward_pre_hook( + lambda module, _args: setattr(module, "hook_active", True) + ) + teacher.lm_head.register_forward_hook( + lambda module, _args, _output: setattr(module, "hook_active", False) + ) + + projected = _project_teacher_hidden_on_reference_mesh( + hidden, + teacher.lm_head, + torch.empty_like(hidden), + ) + torch.testing.assert_close(projected, hidden + 1) + assert not teacher.lm_head.hook_active + + +def test_dp_shard_teacher_projection_preserves_rank_local_hidden(): + spawn_multiprocess_job(size=2, job=_dp_shard_hidden_alignment_job, backend="gloo")