From a2f7ad8324800d3471775e0218ded408c71aee64 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Mon, 24 Aug 2026 21:22:57 +0200 Subject: [PATCH 1/4] Add Puzzletron automation runtime controls Signed-off-by: Johannes Rausch --- CHANGELOG.rst | 2 +- examples/puzzletron/README.md | 33 ++- examples/puzzletron/ci_environment.json | 8 +- examples/puzzletron/ci_environment.py | 68 +++++ .../distributed_eval/run_depth_pool.sh | 24 +- .../distributed_eval/run_replacement_pool.sh | 24 +- .../puzzletron/distributed_eval/run_worker.sh | 68 +++-- .../puzzletron/run_profile_aiperf_worker.py | 27 ++ .../torch/puzzletron/benchmarks/aiperf.py | 140 ++++++++-- .../distillation/global_kd_recipe.py | 239 ++++++++++------ .../distributed_eval/automodel_executor.py | 106 ++++--- .../puzzletron/orchestration/adapters/pool.py | 35 ++- .../orchestration/adapters/sharded.py | 101 +++++-- .../puzzletron/orchestration/task_launcher.py | 42 ++- modelopt/torch/puzzletron/post_mip/runner.py | 16 ++ modelopt/torch/puzzletron/security_policy.py | 36 +++ modelopt/torch/puzzletron/stages/future.py | 52 +++- .../torch/puzzletron/utils/vllm_adapter.py | 29 +- noxfile.py | 81 ++++-- puzzletron_setup/bundle.py | 3 +- puzzletron_setup/v2/cli.py | 32 ++- puzzletron_setup/v2/prompts.py | 63 ++++- puzzletron_setup/v2/session.py | 8 +- puzzletron_setup/v2/wizard.py | 127 ++++++--- .../test_aiperf_context_capacity.py | 116 +++++++- .../test_automodel_solution_scoring.py | 44 +++ .../torch/puzzletron/test_ci_environment.py | 185 +++++++++++++ .../torch/puzzletron/test_future_stages.py | 53 +++- .../puzzletron/test_global_kd_canonical.py | 223 +++++++++------ .../test_orchestration_executors.py | 224 ++++++++++++++- .../test_orchestration_task_topology.py | 158 ++++++++++- .../torch/puzzletron/test_post_mip_runner.py | 5 +- .../puzzletron/test_profile_aiperf_worker.py | 80 ++++++ .../torch/puzzletron/test_setup_bundle.py | 7 + .../torch/puzzletron/test_setup_v2_quick.py | 258 +++++++++++++----- .../puzzletron/test_vllm_axis_contract.py | 52 +++- 36 files changed, 2267 insertions(+), 502 deletions(-) create mode 100644 examples/puzzletron/ci_environment.py create mode 100644 modelopt/torch/puzzletron/security_policy.py create mode 100644 tests/unit/torch/puzzletron/test_ci_environment.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7383b652554..8360f03ab21 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,7 +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 Puzzletron dynamic post-MIP downstream evaluation through ``lmms-eval`` with vLLM-backed checkpoint evaluation, setup-wizard topology/resource prompts, non-interactive setup automation, 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/README.md b/examples/puzzletron/README.md index 86af42abd85..43e2ed0d3c6 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -85,6 +85,21 @@ The defaults file is loaded only when passed explicitly and takes precedence over the selected profile. To expose every per-section and nested setting, use the advanced flow explicitly: +Automation can use the same setup entry point without answering prompts. The +defaults file must provide every required value that has no resolved default: + +```bash +python examples/puzzletron/puzzletron_setup_v2.py \ + --defaults /path/to/setup-v2-defaults.yaml \ + --campaign-dir /path/to/campaign \ + --profile smoke \ + --non-interactive +``` + +Non-interactive setup fails instead of guessing when a required answer has no +resolved default. It generates and validates the same smoke and production +bundles as the interactive wizard. + ```bash python examples/puzzletron/puzzletron_setup_v2.py --full ``` @@ -269,13 +284,15 @@ git -C "${AUTOMODEL_ROOT}" rev-parse HEAD ``` ```bash -python - <<'PY' +PYTHONPATH="${MODEL_OPT_ROOT}" python - <<'PY' import importlib.metadata as metadata import json import os from packaging.version import Version +from examples.puzzletron.ci_environment import verify_installed_vcs_source + import aiperf import lmms_eval import modelopt @@ -307,10 +324,17 @@ assert Version(metadata.version("torchvision")).release == Version( ci_environment["torchvision"] ).release assert transformers.__version__ == ci_environment["transformers"] -assert metadata.version("lmms-eval") == ci_environment["lmms_eval"] +assert Version(metadata.version("lmms-eval")).base_version == ( + ci_environment["lmms_eval"]["base_version"] +) assert Version(metadata.version("nemo-automodel")).base_version == ( ci_environment["nemo_automodel"]["base_version"] ) +for package, source in ( + ("lmms-eval", ci_environment["lmms_eval"]), + ("nemo-automodel", ci_environment["nemo_automodel"]), +): + verify_installed_vcs_source(package, source) assert torch.version.cuda == "12.9" assert torch.cuda.is_available() PY @@ -492,6 +516,11 @@ The setup wizard can also add downstream evaluation for materialized campaign candidates. See [post-MIP pipelines](docs/post_mip_pipeline.md) to configure it or add it to an existing campaign. +Remote model code and AIPerf v0.11 online tokenizer resolution are disabled by +default. Enable remote code only for a trusted model source. The tokenizer +compatibility option permits the AIPerf child process to resolve its tokenizer +online even when the surrounding campaign is configured for offline loading. + ### Legacy checked-in Nano campaign The checked-in Nano experiment uses the legacy `zero_shot_evaluation`, diff --git a/examples/puzzletron/ci_environment.json b/examples/puzzletron/ci_environment.json index 23f43e28623..9e39fd5d6b8 100644 --- a/examples/puzzletron/ci_environment.json +++ b/examples/puzzletron/ci_environment.json @@ -1,11 +1,15 @@ { "schema_version": 1, - "scope": "puzzletron_v2_cpu_ci", + "scope": "puzzletron_v2_ci", "python": "3.12", "torch": "2.11.0", "torchvision": "0.26.0", "transformers": "5.8.1", - "lmms_eval": "0.7.0", + "lmms_eval": { + "base_version": "0.7.0", + "repository": "https://github.com/EvolvingLMMs-Lab/lmms-eval.git", + "commit": "15c32bfec165df13c269ddd3cda03b2ed9137825" + }, "nemo_automodel": { "base_version": "0.5.0", "repository": "https://github.com/Separius/Automodel.git", diff --git a/examples/puzzletron/ci_environment.py b/examples/puzzletron/ci_environment.py new file mode 100644 index 00000000000..fe6924e9841 --- /dev/null +++ b/examples/puzzletron/ci_environment.py @@ -0,0 +1,68 @@ +# 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. + +"""Verification helpers for the pinned Puzzletron CI environment.""" + +from __future__ import annotations + +import json +import subprocess +from importlib import metadata +from typing import Any +from urllib.parse import unquote, urlparse + +__all__ = ["verify_installed_vcs_source"] + + +def _normalized_repository(url: object) -> str: + return str(url or "").removesuffix(".git").rstrip("/") + + +def _installed_vcs_source(package: str) -> tuple[str | None, str | None]: + payload = json.loads(metadata.distribution(package).read_text("direct_url.json") or "{}") + vcs_info = payload.get("vcs_info") or {} + if vcs_info.get("commit_id"): + return payload.get("url"), vcs_info["commit_id"] + if (payload.get("dir_info") or {}).get("editable") and str(payload.get("url", "")).startswith( + "file:" + ): + root = unquote(urlparse(payload["url"]).path) + repository = subprocess.check_output( + ["git", "-C", root, "remote", "get-url", "origin"], text=True + ).strip() + commit = subprocess.check_output( + ["git", "-C", root, "rev-parse", "HEAD"], text=True + ).strip() + dirty = subprocess.check_output( + ["git", "-C", root, "status", "--porcelain", "--untracked-files=all"], + text=True, + ).strip() + if dirty: + raise RuntimeError(f"Pinned Puzzletron dependency {package!r} is dirty: {dirty}") + return repository, commit + return payload.get("url"), vcs_info.get("commit_id") + + +def verify_installed_vcs_source(package: str, expected: dict[str, Any]) -> None: + """Require an installed VCS dependency to match its repository and commit.""" + + repository, commit = _installed_vcs_source(package) + expected_source = (_normalized_repository(expected["repository"]), expected["commit"]) + actual_source = (_normalized_repository(repository), commit) + if actual_source != expected_source: + raise RuntimeError( + f"Pinned Puzzletron dependency {package!r} source mismatch: " + f"actual={actual_source!r}, expected={expected_source!r}" + ) diff --git a/examples/puzzletron/distributed_eval/run_depth_pool.sh b/examples/puzzletron/distributed_eval/run_depth_pool.sh index 3b8b663c9b4..2ab40f3796b 100644 --- a/examples/puzzletron/distributed_eval/run_depth_pool.sh +++ b/examples/puzzletron/distributed_eval/run_depth_pool.sh @@ -1,6 +1,18 @@ #!/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 @@ -10,10 +22,12 @@ set -Eeuo pipefail : "${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}" +: "${PUZZLETRON_GROUP_RANK:=0}" PYTHON_BIN="${PYTHON_BIN:-python}" SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" GROUP_INDEX="${PUZZLETRON_GROUP_INDEX}" +GROUP_RANK="${PUZZLETRON_GROUP_RANK}" JOB_ID="${SLURM_JOB_ID:-local}" WORKER_PREFIX="${JOB_ID}-depth-" MANIFEST_PATH="${CAMPAIGN_DIR}/manifest.json" @@ -38,7 +52,7 @@ cleanup() { local rc=$? trap - EXIT INT TERM set +e - if [[ "${GROUP_INDEX}" == "0" ]]; then + if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" ]]; then drain_workers fi if [[ -n "${worker_pid}" ]] && kill -0 "${worker_pid}" 2>/dev/null; then @@ -50,7 +64,7 @@ cleanup() { trap cleanup EXIT INT TERM # Rank 0 creates the shared campaign before any worker attempts to open it. -if [[ "${GROUP_INDEX}" == "0" && ! -f "${MANIFEST_PATH}" ]]; then +if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" && ! -f "${MANIFEST_PATH}" ]]; then CUDA_VISIBLE_DEVICES="" "${PYTHON_BIN}" \ -m modelopt.torch.puzzletron.distributed_eval.cli init \ --campaign-dir "${CAMPAIGN_DIR}" \ @@ -72,21 +86,17 @@ done # Every scheduler task owns one GPU slice and starts one worker group. Multiple # independent worker groups may share a node. -export NNODES=1 -export NODE_RANK=0 export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORLD_SIZE}}" export WORKER_GROUP_INDEX="${GROUP_INDEX}" export WORKER_ID="${WORKER_PREFIX}${GROUP_INDEX}" export WORKER_HOST="${WORKER_HOST:-$(hostname -f)}" export WORKER_PORT="${WORKER_PORT:-$((5010 + GROUP_INDEX))}" -export RDZV_ENDPOINT="127.0.0.1:$((29500 + GROUP_INDEX))" -export RDZV_ID="depth-${JOB_ID}-${GROUP_INDEX}" bash "${SCRIPT_DIR}/run_worker.sh" & worker_pid=$! coordinator_rc=0 -if [[ "${GROUP_INDEX}" == "0" ]]; then +if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" ]]; then # Do not start depth iteration zero until every resident model is ready. CUDA_VISIBLE_DEVICES="" "${PYTHON_BIN}" - \ "${CAMPAIGN_DIR}" \ diff --git a/examples/puzzletron/distributed_eval/run_replacement_pool.sh b/examples/puzzletron/distributed_eval/run_replacement_pool.sh index 283a9ef11a3..c5ebc9ebc56 100755 --- a/examples/puzzletron/distributed_eval/run_replacement_pool.sh +++ b/examples/puzzletron/distributed_eval/run_replacement_pool.sh @@ -1,6 +1,18 @@ #!/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 @@ -10,10 +22,12 @@ set -Eeuo pipefail : "${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}" +: "${PUZZLETRON_GROUP_RANK:=0}" PYTHON_BIN="${PYTHON_BIN:-python}" SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" GROUP_INDEX="${PUZZLETRON_GROUP_INDEX}" +GROUP_RANK="${PUZZLETRON_GROUP_RANK}" JOB_ID="${SLURM_JOB_ID:-local}" WORKER_PREFIX="${JOB_ID}-replacement-" MANIFEST_PATH="${CAMPAIGN_DIR}/manifest.json" @@ -38,7 +52,7 @@ cleanup() { local rc=$? trap - EXIT INT TERM set +e - if [[ "${GROUP_INDEX}" == "0" ]]; then + if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" ]]; then drain_workers fi if [[ -n "${worker_pid}" ]] && kill -0 "${worker_pid}" 2>/dev/null; then @@ -49,7 +63,7 @@ cleanup() { } trap cleanup EXIT INT TERM -if [[ "${GROUP_INDEX}" == "0" && ! -f "${MANIFEST_PATH}" ]]; then +if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" && ! -f "${MANIFEST_PATH}" ]]; then CUDA_VISIBLE_DEVICES="" "${PYTHON_BIN}" \ -m modelopt.torch.puzzletron.distributed_eval.cli init \ --campaign-dir "${CAMPAIGN_DIR}" \ @@ -69,21 +83,17 @@ while [[ ! -f "${MANIFEST_PATH}" ]]; do sleep 1 done -export NNODES=1 -export NODE_RANK=0 export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORLD_SIZE}}" export WORKER_GROUP_INDEX="${GROUP_INDEX}" export WORKER_ID="${WORKER_PREFIX}${GROUP_INDEX}" export WORKER_HOST="${WORKER_HOST:-$(hostname -f)}" export WORKER_PORT="${WORKER_PORT:-$((5010 + GROUP_INDEX))}" -export RDZV_ENDPOINT="127.0.0.1:$((29500 + GROUP_INDEX))" -export RDZV_ID="replacement-${JOB_ID}-${GROUP_INDEX}" bash "${SCRIPT_DIR}/run_worker.sh" & worker_pid=$! coordinator_rc=0 -if [[ "${GROUP_INDEX}" == "0" ]]; then +if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" ]]; then CUDA_VISIBLE_DEVICES="" "${PYTHON_BIN}" - \ "${CAMPAIGN_DIR}" \ "${WORKER_COUNT}" \ diff --git a/examples/puzzletron/distributed_eval/run_worker.sh b/examples/puzzletron/distributed_eval/run_worker.sh index b95b77f3119..911c0503f6f 100755 --- a/examples/puzzletron/distributed_eval/run_worker.sh +++ b/examples/puzzletron/distributed_eval/run_worker.sh @@ -1,4 +1,19 @@ #!/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}" @@ -6,11 +21,21 @@ set -Eeuo pipefail PYTHON_BIN="${PYTHON_BIN:-python}" TORCHRUN="${TORCHRUN:-torchrun}" -NNODES="${NNODES:-1}" NPROC_PER_NODE="${NPROC_PER_NODE:-8}" -NODE_RANK="${NODE_RANK:-0}" -RDZV_ID="${RDZV_ID:-distributed-eval-${SLURM_JOB_ID:-local}}" -RDZV_ENDPOINT="${RDZV_ENDPOINT:-127.0.0.1:29500}" +if [[ -n "${PUZZLETRON_GROUP_SIZE:-}" ]]; then + : "${PUZZLETRON_GROUP_RANK:?set PUZZLETRON_GROUP_RANK with task identity}" + : "${PUZZLETRON_RENDEZVOUS_ENDPOINT:?set PUZZLETRON_RENDEZVOUS_ENDPOINT with task identity}" + : "${PUZZLETRON_RENDEZVOUS_ID:?set PUZZLETRON_RENDEZVOUS_ID with task identity}" + NNODES="${PUZZLETRON_GROUP_SIZE}" + NODE_RANK="${PUZZLETRON_GROUP_RANK}" + RDZV_ENDPOINT="${PUZZLETRON_RENDEZVOUS_ENDPOINT}" + RDZV_ID="${PUZZLETRON_RENDEZVOUS_ID}" +else + NNODES="${NNODES:-1}" + NODE_RANK="${NODE_RANK:-0}" + RDZV_ID="${RDZV_ID:-distributed-eval-${SLURM_JOB_ID:-local}}" + RDZV_ENDPOINT="${RDZV_ENDPOINT:-127.0.0.1:29500}" +fi WORKER_HOST="${WORKER_HOST:-$(hostname -f)}" WORKER_PORT="${WORKER_PORT:-5010}" WORKER_ID="${WORKER_ID:-${SLURM_JOB_ID:-local}-group-${WORKER_GROUP_INDEX:-0}}" @@ -25,18 +50,23 @@ fi export TORCH_NCCL_ASYNC_ERROR_HANDLING=1 export PYTHONUNBUFFERED=1 -exec "${TORCHRUN}" \ - --nnodes "${NNODES}" \ - --nproc-per-node "${NPROC_PER_NODE}" \ - --node-rank "${NODE_RANK}" \ - --rdzv-backend c10d \ - --rdzv-id "${RDZV_ID}" \ - --rdzv-endpoint "${RDZV_ENDPOINT}" \ - -m modelopt.torch.puzzletron.distributed_eval.cli worker \ - --campaign-dir "${CAMPAIGN_DIR}" \ - --config "${CONFIG_PATH}" \ - --host "${WORKER_HOST}" \ - --port "${WORKER_PORT}" \ - --worker-id "${WORKER_ID}" \ - --heartbeat-seconds "${HEARTBEAT_SECONDS:-10}" \ - "${override_args[@]}" +worker_command=( + "${TORCHRUN}" + --nnodes "${NNODES}" + --nproc-per-node "${NPROC_PER_NODE}" + --node-rank "${NODE_RANK}" + --rdzv-backend c10d + --rdzv-id "${RDZV_ID}" + --rdzv-endpoint "${RDZV_ENDPOINT}" + -m modelopt.torch.puzzletron.distributed_eval.cli worker + --campaign-dir "${CAMPAIGN_DIR}" + --config "${CONFIG_PATH}" + --host "${WORKER_HOST}" + --port "${WORKER_PORT}" + --worker-id "${WORKER_ID}" + --heartbeat-seconds "${HEARTBEAT_SECONDS:-10}" +) +if ((${#override_args[@]})); then + worker_command+=("${override_args[@]}") +fi +exec "${worker_command[@]}" diff --git a/examples/puzzletron/run_profile_aiperf_worker.py b/examples/puzzletron/run_profile_aiperf_worker.py index 8fa381ca1ec..930813a3b61 100644 --- a/examples/puzzletron/run_profile_aiperf_worker.py +++ b/examples/puzzletron/run_profile_aiperf_worker.py @@ -188,7 +188,16 @@ def run_worker( concurrencies: tuple[int, ...] | None = None, request_count: int | None = None, benchmark_timeout: float = 7200, + trust_remote_code: bool = False, + allow_aiperf_v011_online_tokenizer_resolution: bool = False, ) -> Path: + """Run one AIPerf shard with security-sensitive behavior disabled by default. + + ``trust_remote_code`` is only appropriate for trusted model sources. The + AIPerf v0.11 compatibility option permits online tokenizer resolution for + the AIPerf child process even when the campaign otherwise runs offline. + """ + # Worker execution needs the GPU stack; result merging intentionally remains # usable by the dependency-light login-node orchestrator. from modelopt.torch.puzzletron.benchmarks import run_aiperf_sweep @@ -241,6 +250,10 @@ def run_worker( seed=42, gpu_telemetry="pynvml", benchmark_timeout=benchmark_timeout, + trust_remote_code=trust_remote_code, + allow_aiperf_v011_online_tokenizer_resolution=( + allow_aiperf_v011_online_tokenizer_resolution + ), ) rows.extend(result.model_dump(mode="json") for result in results) selection_suffix = "" @@ -335,6 +348,16 @@ def main() -> None: parser.add_argument("--concurrency", type=int, action="append", default=[]) parser.add_argument("--request-count", type=int) parser.add_argument("--benchmark-timeout", type=float, default=7200) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Allow remote model code; use only with trusted model sources.", + ) + parser.add_argument( + "--allow-aiperf-v011-online-tokenizer-resolution", + action="store_true", + help="Permit online tokenizer resolution in the AIPerf v0.11 child process.", + ) parser.add_argument("--preflight", action="store_true") parser.add_argument("--merge", action="store_true") args = parser.parse_args() @@ -370,6 +393,10 @@ def main() -> None: concurrencies=tuple(args.concurrency) or None, request_count=args.request_count, benchmark_timeout=args.benchmark_timeout, + trust_remote_code=args.trust_remote_code, + allow_aiperf_v011_online_tokenizer_resolution=( + args.allow_aiperf_v011_online_tokenizer_resolution + ), ) print(output) diff --git a/modelopt/torch/puzzletron/benchmarks/aiperf.py b/modelopt/torch/puzzletron/benchmarks/aiperf.py index 872822777bd..0bb2458b482 100644 --- a/modelopt/torch/puzzletron/benchmarks/aiperf.py +++ b/modelopt/torch/puzzletron/benchmarks/aiperf.py @@ -35,18 +35,25 @@ from datetime import datetime, timezone from pathlib import Path from threading import Lock -from typing import Any, Iterable +from typing import TYPE_CHECKING, Any, Iterable, cast from ..identity import stable_hash from ..orchestration.mesh import normalize_vllm_topology from .schema import BenchmarkResult +if TYPE_CHECKING: + from ..anymodel.model_descriptor import ModelDescriptor + __all__ = ["run_aiperf_benchmark", "run_aiperf_sweep"] _CHECKPOINT_PREPARE_LOCK = Lock() -def _prepare_vllm_checkpoint(checkpoint_dir: Path) -> bool: +def _prepare_vllm_checkpoint( + checkpoint_dir: Path, + *, + trust_remote_code: bool = False, +) -> bool: """Restore AnyModel metadata lost by generic HF checkpoint consolidation.""" with _CHECKPOINT_PREPARE_LOCK: config = json.loads((checkpoint_dir / "config.json").read_text()) @@ -57,7 +64,10 @@ def _prepare_vllm_checkpoint(checkpoint_dir: Path) -> bool: return False from ..utils.vllm_adapter import refresh_realized_checkpoint_config - refresh_realized_checkpoint_config(checkpoint_dir) + refresh_realized_checkpoint_config( + checkpoint_dir, + trust_remote_code=trust_remote_code, + ) return True @@ -67,7 +77,8 @@ def _descriptor_vllm_args(checkpoint_dir: Path) -> list[str]: config = json.loads((checkpoint_dir / "config.json").read_text()) resolution = resolve_descriptor(config) - return [str(arg) for arg in resolution.descriptor.runtime_vllm_benchmark_args(config)] + descriptor = cast("type[ModelDescriptor]", resolution.descriptor) + return [str(arg) for arg in descriptor.runtime_vllm_benchmark_args(config)] def _free_port() -> int: @@ -309,6 +320,53 @@ def _profile_command( return command +def _vllm_server_command( + *, + checkpoint_dir: Path, + port: int, + model_name: str, + input_tokens: int, + output_tokens: int, + topology: dict[str, Any], + trust_remote_code: bool, +) -> list[str]: + """Build the vLLM command under the caller's explicit code-trust policy.""" + + command = [ + "vllm", + "serve", + str(checkpoint_dir), + "--host", + "127.0.0.1", + "--port", + str(port), + "--served-model-name", + model_name, + "--max-model-len", + str(_server_max_model_len(input_tokens, output_tokens, topology)), + ] + if trust_remote_code: + command.append("--trust-remote-code") + command.extend(_topology_vllm_args(topology)) + command.extend(_descriptor_vllm_args(checkpoint_dir)) + extra_vllm_args = tuple(str(arg) for arg in topology.get("extra_vllm_args", ())) + reserved_options = {"--config", "--trust-remote-code"} + normalized_options = {arg.partition("=")[0].replace("_", "-") for arg in extra_vllm_args} + overridden_options = { + reserved + for option in normalized_options + for reserved in reserved_options + if reserved.startswith(option) + } + if overridden_options: + raise ValueError( + "topology.extra_vllm_args cannot set policy-owned vLLM options: " + + ", ".join(sorted(overridden_options)) + ) + command.extend(extra_vllm_args) + return command + + def _clean_subprocess_environment( gpu_ids: str, *, architecture_id: str, topology_id: str ) -> dict[str, str]: @@ -355,6 +413,24 @@ def _clean_subprocess_environment( return env +def _aiperf_subprocess_environment( + env: dict[str, str], + *, + allow_aiperf_v011_online_tokenizer_resolution: bool = False, +) -> dict[str, str]: + """Work around AIPerf v0.11's broken offline local-tokenizer resolution. + + Remove this compatibility option after the pinned AIPerf resolver accepts + absolute local tokenizer directories while offline. + """ + + resolved = dict(env) + if allow_aiperf_v011_online_tokenizer_resolution: + resolved.pop("HF_HUB_OFFLINE", None) + resolved.pop("TRANSFORMERS_OFFLINE", None) + return resolved + + def run_aiperf_sweep( checkpoint_dir: str | Path, *, @@ -376,13 +452,18 @@ def run_aiperf_sweep( readiness_timeout: float = 1200, benchmark_timeout: float = 600, gpu_telemetry: str | None = "pynvml", + trust_remote_code: bool = False, + allow_aiperf_v011_online_tokenizer_resolution: bool = False, ) -> list[BenchmarkResult]: - """Run multiple concurrencies against one persistent vLLM server.""" + """Run a serving sweep while preserving offline and remote-code policy by default.""" checkpoint_dir = Path(checkpoint_dir).resolve() artifact_dir = Path(artifact_dir).resolve() artifact_dir.mkdir(parents=True, exist_ok=True) - _prepare_vllm_checkpoint(checkpoint_dir) + _prepare_vllm_checkpoint( + checkpoint_dir, + trust_remote_code=trust_remote_code, + ) concurrency_values = tuple(int(value) for value in concurrencies) if not concurrency_values or len(set(concurrency_values)) != len(concurrency_values): raise ValueError("AIPerf concurrencies must be non-empty and unique") @@ -405,23 +486,15 @@ def run_aiperf_sweep( model_name = f"puzzletron-{architecture_id[:16]}" tokenizer_dir = _short_tokenizer_alias(checkpoint_dir, artifact_dir) server_log = artifact_dir / "vllm_server.log" - server_cmd = [ - "vllm", - "serve", - str(checkpoint_dir), - "--host", - "127.0.0.1", - "--port", - str(port), - "--served-model-name", - model_name, - "--max-model-len", - str(_server_max_model_len(input_tokens, output_tokens, topology)), - "--trust-remote-code", - ] - server_cmd.extend(_topology_vllm_args(topology)) - server_cmd.extend(_descriptor_vllm_args(checkpoint_dir)) - server_cmd.extend(str(arg) for arg in topology.get("extra_vllm_args", ())) + server_cmd = _vllm_server_command( + checkpoint_dir=checkpoint_dir, + port=port, + model_name=model_name, + input_tokens=input_tokens, + output_tokens=output_tokens, + topology=topology, + trust_remote_code=trust_remote_code, + ) executable = _resolve_executable(executable) env = _clean_subprocess_environment( gpu_ids, @@ -430,6 +503,12 @@ def run_aiperf_sweep( ) for key, value in (topology.get("env") or {}).items(): env[str(key)] = str(value) + aiperf_env = _aiperf_subprocess_environment( + env, + allow_aiperf_v011_online_tokenizer_resolution=( + allow_aiperf_v011_online_tokenizer_resolution + ), + ) cached: dict[int, BenchmarkResult] = {} missing: list[tuple[int, Path, list[str], str]] = [] for concurrency in concurrency_values: @@ -463,6 +542,10 @@ def run_aiperf_sweep( "endpoint_type": endpoint_type, "extra_inputs": _exact_length_extra_inputs(extra_inputs, output_tokens), "use_server_token_count": use_server_token_count, + "trust_remote_code": trust_remote_code, + "allow_aiperf_v011_online_tokenizer_resolution": ( + allow_aiperf_v011_online_tokenizer_resolution + ), "revisions": revisions, }, prefix="aiperf_result", @@ -470,7 +553,7 @@ def run_aiperf_sweep( metadata_path = run_dir / "puzzletron_aiperf_result.json" export = run_dir / "profile_export_aiperf.json" if metadata_path.is_file() and export.is_file(): - result = BenchmarkResult.model_validate(json.loads(metadata_path.read_text())) + result = BenchmarkResult(**json.loads(metadata_path.read_text())) if result.cache_identity == cache_identity: cached[concurrency] = result continue @@ -496,7 +579,7 @@ def run_aiperf_sweep( command, check=True, timeout=benchmark_timeout, - env=env, + env=aiperf_env, ) export = run_dir / "profile_export_aiperf.json" if not export.is_file(): @@ -542,8 +625,13 @@ def run_aiperf_sweep( command=tuple(command), started_at=started_at, ) + result_payload = ( + result.model_dump(mode="json") + if hasattr(result, "model_dump") + else json.loads(result.json()) + ) (run_dir / "puzzletron_aiperf_result.json").write_text( - json.dumps(result.model_dump(mode="json"), indent=2, sort_keys=True) + "\n" + json.dumps(result_payload, indent=2, sort_keys=True) + "\n" ) cached[concurrency] = result finally: diff --git a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py index c874ff0e332..9aceb6a9441 100644 --- a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py +++ b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py @@ -1,3 +1,18 @@ +# 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. + # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,11 +28,13 @@ import dataclasses import hashlib +import json import os import types from collections import deque from contextlib import nullcontext -from typing import Any +from pathlib import Path +from typing import Any, Callable import torch from nemo_automodel.components.distributed.config import DistributedSetup @@ -61,6 +78,7 @@ from ..plugins.automodel.batch_adapter import VisionForwardMonitor from ..plugins.automodel.pp_utils import set_pp_vlm_chunk_specs +from ..security_policy import require_boolean_policy from .flash_kld import TrainingFlashKLD @@ -109,9 +127,7 @@ def _global_kd_checkpoint_adapter_context(model_parts, descriptor_name: str | No block_configs = _config_value(text_config, "block_configs") if not block_configs: continue - active_descriptor = descriptor_name or _config_value( - config, "anymodel_descriptor" - ) + active_descriptor = descriptor_name or _config_value(config, "anymodel_descriptor") if not active_descriptor: continue descriptor = AutoModelDescriptorFactory.get(str(active_descriptor)) @@ -159,13 +175,17 @@ def non_strict(*args, _original=original, _name=name, options=None, **kwargs): and len(args) >= 2 ): model, optimizer = args[:2] - model_parameters = {id(parameter): fqn for fqn, parameter in model.named_parameters()} + model_parameters = { + id(parameter): fqn for fqn, parameter in model.named_parameters() + } unmatched = [] for group_index, group in enumerate(optimizer.param_groups): for parameter_index, parameter in enumerate(group["params"]): if id(parameter) in model_parameters: continue - local = parameter.to_local() if isinstance(parameter, DTensor) else parameter + local = ( + parameter.to_local() if isinstance(parameter, DTensor) else parameter + ) unmatched.append( { "group": group_index, @@ -174,7 +194,9 @@ def non_strict(*args, _original=original, _name=name, options=None, **kwargs): "local_shape": tuple(local.shape), "requires_grad": bool(parameter.requires_grad), "has_grad": parameter.grad is not None, - "state": sorted(str(key) for key in optimizer.state.get(parameter, {})), + "state": sorted( + str(key) for key in optimizer.state.get(parameter, {}) + ), } ) if unmatched: @@ -185,7 +207,7 @@ def non_strict(*args, _original=original, _name=name, options=None, **kwargs): ) return _original(*args, options=relaxed(options), **kwargs) - non_strict._puzzletron_pp_non_strict = True + setattr(non_strict, "_puzzletron_pp_non_strict", True) setattr(stateful_wrappers, name, non_strict) @@ -213,23 +235,23 @@ def _attach_global_kd_gdn_traces(parts, *, prefix: str, trace_backward: bool = F def _forward_end(_module, _args, output, *, layer_idx=layer_idx): _trace_global_kd_phase(f"{prefix}_gdn_{layer_idx}_forward_end") if trace_backward and isinstance(output, torch.Tensor) and output.requires_grad: - output.register_hook( - lambda grad, layer_idx=layer_idx: ( - _trace_global_kd_phase(f"{prefix}_gdn_{layer_idx}_backward_begin"), - grad, - )[1] - ) + + def trace_backward_begin(grad, *, layer_idx=layer_idx): + _trace_global_kd_phase(f"{prefix}_gdn_{layer_idx}_backward_begin") + return grad + + output.register_hook(trace_backward_begin) module.register_forward_hook(_forward_end) if trace_backward: parameter = next(module.parameters(), None) if parameter is not None and parameter.requires_grad: - parameter.register_hook( - lambda grad, layer_idx=layer_idx: ( - _trace_global_kd_phase(f"{prefix}_gdn_{layer_idx}_parameter_grad"), - grad, - )[1] - ) + + def trace_parameter_grad(grad, *, layer_idx=layer_idx): + _trace_global_kd_phase(f"{prefix}_gdn_{layer_idx}_parameter_grad") + return grad + + parameter.register_hook(trace_parameter_grad) def _instantiate(node): @@ -312,9 +334,7 @@ def _project_teacher_hidden_on_reference_mesh(hidden, teacher_head, reference_lo return projection(hidden) if projection is not None else teacher_head(hidden) local_hidden = hidden.to_local() if isinstance(hidden, DTensor) else hidden reference_local = ( - reference_logits.to_local() - if isinstance(reference_logits, DTensor) - else reference_logits + reference_logits.to_local() if isinstance(reference_logits, DTensor) else reference_logits ) expected_vocab = int(reference_local.shape[-1]) @@ -550,9 +570,7 @@ def _set_teacher_mtp_enabled(teacher): def _split_output(output, model): seq_idx = None - main_is_hidden = bool( - getattr(model, "_puzzletron_distillation_hidden_output", False) - ) + main_is_hidden = bool(getattr(model, "_puzzletron_distillation_hidden_output", False)) if isinstance(output, tuple): values = list(output) if values and isinstance(values[-1], torch.Tensor) and values[-1].dtype == torch.int32: @@ -577,6 +595,23 @@ def _split_output(output, model): class _WeightedObjectiveMixin: + """Objective behavior shared by the LLM and VLM AutoModel recipe bases.""" + + cfg: Any + checkpointer: Any + device_mesh: Any + dist_env: Any + loss_fn: Any + metric_logger_train: Any + model_parts: Any + optimizer: Any + pp: Any + pp_enabled: bool + teacher_model: Any + _ce_loss_buffer: list[torch.Tensor] + _kd_loss_buffer: list[torch.Tensor] + _dp_allreduce: Callable[..., torch.Tensor] + def _configure_objective(self): objective = self.cfg.get("objective", {}) self.objective = { @@ -595,8 +630,7 @@ def _configure_objective(self): self._objective_step_cursor = {name: 0 for name in self.objective} self._loss_topology_logged = False self._gradient_squared = { - name: torch.tensor(0.0) - for name in ("vision", "projector", "language", "mtp") + name: torch.tensor(0.0) for name in ("vision", "projector", "language", "mtp") } self._gradient_hook_handles = [] self._vision_monitors = [] @@ -619,25 +653,51 @@ def save_checkpoint( ): """Publish a completion marker only after model and optimizer DCP succeed.""" - result = super().save_checkpoint( + result = super().save_checkpoint( # type: ignore[misc] epoch, step, train_loss, val_loss, best_metric_key=best_metric_key, ) - checkpoint_path = ( - os.path.join( - str(self.checkpointer.config.checkpoint_dir), - f"epoch_{epoch}_step_{step}", - ) + checkpoint_path = os.path.join( + str(self.checkpointer.config.checkpoint_dir), + f"epoch_{epoch}_step_{step}", ) + publication_error: Exception | None = None + publication_error_text: str | None = None if self.dist_env.is_main: - from pathlib import Path - - Path(checkpoint_path, "saving_completed").touch() + try: + consolidated = Path(checkpoint_path, "model", "consolidated") + config_path = consolidated / "config.json" + config = json.loads(config_path.read_text()) if config_path.is_file() else {} + if config.get("block_configs"): + from ..utils.vllm_adapter import refresh_realized_checkpoint_config + + model_config = _config_value(getattr(self, "cfg", None), "model") + configured_trust = _config_value(model_config, "trust_remote_code") + refresh_realized_checkpoint_config( + consolidated, + trust_remote_code=require_boolean_policy( + configured_trust, + path="model.trust_remote_code", + default=False, + ), + ) + Path(checkpoint_path, "saving_completed").touch() + except Exception as error: # noqa: BLE001 - all ranks must reach the collective + publication_error = error + publication_error_text = f"{type(error).__name__}: {error}" if torch.distributed.is_initialized(): - torch.distributed.barrier() + publication_status = [publication_error_text] + torch.distributed.broadcast_object_list(publication_status, src=0) + publication_error_text = publication_status[0] + if publication_error is not None: + raise publication_error + if publication_error_text is not None: + raise RuntimeError( + f"global KD checkpoint publication failed on rank 0: {publication_error_text}" + ) return result def _install_vision_observers(self, parts, *, role: str): @@ -692,26 +752,23 @@ def observability_metadata(self): } if not torch.distributed.is_initialized(): return local - gathered = [None] * torch.distributed.get_world_size() + gathered: list[dict[str, Any] | None] = [None] * torch.distributed.get_world_size() torch.distributed.all_gather_object(gathered, local) - roles = set().union(*(item["vision_by_role"] for item in gathered)) + observations = [item for item in gathered if item is not None] + if len(observations) != len(gathered): + raise RuntimeError("Missing global KD observability metadata from a distributed rank") + roles = set().union(*(item["vision_by_role"] for item in observations)) return { - "vision_forward_count": sum(item["vision_forward_count"] for item in gathered), + "vision_forward_count": sum(item["vision_forward_count"] for item in observations), "vision_by_role": { - role: sum(item["vision_by_role"].get(role, 0) for item in gathered) + role: sum(item["vision_by_role"].get(role, 0) for item in observations) for role in sorted(roles) }, "vision_output_checksums": sorted( - checksum - for item in gathered - for checksum in item["vision_output_checksums"] + checksum for item in observations for checksum in item["vision_output_checksums"] ), "media_input_checksums": sorted( - set( - checksum - for item in gathered - for checksum in item["media_input_checksums"] - ) + set(checksum for item in observations for checksum in item["media_input_checksums"]) ), } @@ -765,7 +822,9 @@ def _remove_text_inactive_optimizer_parameters(self) -> None: if not inactive_ids: return - optimizers = self.optimizer if isinstance(self.optimizer, (list, tuple)) else [self.optimizer] + optimizers = ( + self.optimizer if isinstance(self.optimizer, (list, tuple)) else [self.optimizer] + ) removed = 0 for optimizer in optimizers: for group in optimizer.param_groups: @@ -792,7 +851,7 @@ def load_checkpoint(self, restore_from=None): return None if getattr(self, "_puzzletron_global_kd_domain", None) == "llm": self._remove_text_inactive_optimizer_parameters() - return super().load_checkpoint(restore_from or "LATEST") + return super().load_checkpoint(restore_from or "LATEST") # type: ignore[misc] def _install_gradient_norm_observers(self): self._gradient_squared = { @@ -818,11 +877,11 @@ def observe_optimizer_step(optimizer, _args, _kwargs): continue gradient = parameter.grad value = gradient.to_local() if isinstance(gradient, DTensor) else gradient - self._gradient_squared[group].add_( - value.detach().float().square().sum() - ) + self._gradient_squared[group].add_(value.detach().float().square().sum()) - optimizers = self.optimizer if isinstance(self.optimizer, (list, tuple)) else [self.optimizer] + optimizers = ( + self.optimizer if isinstance(self.optimizer, (list, tuple)) else [self.optimizer] + ) for optimizer in optimizers: self._gradient_hook_handles.append( optimizer.register_step_pre_hook(observe_optimizer_step) @@ -852,12 +911,12 @@ def _rebind_optimizer_to_current_model_parameters(self, model=None) -> None: seen.add(id(parameter)) current_parameters.append(parameter) - optimizers = self.optimizer if isinstance(self.optimizer, (list, tuple)) else [self.optimizer] + optimizers = ( + self.optimizer if isinstance(self.optimizer, (list, tuple)) else [self.optimizer] + ) for optimizer in optimizers: optimizer_parameters = [ - parameter - for group in optimizer.param_groups - for parameter in group["params"] + parameter for group in optimizer.param_groups for parameter in group["params"] ] current_ids = {id(parameter) for parameter in current_parameters} if all(id(parameter) in current_ids for parameter in optimizer_parameters): @@ -994,9 +1053,7 @@ def _teacher_loss_model(self): if teacher_pp is None: return self.teacher_model return next( - part - for part, stage in zip(teacher_pp.parts, teacher_pp.info.stages) - if stage.is_last + part for part, stage in zip(teacher_pp.parts, teacher_pp.info.stages) if stage.is_last ) @staticmethod @@ -1159,9 +1216,7 @@ def _mtp_objective_losses( student_head = _get_lm_head_module(student_model) if student_is_hidden else None teacher_head = ( - _get_lm_head_module(teacher_model) - if needs_mtp_kd and teacher_is_hidden - else None + _get_lm_head_module(teacher_model) if needs_mtp_kd and teacher_is_hidden else None ) if student_is_hidden and student_head is None: raise ValueError("MTP losses require an accessible student lm_head") @@ -1187,9 +1242,7 @@ def _mtp_objective_losses( depth_labels = torch.where(rolled == seq_idx, depth_labels, -100) flat_student = self._flatten_tokens(student_value) - flat_teacher = ( - self._flatten_tokens(teacher_values[depth]) if needs_mtp_kd else None - ) + flat_teacher = self._flatten_tokens(teacher_values[depth]) if needs_mtp_kd else None flat_labels = depth_labels.reshape(-1) for start in range(0, flat_student.shape[0], chunk_size): stop = min(start + chunk_size, flat_student.shape[0]) @@ -1205,8 +1258,12 @@ def _chunk_objectives(s_chunk, t_chunk, chunk_labels): phase = f"mtp_depth_{depth}_chunk_{start}_{stop}" _trace_global_kd_phase(f"{phase}_student_head_begin") if student_is_hidden: + if student_head is None: + raise RuntimeError("MTP hidden-state projection is missing its lm_head") s_chunk = _align_dtensor_to_module_mesh(s_chunk, student_head) - s_logits = student_head(s_chunk) if student_is_hidden else s_chunk + s_logits = student_head(s_chunk) + else: + s_logits = s_chunk _trace_global_kd_phase(f"{phase}_student_head_end") zero = self._local_zero(s_logits) _trace_global_kd_phase(f"{phase}_ce_begin") @@ -1273,7 +1330,10 @@ def _objective_loss(self, student_out, teacher_out, labels, model, num_label_tok rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 if rank == 0: placements = ( - tuple(type(item).__name__ + f"({getattr(item, 'dim', '')})" for item in student_logits.placements) + tuple( + type(item).__name__ + f"({getattr(item, 'dim', '')})" + for item in student_logits.placements + ) if isinstance(student_logits, DTensor) else ("replicated_tensor",) ) @@ -1387,9 +1447,7 @@ def loss_wrapper(student_out, target, **_kwargs): if teacher_out is None: raise RuntimeError("Teacher PP output queue is empty") model = next( - part - for part, stage in zip(self.model_parts, self.pp.info.stages) - if stage.is_last + part for part, stage in zip(self.model_parts, self.pp.info.stages) if stage.is_last ) # PP schedules rescale by the optimizer-step label count after # backward, so every microbatch loss must remain an unnormalized sum. @@ -1401,9 +1459,7 @@ def loss_wrapper(student_out, target, **_kwargs): return loss_wrapper -class KnowledgeDistillationRecipeForNextTokenPrediction( - _WeightedObjectiveMixin, _AutoModelLLMKD -): +class KnowledgeDistillationRecipeForNextTokenPrediction(_WeightedObjectiveMixin, _AutoModelLLMKD): """AutoModel LLM KD with independently weighted main/MTP objectives.""" def setup(self): @@ -1476,7 +1532,9 @@ def _forward_backward_step( is_train=is_train, loss_buffer=loss_buffer, ) - batch = {key: value.to(self.dist_env.device, non_blocking=True) for key, value in batch.items()} + batch = { + key: value.to(self.dist_env.device, non_blocking=True) for key, value in batch.items() + } # Current AutoModel CP owns label sharding through the batch mapping. # Keep labels present until CP has padded/sharded every sequence tensor, # then remove the CP-local labels for the weighted objective. @@ -1496,7 +1554,9 @@ def _forward_backward_step( teacher_out = None if self.needs_teacher: with ScopedModuleOffloading(self.teacher_model, enabled=False), torch.no_grad(): - teacher_out = self.teacher_model(**filter_forward_kwargs(self.teacher_model, batch)) + teacher_out = self.teacher_model( + **filter_forward_kwargs(self.teacher_model, batch) + ) student_out = model(**filter_forward_kwargs(model, batch)) total, terms = self._objective_loss( student_out, teacher_out, labels, model, num_label_tokens @@ -1739,7 +1799,9 @@ def _forward_backward_step( with torch.no_grad(): prepared = model(_pre_embed_only=True, **media) if self.needs_teacher and "inputs_embeds" in prepared: - _validate_cp_pre_embed_teacher_compatibility(prepared["inputs_embeds"], self.teacher_model) + _validate_cp_pre_embed_teacher_compatibility( + prepared["inputs_embeds"], self.teacher_model + ) for key in VLM_INPUT_KEYS: batch.pop(key, None) batch.update(prepared) @@ -1796,9 +1858,7 @@ def prepare_cp_inputs(pp, parts, values): batch = prepare_cp_inputs(self.pp, self.model_parts, batch) if self.needs_teacher: - teacher_batch = prepare_cp_inputs( - self.teacher_pp, self.teacher_pp.parts, teacher_batch - ) + teacher_batch = prepare_cp_inputs(self.teacher_pp, self.teacher_pp.parts, teacher_batch) train_ctx, batch = make_cp_batch_and_ctx(self.device_mesh, batch) labels = batch.pop("labels") model_input_key = "inputs_embeds" if "inputs_embeds" in batch else "input_ids" @@ -1824,9 +1884,7 @@ def prepare_cp_inputs(pp, parts, values): if self.needs_teacher: teacher_ctx, teacher_batch = make_cp_batch_and_ctx(self.device_mesh, teacher_batch) teacher_labels = teacher_batch.pop("labels") - teacher_input_key = ( - "inputs_embeds" if "inputs_embeds" in teacher_batch else "input_ids" - ) + teacher_input_key = "inputs_embeds" if "inputs_embeds" in teacher_batch else "input_ids" teacher_input = teacher_batch.pop(teacher_input_key) with teacher_ctx(): teacher_targets = ( @@ -1837,8 +1895,9 @@ def prepare_cp_inputs(pp, parts, values): set_pp_vlm_chunk_specs(self.teacher_pp.info.schedule, teacher_batch) capture = self.teacher_model._teacher_logits_capture capture.clear() - with torch.no_grad(), stage_vlm_media_for_pp( - self.teacher_pp, self.teacher_pp.parts, teacher_batch + with ( + torch.no_grad(), + stage_vlm_media_for_pp(self.teacher_pp, self.teacher_pp.parts, teacher_batch), ): teacher_losses = [] if self.teacher_pp.info.has_last_stage else None if self.teacher_pp.info.has_first_stage: @@ -1880,9 +1939,7 @@ def prepare_cp_inputs(pp, parts, values): ) def _run_train_optim_step(self, batches, max_grad_norm=None): - log_data = FinetuneRecipeForVLM._run_train_optim_step( - self, batches, max_grad_norm - ) + log_data = FinetuneRecipeForVLM._run_train_optim_step(self, batches, max_grad_norm) # The shared publisher normalizes last-stage PP microbatch sums and # forwards every objective term to rank zero. diff --git a/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py b/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py index 531070d2620..5df3f7df23a 100644 --- a/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py +++ b/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py @@ -1,3 +1,18 @@ +# 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. + """Long-lived AutoModel executor for distributed replace-block evaluation.""" from __future__ import annotations @@ -5,6 +20,7 @@ import time from contextlib import ExitStack, nullcontext from pathlib import Path +from typing import Any from ..anymodel.model_descriptor import ModelDescriptorFactory from ..anymodel.registry import resolve_descriptor_from_pretrained @@ -34,17 +50,19 @@ class AutoModelReplaceBlockExecutor: def __init__(self, hydra_cfg): self.cfg = hydra_cfg - self.recipe = None - self.cache = None - self.params = None - self.teacher_block_configs = None - self.num_q = None - self.head_dim = None - self.bypass_checkpoint_dir = None + self.recipe: Any | None = None + self.cache: Any | None = None + self.params: dict[str, Any] | None = None + self.teacher_block_configs: Any | None = None + self.num_q: int | None = None + self.head_dim: int | None = None + self.bypass_checkpoint_dir: Path | None = None self.is_output_writer = False - self.source_hidden_width = None - self.sliced_teacher_baseline = None - self.latest_observability = None + self.source_hidden_width: int | None = None + self.sliced_teacher_baseline: dict[str, Any] | None = None + self.latest_observability: dict[str, Any] | None = None + self.latest_score_device_type: str | None = None + self.visible_cuda_device_count: int | None = None self._setup_complete = False def capabilities(self) -> dict: @@ -74,7 +92,8 @@ def setup(self) -> None: from ..tools.checkpoint_utils import load_model_config scoring = self.cfg.scoring - self.params = solution_scoring_params(self.cfg) + params = solution_scoring_params(self.cfg) + self.params = params apply_patch() teacher_dir = Path( scoring.get("teacher_dir", None) or f"{self.cfg.puzzle_dir}/ckpts/teacher" @@ -115,18 +134,19 @@ def setup(self) -> None: recipe_dict = build_solution_recipe_config(self.cfg, target_dir) distributed = recipe_dict.get("distributed", {}) validate_force_hf_ep( - self.params["force_hf"], + params["force_hf"], int(distributed.get("ep_size", 1) or 1), ) target_recipe = _run_recipe( recipe_dict, scoring, - self.params["eval_iters"], - self.params["use_puzzletron_dataloader"], - self.params["data_cfg"], + params["eval_iters"], + params["use_puzzletron_dataloader"], + params["data_cfg"], ) - self.cache = TeacherTargetCache(device=self.params["teacher_cache_device"]) - _extract_teacher_targets(target_recipe, self.cache, self.params) + cache = TeacherTargetCache(device=params["teacher_cache_device"]) + self.cache = cache + _extract_teacher_targets(target_recipe, cache, params) dist.barrier() if source_dir.resolve() == target_dir.resolve(): @@ -137,9 +157,9 @@ def setup(self) -> None: self.recipe = _run_recipe( build_solution_recipe_config(self.cfg, source_dir), scoring, - self.params["eval_iters"], - self.params["use_puzzletron_dataloader"], - self.params["data_cfg"], + params["eval_iters"], + params["use_puzzletron_dataloader"], + params["data_cfg"], ) # AutoModel PP containers can retain a final norm/LM head on a rank # where the pipeline stage does not actually execute them. Elect from @@ -149,12 +169,15 @@ def setup(self) -> None: import torch.distributed as torch_dist rank = torch_dist.get_rank() if torch_dist.is_initialized() else 0 - observed = bool(len(self.cache)) - observed_by_rank = [(rank, observed)] + observed = bool(len(cache)) + observed_by_rank: list[tuple[int, bool] | None] = [(rank, observed)] if torch_dist.is_initialized(): observed_by_rank = [None] * torch_dist.get_world_size() torch_dist.all_gather_object(observed_by_rank, (rank, observed)) - output_ranks = [item_rank for item_rank, has_capture in observed_by_rank if has_capture] + observations = [item for item in observed_by_rank if item is not None] + if len(observations) != len(observed_by_rank): + raise RuntimeError("Missing AutoModel output-rank observation from a distributed rank") + output_ranks = [item_rank for item_rank, has_capture in observations if has_capture] if not output_ranks: raise RuntimeError("No AutoModel rank captured teacher final hidden states") self.is_output_writer = observed and rank == min(output_ranks) @@ -172,6 +195,9 @@ def evaluate(self, request: EvaluationRequest) -> EvaluationResult | None: raise NotImplementedError(f"Unsupported evaluation handler {request.handler!r}") if not self._setup_complete: raise RuntimeError("AutoModelReplaceBlockExecutor.setup() was not called") + params = self.params + if params is None: + raise RuntimeError("AutoModelReplaceBlockExecutor setup state is incomplete") from ..plugins.automodel.solution_launch import _solution_prune_target from ..replacement_library.replacement_utils import parse_layer_replacement @@ -229,8 +255,8 @@ def evaluate(self, request: EvaluationRequest) -> EvaluationResult | None: f"seconds={time.perf_counter() - started:.3f}", flush=True, ) - counts = { - name: len(value.get("per_sample", [])) + counts: dict[str, int | float] = { + str(name): len(value.get("per_sample", [])) for name, value in metrics.items() if isinstance(value, dict) } @@ -243,14 +269,18 @@ def evaluate(self, request: EvaluationRequest) -> EvaluationResult | None: provenance={ "handler": request.handler, "evaluator_revision": request.evaluator_revision, - "micro_batch_size": self.params.get("micro_batch_size"), + "micro_batch_size": params.get("micro_batch_size"), "hidden_width": self.source_hidden_width, "sliced_teacher_baseline": self.sliced_teacher_baseline, "observability": self.latest_observability, + "score_device_type": getattr(self, "latest_score_device_type", None), + "visible_cuda_device_count": getattr(self, "visible_cuda_device_count", None), }, ) - def _score(self, prune_target: dict | list[dict] | None) -> dict | None: + def _score(self, prune_target: dict | list[dict] | None) -> dict[str, Any] | None: + # Keep framework imports lazy so this executor remains dependency-light until setup. + import torch import torch.distributed as torch_dist import modelopt.torch.utils.distributed as dist @@ -264,10 +294,17 @@ def _score(self, prune_target: dict | list[dict] | None) -> dict | None: recipe = self.recipe cache = self.cache params = self.params + if recipe is None or cache is None or params is None: + raise RuntimeError("AutoModelReplaceBlockExecutor setup state is incomplete") + self.visible_cuda_device_count = torch.cuda.device_count() per_batch = [] tp_group = recipe.tensor_parallel_group() candidate_lm_head = recipe.lm_head_weight() if recipe.has_outputs else None - raw_targets = prune_target if isinstance(prune_target, list) else [prune_target] + raw_targets: list[dict | None] = [] + if isinstance(prune_target, list): + raw_targets.extend(prune_target) + else: + raw_targets.append(prune_target) prune_targets = [dict(target) for target in raw_targets if target is not None] layer_indices = [int(target["layer_idx"]) for target in prune_targets] owned_layers = [ @@ -279,9 +316,7 @@ def _score(self, prune_target: dict | list[dict] | None) -> dict | None: with ExitStack() as stack: for target in prune_targets: layer_idx = int(target["layer_idx"]) - bypass_dir = target.pop( - "bypass_checkpoint_dir", self.bypass_checkpoint_dir - ) + bypass_dir = target.pop("bypass_checkpoint_dir", self.bypass_checkpoint_dir) stack.enter_context( recipe.block_checkpoint_overlay_context(bypass_dir, layer_idx) if bypass_dir is not None @@ -298,6 +333,15 @@ def _score(self, prune_target: dict | list[dict] | None) -> dict | None: for batch_idx, (hidden, targets) in enumerate(recipe.iterate_captures()): if hidden is None: continue + if candidate_lm_head is None: + raise RuntimeError("AutoModel scoring requires a candidate LM-head weight") + device_type = hidden.device.type + if self.latest_score_device_type not in (None, device_type): + raise RuntimeError( + "AutoModel scoring tensors changed device type within one executor: " + f"{self.latest_score_device_type!r} -> {device_type!r}" + ) + self.latest_score_device_type = device_type teacher_hidden = cache.hidden( batch_idx, device=hidden.device, diff --git a/modelopt/torch/puzzletron/orchestration/adapters/pool.py b/modelopt/torch/puzzletron/orchestration/adapters/pool.py index 1502a503a21..07e6364afa9 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/pool.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/pool.py @@ -18,7 +18,6 @@ from __future__ import annotations from pathlib import Path -from typing import cast from ..identity import stable_hash from ..schema import ( @@ -68,6 +67,14 @@ def _replacement_puzzle_dir(plan: CampaignPlan, width: int | None) -> Path: return plan.puzzle_dir / "scenarios" / f"width-{int(width):04d}" / "depth-00" +def _replacement_work_id(stage_id: str, width: int | None, width_count: int) -> str: + if width_count == 1: + return f"{stage_id}:gang" + if width is None: + raise RuntimeError("multi-width replacement scoring requires concrete widths") + return f"{stage_id}:width-{width:04d}" + + def _replacement_environment(plan: CampaignPlan, puzzle_dir: Path) -> dict[str, str]: scoring = plan.experiment_config.get("replacement_scoring") or {} granularity = str(scoring.get("granularity", "block")) @@ -125,16 +132,16 @@ def _replacement_overrides(plan: CampaignPlan, puzzle_dir: Path) -> tuple[str, . ) teacher = puzzle_dir / "ckpts" / "sorted_teacher" overrides = [ - f"puzzle_dir={puzzle_dir}", + f"++puzzle_dir={puzzle_dir}", f"experiment.dir={puzzle_dir}", f"teacher_dir={teacher}", f"convert.teacher_dir={teacher}", "bypass.enabled=false", f"replacement_library_path={puzzle_dir / 'replacement_library.json'}", - f"build_replacement_library.source_checkpoint_dir={teacher}", + f"++build_replacement_library.source_checkpoint_dir={teacher}", f"replacement_scoring.teacher_dir={teacher}", - f"replacement_scoring.source_checkpoint_dir={teacher}", - f"replacement_scoring.target_teacher_dir={teacher}", + f"++replacement_scoring.source_checkpoint_dir={teacher}", + f"++replacement_scoring.target_teacher_dir={teacher}", f"replacement_scoring.solutions_path={puzzle_dir / f'{stem}.json'}", f"replacement_scoring.output_dir={puzzle_dir / f'{stem}--validation'}", ] @@ -178,11 +185,7 @@ def plan(self, plan: CampaignPlan, node: StagePlanNode) -> WorkPlan: workers_per_width, remainder = divmod(node.instances, len(widths)) items = tuple( WorkItem( - work_id=( - f"{node.stage_id}:gang" - if len(widths) == 1 - else f"{node.stage_id}:width-{cast('int', width):04d}" - ), + work_id=_replacement_work_id(node.stage_id, width, len(widths)), stage_id=node.stage_id, shard_index=index, shard_count=len(widths), @@ -271,12 +274,15 @@ def command( effective_overrides.extend(_replacement_overrides(plan, replacement_puzzle_dir)) if role == "gang": worker_count = int(item.metadata.get("worker_count", node.instances)) + allocation_nodes, allocation_gpus, topology = packed_allocation( + node, instances=worker_count + ) env = { "CAMPAIGN_DIR": str(campaign_dir), "CONFIG_PATH": plan.experiment_config_path, "PUZZLE_DIR": str(replacement_puzzle_dir), "WORLD_SIZE": str(node.gpus_per_instance), - "NPROC_PER_NODE": str(node.gpus_per_instance), + "NPROC_PER_NODE": str(topology.gpus_per_task), "WORKER_COUNT": str(worker_count), } if node.stage_id == "depth_importance": @@ -309,9 +315,6 @@ def command( existing = env.get("DISTRIBUTED_EVAL_OVERRIDES", "") env["DISTRIBUTED_EVAL_OVERRIDES"] = f"{existing}\n{override}".strip() log_path = str(log_dir / f"{node.stage_id}_gang_{attempt_id}.log") - allocation_nodes, allocation_gpus, topology = packed_allocation( - node, instances=worker_count - ) return AttemptSpec( attempt_id=attempt_id, work_id=item.work_id, @@ -367,13 +370,9 @@ def command( if role == "worker": env["CUDA_VISIBLE_DEVICES"] = ",".join(str(gpu) for gpu in item.local_gpu_ids) env["NPROC_PER_NODE"] = str(node.gpus_per_instance) - env["NNODES"] = "1" - env["NODE_RANK"] = "0" worker_id = int(item.metadata.get("worker_id", 0)) env["WORKER_GROUP_INDEX"] = str(worker_id) env["WORKER_PORT"] = str(5010 + worker_id) - env["RDZV_ENDPOINT"] = f"127.0.0.1:{29500 + worker_id}" - env["RDZV_ID"] = f"{node.stage_id}-{attempt_id}" argv = ("bash", str(script)) # GPU partitions reject zero-GPU jobs; coordinators still need one GPU slot. allocation_gpus = 1 if role == "coordinator" else node.gpus_per_instance diff --git a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py index cce0df7fc18..692859cbace 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py @@ -1,16 +1,29 @@ # 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. """Sharded stage adapter for independent worker instances.""" from __future__ import annotations -import subprocess import time import uuid from dataclasses import replace from pathlib import Path +from typing import TYPE_CHECKING +from ..executors.local import LocalExecutor from ..executors.slurm import SlurmExecutor from ..schema import ( AttemptSpec, @@ -30,9 +43,16 @@ from .packing import packed_allocation from .stage_compat import stage_is_complete, stage_output_patterns +if TYPE_CHECKING: + from ...security_policy import require_boolean_policy +elif __package__.startswith("puzzletron_orchestrator."): + from puzzletron_orchestrator.security_policy import require_boolean_policy +else: + from ...security_policy import require_boolean_policy + __all__ = ["ShardedStageAdapter"] -_SHARDED_ENTRYPOINTS = { +_SHARDED_ENTRYPOINTS: dict[str, tuple[str, list[str]]] = { "vllm_stats": ("examples/puzzletron/run_runtime_stats_shard.py", []), "zero_shot_evaluation": ( "examples/puzzletron/run_profile_online_evaluation.py", @@ -109,6 +129,46 @@ def _run_slurm_aggregate( time.sleep(2) +def _run_local_aggregate( + *, + plan: CampaignPlan, + node: StagePlanNode, + command: tuple[str, ...], +) -> None: + """Run a controller-side merge through the reviewed local executor.""" + + attempt_id = str(uuid.uuid4()) + log_path = plan.puzzle_dir / "logs" / f"{node.stage_id}_merge_{attempt_id}.log" + attempt = AttemptSpec( + attempt_id=attempt_id, + work_id=f"{node.stage_id}:aggregate", + stage_id=f"{node.stage_id}_merge", + command=CommandSpec( + argv=command, + cwd=plan.runner.contract.repository, + log_path=str(log_path), + ), + allocation_nodes=1, + allocation_gpus=0, + contract_hash=plan.contract_hash, + metadata={"gpus_per_node": 0}, + task_topology=TaskTopology(task_count=1, gpus_per_task=0), + ) + executor = LocalExecutor() + handle = executor.submit(attempt) + while True: + status = executor.poll([handle])[0] + if status.state is JobState.COMPLETED: + return + if status.state not in {JobState.PENDING, JobState.RUNNING}: + detail = _read_log_tail(str(log_path)) + raise RuntimeError( + f"{node.stage_id} aggregation {status.state.value}: " + f"{detail or status.reason or 'no log output'}" + ) + time.sleep(0.1) + + def _gang_item( node: StagePlanNode, *, @@ -217,11 +277,11 @@ def command( measurement_id = item.metadata.get("measurement_id") name_suffix = f"_{measurement_id}" if measurement_id else "" log_path = str( - log_dir - / f"{node.stage_id}{name_suffix}_shard{item.shard_index}_{attempt_id}.log" + log_dir / f"{node.stage_id}{name_suffix}_shard{item.shard_index}_{attempt_id}.log" ) if node.stage_id == "aiperf": aiperf = plan.experiment_config.get("aiperf") or {} + model = plan.experiment_config.get("model") or {} argv = [ "python", str(script_path), @@ -236,6 +296,26 @@ def command( "--output-tokens", str(aiperf.get("output_tokens", 1024)), ] + if "trust_remote_code" in aiperf: + trust_remote_code_path = "aiperf.trust_remote_code" + trust_remote_code_value = aiperf["trust_remote_code"] + else: + trust_remote_code_path = "model.trust_remote_code" + trust_remote_code_value = model.get("trust_remote_code", False) + trust_remote_code = require_boolean_policy( + trust_remote_code_value, + path=trust_remote_code_path, + default=False, + ) + allow_online_tokenizer_resolution = require_boolean_policy( + aiperf.get("allow_aiperf_v011_online_tokenizer_resolution", False), + path="aiperf.allow_aiperf_v011_online_tokenizer_resolution", + default=False, + ) + if trust_remote_code: + argv.append("--trust-remote-code") + if allow_online_tokenizer_resolution: + argv.append("--allow-aiperf-v011-online-tokenizer-resolution") else: argv = ["python", str(script_path), "--config", plan.experiment_config_path] argv.extend(extra_args) @@ -353,18 +433,7 @@ def aggregate( command=command, ) else: - result = subprocess.run( - command, - cwd=plan.runner.contract.repository, - capture_output=True, - text=True, - check=False, - ) - if result.returncode: - raise RuntimeError( - f"{node.stage_id} aggregation failed: " - f"{result.stderr.strip() or result.stdout.strip()}" - ) + _run_local_aggregate(plan=plan, node=node, command=command) return PublishedOutput( stage_id=node.stage_id, artifacts=stage_output_patterns(plan.experiment_config, node.stage_id), diff --git a/modelopt/torch/puzzletron/orchestration/task_launcher.py b/modelopt/torch/puzzletron/orchestration/task_launcher.py index 4d4b21c414a..6245435e902 100644 --- a/modelopt/torch/puzzletron/orchestration/task_launcher.py +++ b/modelopt/torch/puzzletron/orchestration/task_launcher.py @@ -1,5 +1,17 @@ # 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. """Dependency-light launcher for one task in an orchestration attempt.""" @@ -22,6 +34,7 @@ "TaskBinding", "build_task_command", "main", + "rendezvous_endpoint", "rendezvous_port", "resolve_task_binding", ] @@ -38,6 +51,7 @@ "PUZZLETRON_LOCAL_TASK_INDEX", "PUZZLETRON_MASTER_ADDR", "PUZZLETRON_MASTER_PORT", + "PUZZLETRON_RENDEZVOUS_ENDPOINT", "PUZZLETRON_RENDEZVOUS_ID", "PUZZLETRON_TASK_HOSTS", "PUZZLETRON_TASK_INDEX", @@ -76,6 +90,14 @@ def rendezvous_port(attempt_id: str, group_index: int, group_count: int) -> int: return port_start + (seed % (port_span - group_count + 1)) + group_index +def rendezvous_endpoint(binding: TaskBinding) -> str: + """Return a collision-free local endpoint or the shared multi-node endpoint.""" + + if binding.group_size == 1: + return "localhost:0" + return f"{binding.master_addr}:{binding.master_port}" + + def resolve_task_binding( *, attempt_id: str, @@ -131,7 +153,7 @@ def build_task_command( command = tuple(str(part) for part in payload) if launcher is TaskLauncher.DIRECT: return command - rendezvous_host = "localhost" if binding.group_size == 1 else binding.master_addr + endpoint = rendezvous_endpoint(binding) return ( "python", "-m", @@ -139,7 +161,7 @@ def build_task_command( f"--nnodes={binding.group_size}", f"--nproc-per-node={gpus_per_task}", "--rdzv-backend=c10d", - f"--rdzv-endpoint={rendezvous_host}:{binding.master_port}", + f"--rdzv-endpoint={endpoint}", f"--rdzv-id={binding.rendezvous_id}", "--no-python", *command, @@ -203,9 +225,7 @@ def main(argv: Sequence[str] | None = None) -> int: f"task {task_index} expected {args.gpus_per_task} visible GPUs, got {visible_gpus}" ) tasks_per_node = ( - args.task_count - if args.gpus_per_task == 0 - else args.gpus_per_node // args.gpus_per_task + args.task_count if args.gpus_per_task == 0 else args.gpus_per_node // args.gpus_per_task ) expected_hosts = math.ceil(args.task_count / tasks_per_node) if expected_hosts > args.nodes: @@ -234,6 +254,7 @@ def main(argv: Sequence[str] | None = None) -> int: PUZZLETRON_GROUP_SIZE=str(binding.group_size), PUZZLETRON_MASTER_ADDR=binding.master_addr, PUZZLETRON_MASTER_PORT=str(binding.master_port), + PUZZLETRON_RENDEZVOUS_ENDPOINT=rendezvous_endpoint(binding), PUZZLETRON_RENDEZVOUS_ID=binding.rendezvous_id, ) print( @@ -241,7 +262,7 @@ def main(argv: Sequence[str] | None = None) -> int: f"host={binding.hostname} task={binding.task_index} " f"local={binding.local_task_index} gpus={','.join(visible_gpus)} " f"group={binding.group_index} rank={binding.group_rank}/{binding.group_size} " - f"endpoint={binding.master_addr}:{binding.master_port} " + f"endpoint={rendezvous_endpoint(binding)} " f"rdzv_id={binding.rendezvous_id}", flush=True, ) @@ -251,8 +272,13 @@ def main(argv: Sequence[str] | None = None) -> int: binding=binding, gpus_per_task=args.gpus_per_task, ) - os.execvpe(command[0], command, env) - return 0 + # The command is explicit argv from the compiled campaign plan, and no shell + # is involved. The payload inherits this task's process group so scheduler + # group signals reach both launcher and payload. + pid = os.posix_spawnp(command[0], command, env) + _, status = os.waitpid(pid, 0) + exit_code = os.waitstatus_to_exitcode(status) + return exit_code if exit_code >= 0 else 128 - exit_code if __name__ == "__main__": diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 159ab200e9b..a040b8c0f53 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -31,6 +31,7 @@ from ..evaluation import DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS, run_lmms_eval_checkpoint from ..identity import canonicalize, stable_hash +from ..security_policy import require_boolean_policy from .base import CompiledPostMIPNode, NodeKind, compile_post_mip_flows from .filters import apply_filter from .identity import ( @@ -445,6 +446,19 @@ def _aiperf( for concurrency in concurrencies } topology = dict(settings.pop("topology", {}) or {}) + trust_remote_code = require_boolean_policy( + settings.pop( + "trust_remote_code", + (config.get("model") or {}).get("trust_remote_code", False), + ), + path="post_mip.aiperf.config.trust_remote_code", + default=False, + ) + allow_online_tokenizer_resolution = require_boolean_policy( + settings.pop("allow_aiperf_v011_online_tokenizer_resolution", False), + path="post_mip.aiperf.config.allow_aiperf_v011_online_tokenizer_resolution", + default=False, + ) gpu_ids = os.environ.get("CUDA_VISIBLE_DEVICES", "") if not gpu_ids: gpu_ids = ",".join(str(index) for index in range(int(topology.get("gpu_group_size", 1)))) @@ -461,6 +475,8 @@ def _aiperf( request_counts=request_counts, solution_id=source.architecture_id, profile_id=node.flow_id, + trust_remote_code=trust_remote_code, + allow_aiperf_v011_online_tokenizer_resolution=allow_online_tokenizer_resolution, **settings, ) metrics = {} diff --git a/modelopt/torch/puzzletron/security_policy.py b/modelopt/torch/puzzletron/security_policy.py new file mode 100644 index 00000000000..529f1589ff3 --- /dev/null +++ b/modelopt/torch/puzzletron/security_policy.py @@ -0,0 +1,36 @@ +# 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. + +"""Validation helpers for security-sensitive Puzzletron configuration.""" + +from typing import Any + +__all__ = ["require_boolean_policy"] + + +def require_boolean_policy( + value: Any, + *, + path: str, + default: bool | None = None, +) -> bool: + """Return a policy boolean, resolving ``None`` only to an explicit default.""" + if default is not None and not isinstance(default, bool): + raise ValueError(f"{path} default must be a boolean") + if value is None and default is not None: + return default + if not isinstance(value, bool): + raise ValueError(f"{path} must be a boolean") + return value diff --git a/modelopt/torch/puzzletron/stages/future.py b/modelopt/torch/puzzletron/stages/future.py index 39253eb58b8..1682f3a099c 100644 --- a/modelopt/torch/puzzletron/stages/future.py +++ b/modelopt/torch/puzzletron/stages/future.py @@ -23,7 +23,7 @@ from dataclasses import asdict from pathlib import Path from queue import Queue -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Sequence from ..anymodel.model_descriptor import ModelDescriptorFactory from ..anymodel.registry import resolve_descriptor_from_pretrained @@ -33,6 +33,7 @@ build_global_kd_config, run_global_kd, ) +from ..security_policy import require_boolean_policy from .common import complete_stage from .graph import StageSkipReason @@ -197,11 +198,11 @@ def _select_evaluated_candidates( def _with_teacher_checkpoint( - teacher_dir: str | Path | None, candidates: list[tuple[str, str | Path]] + teacher_dir: str | Path | None, candidates: Sequence[tuple[str, str | Path]] ) -> list[tuple[str, str | Path]]: """Prepend the configured teacher while keeping downstream checkpoints unique.""" if teacher_dir is None: - return candidates + return list(candidates) teacher_key = str(teacher_dir) return [("teacher", teacher_dir)] + [ (name, checkpoint) @@ -300,9 +301,23 @@ def aiperf_stage(config: dict[str, Any], manifest: StageManifest): skip_reason=StageSkipReason.DISABLED, message="AIPerf is disabled.", ) + model_cfg = dict(config.get("model") or {}) + trust_remote_code = require_boolean_policy( + stage_cfg.get("trust_remote_code", model_cfg.get("trust_remote_code", False)), + path="aiperf.trust_remote_code", + default=False, + ) + allow_aiperf_v011_online_tokenizer_resolution = require_boolean_policy( + stage_cfg.get("allow_aiperf_v011_online_tokenizer_resolution", False), + path="aiperf.allow_aiperf_v011_online_tokenizer_resolution", + default=False, + ) from ..benchmarks import run_aiperf_sweep, write_aiperf_report - puzzle_dir = Path((config.get("experiment") or {}).get("dir")) + experiment_dir = (config.get("experiment") or {}).get("dir") + if experiment_dir is None: + raise ValueError("AIPerf requires experiment.dir") + puzzle_dir = Path(experiment_dir) teacher_dir = (config.get("convert") or {}).get("teacher_dir") checkpoint_root = Path( stage_cfg.get( @@ -310,7 +325,7 @@ def aiperf_stage(config: dict[str, Any], manifest: StageManifest): puzzle_dir / "mip" / "puzzle_solutions" / "depth_tournament" / "solutions--checkpoints", ) ) - checkpoints = [] + checkpoints: list[tuple[str, str | Path]] = [] if stage_cfg.get("checkpoint_source") == "global_kd": checkpoints.extend(_scenario_grid_global_kd_checkpoints(puzzle_dir)) elif stage_cfg.get("checkpoint_source") == "scenario_grid": @@ -414,6 +429,10 @@ def _run_checkpoint(item): extra_inputs=dict(stage_cfg.get("extra_inputs") or {}), use_server_token_count=bool(stage_cfg.get("use_server_token_count", True)), seed=int(stage_cfg.get("seed", 42)), + trust_remote_code=trust_remote_code, + allow_aiperf_v011_online_tokenizer_resolution=( + allow_aiperf_v011_online_tokenizer_resolution + ), ) finally: pool.put(gpu_ids) @@ -439,6 +458,9 @@ def evaluation_stage(config: dict[str, Any], manifest: StageManifest): skip_reason=StageSkipReason.DISABLED, message="Zero-shot evaluation is disabled.", ) + configured = stage_cfg.get("checkpoints") + if configured is not None and not isinstance(configured, (list, tuple)): + raise ValueError("zero_shot_evaluation.checkpoints must be a list or tuple") from omegaconf import OmegaConf import modelopt.torch.utils.distributed as dist @@ -452,21 +474,21 @@ def evaluation_stage(config: dict[str, Any], manifest: StageManifest): from .pipeline import _distributed puzzle_dir = Path((config.get("experiment") or {})["dir"]) - configured = stage_cfg.get("checkpoints") + raw_checkpoint_entries: list[tuple[str, str | Path]] if configured: - checkpoint_entries = [(Path(path).name, Path(path)) for path in configured] + raw_checkpoint_entries = [(Path(path).name, Path(path)) for path in configured] elif stage_cfg.get("checkpoint_source") == "global_kd": - checkpoint_entries = _scenario_grid_global_kd_checkpoints(puzzle_dir) + raw_checkpoint_entries = _scenario_grid_global_kd_checkpoints(puzzle_dir) else: - checkpoint_entries = [ + raw_checkpoint_entries = [ (name, checkpoint) for name, checkpoint in _profile_solution_checkpoints( puzzle_dir, stage_cfg.get("profile_id") ) if name != "teacher" ] - if not checkpoint_entries: - checkpoint_entries = [ + if not raw_checkpoint_entries: + raw_checkpoint_entries = [ (path.parent.name, path.parent) for path in sorted( (puzzle_dir / "scenarios").glob( @@ -477,7 +499,7 @@ def evaluation_stage(config: dict[str, Any], manifest: StageManifest): teacher_dir = (config.get("convert") or {}).get("teacher_dir") checkpoint_entries = [ (name, Path(checkpoint)) - for name, checkpoint in _with_teacher_checkpoint(teacher_dir, checkpoint_entries) + for name, checkpoint in _with_teacher_checkpoint(teacher_dir, raw_checkpoint_entries) ] if len(checkpoint_entries) == 1 and teacher_dir is None: raise FileNotFoundError("exact evaluation found no scenario checkpoints") @@ -496,7 +518,9 @@ def evaluation_stage(config: dict[str, Any], manifest: StageManifest): trust_remote_code=descriptor.requires_trust_remote_code(), ) lm = descriptor.get_language_model_config(checkpoint_config) - blocks = list(maybe_cast_block_configs(checkpoint_config.block_configs)) + blocks = maybe_cast_block_configs(checkpoint_config.block_configs) + if not blocks: + raise ValueError(f"{checkpoint} does not define Puzzletron block_configs") replacements = [ { "weight_paths": [], @@ -727,7 +751,7 @@ def distillation_stage(config: dict[str, Any], manifest: StageManifest): hydra_cfg, recipe_runner=config.get("_global_kd_runner"), ) - return complete_stage(config, manifest, outputs=result) + return complete_stage(config, manifest, outputs=result.to_dict()) kd_config = build_global_kd_config(config) result = run_global_kd(kd_config, recipe_runner=config.get("_global_kd_runner")) summary_path = kd_config.output_dir / "global_distillation_summary.json" diff --git a/modelopt/torch/puzzletron/utils/vllm_adapter.py b/modelopt/torch/puzzletron/utils/vllm_adapter.py index 3622fd0a44e..fa8aec78dc2 100644 --- a/modelopt/torch/puzzletron/utils/vllm_adapter.py +++ b/modelopt/torch/puzzletron/utils/vllm_adapter.py @@ -218,14 +218,21 @@ def configure_anymodel_metadata(hf_config: Any, descriptor: Any) -> bool: return True -def refresh_realized_checkpoint_config(checkpoint_dir: str | Path) -> Path: - """Rebuild the vLLM interchange fields of an already-realized checkpoint.""" +def refresh_realized_checkpoint_config( + checkpoint_dir: str | Path, + *, + trust_remote_code: bool = False, +) -> Path: + """Rebuild vLLM interchange fields without trusting checkpoint code by default.""" from transformers import AutoConfig from ..anymodel.registry import resolve_descriptor checkpoint_dir = Path(checkpoint_dir) - config = AutoConfig.from_pretrained(checkpoint_dir, trust_remote_code=True) + config = AutoConfig.from_pretrained( + checkpoint_dir, + trust_remote_code=trust_remote_code, + ) descriptor = resolve_descriptor(config).descriptor text_config = _get_text_config(config) @@ -338,22 +345,16 @@ def _convert_block_entry( is_full = _get(attn, "sliding_window_size") == "full" head_dim_field = "global_head_dim" if (is_full or k_eq_v) else "head_dim" current_head_dim = ( - global_global_head_dim - if head_dim_field == "global_head_dim" - else global_head_dim + global_global_head_dim if head_dim_field == "global_head_dim" else global_head_dim ) if qk_head_dim != current_head_dim: entry[head_dim_field] = qk_head_dim window = _get(attn, "sliding_window_size") if window is not None: - desired_type = ( - "full_attention" if window == "full" else "sliding_attention" - ) + desired_type = "full_attention" if window == "full" else "sliding_attention" current_type = ( - global_layer_types[layer_idx] - if layer_idx < len(global_layer_types) - else None + global_layer_types[layer_idx] if layer_idx < len(global_layer_types) else None ) if global_layer_types and current_type != desired_type: layer_types = list(global_layer_types) @@ -462,7 +463,9 @@ def _derive_per_layer_config( global_moe_latent_size = ( _get(text_config, moe_latent_field) if moe_latent_field is not None else None ) - global_mamba_values = {hf_field: _get(text_config, hf_field) for hf_field in mamba_fields.values()} + global_mamba_values = { + hf_field: _get(text_config, hf_field) for hf_field in mamba_fields.values() + } global_q_lora_rank = _get(text_config, "q_lora_rank") global_kv_lora_rank = _get(text_config, "kv_lora_rank") global_layer_types = list(_get(text_config, "layer_types") or ()) diff --git a/noxfile.py b/noxfile.py index 2f72b4181be..c572eb01cf3 100644 --- a/noxfile.py +++ b/noxfile.py @@ -58,6 +58,7 @@ with PUZZLETRON_V2_CI_ENVIRONMENT_PATH.open(encoding="utf-8") as environment_file: PUZZLETRON_V2_CI_ENVIRONMENT = json.load(environment_file) PUZZLETRON_V2_AUTOMODEL_SOURCE = PUZZLETRON_V2_CI_ENVIRONMENT["nemo_automodel"] +PUZZLETRON_V2_LMMS_SOURCE = PUZZLETRON_V2_CI_ENVIRONMENT["lmms_eval"] PUZZLETRON_V2_AUTOMODEL = ( "nemo-automodel @ git+" f"{PUZZLETRON_V2_AUTOMODEL_SOURCE['repository']}@" @@ -65,6 +66,55 @@ ) +def _verify_puzzletron_v2_environment(session): + """Fail before collection when the dedicated Puzzletron runtime drifts.""" + expected_versions = { + "python": PUZZLETRON_V2_CI_ENVIRONMENT["python"], + "torch": PUZZLETRON_V2_CI_ENVIRONMENT["torch"], + "torchvision": PUZZLETRON_V2_CI_ENVIRONMENT["torchvision"], + "transformers": PUZZLETRON_V2_CI_ENVIRONMENT["transformers"], + "lmms-eval": PUZZLETRON_V2_LMMS_SOURCE["base_version"], + "nemo-automodel": PUZZLETRON_V2_AUTOMODEL_SOURCE["base_version"], + } + expected_vcs = { + "lmms-eval": PUZZLETRON_V2_LMMS_SOURCE, + "nemo-automodel": PUZZLETRON_V2_AUTOMODEL_SOURCE, + } + session.run( + "python", + "-c", + f""" +import sys +from importlib.metadata import version + +from packaging.version import Version + +from examples.puzzletron.ci_environment import verify_installed_vcs_source + +expected = {expected_versions!r} +expected_vcs = {expected_vcs!r} +actual = {{ + "python": f"{{sys.version_info.major}}.{{sys.version_info.minor}}", + "torch": Version(version("torch")).base_version, + "torchvision": Version(version("torchvision")).base_version, + "transformers": Version(version("transformers")).base_version, + "lmms-eval": Version(version("lmms-eval")).base_version, + "nemo-automodel": Version(version("nemo-automodel")).base_version, +}} +mismatches = {{ + name: (actual[name], expected_version) + for name, expected_version in expected.items() + if actual[name] != expected_version +}} + +for name, source in expected_vcs.items(): + verify_installed_vcs_source(name, source) + +assert not mismatches, f"Pinned Puzzletron CI environment mismatch: {{mismatches}}" +""", + ) + + def _cov_args(): """Return --cov when COVERAGE_PROCESS_START is set (CI only).""" return ["--cov"] if os.environ.get("COVERAGE_PROCESS_START") else [] @@ -109,36 +159,7 @@ def puzzletron_v2(session): PUZZLETRON_V2_AUTOMODEL, ) session.run("uv", "pip", "check") - expected_versions = { - "python": PUZZLETRON_V2_CI_ENVIRONMENT["python"], - "torch": PUZZLETRON_V2_CI_ENVIRONMENT["torch"], - "torchvision": PUZZLETRON_V2_CI_ENVIRONMENT["torchvision"], - "transformers": PUZZLETRON_V2_CI_ENVIRONMENT["transformers"], - "lmms-eval": PUZZLETRON_V2_CI_ENVIRONMENT["lmms_eval"], - "nemo-automodel": PUZZLETRON_V2_AUTOMODEL_SOURCE["base_version"], - } - session.run( - "python", - "-c", - ( - "import sys; " - "from importlib.metadata import version; " - "from packaging.version import Version; " - f"expected = {expected_versions!r}; " - "actual = {" - "'python': f'{sys.version_info.major}.{sys.version_info.minor}', " - "'torch': Version(version('torch')).base_version, " - "'torchvision': Version(version('torchvision')).base_version, " - "'transformers': Version(version('transformers')).base_version, " - "'lmms-eval': Version(version('lmms-eval')).base_version, " - "'nemo-automodel': Version(version('nemo-automodel')).base_version}; " - "mismatches = {name: (actual[name], expected_version) " - "for name, expected_version in expected.items() " - "if actual[name] != expected_version}; " - "assert not mismatches, " - "f'Pinned Puzzletron CI environment mismatch: {mismatches}'" - ), - ) + _verify_puzzletron_v2_environment(session) session.run( "python", "-m", diff --git a/puzzletron_setup/bundle.py b/puzzletron_setup/bundle.py index 3ee3e8f0936..71d0b0c5bd8 100644 --- a/puzzletron_setup/bundle.py +++ b/puzzletron_setup/bundle.py @@ -470,6 +470,7 @@ def render_experiment(state: Mapping[str, Any], budget: str) -> dict[str, Any]: "modality": data["modality"], "layout": data["layout"], "max_sample_length": sequence_length, + "sequence_length": sequence_length, "path": data["source"], "acquisition": deepcopy(data_acquisition) if data_acquisition else None, "packing": ( @@ -854,7 +855,7 @@ def render_execution( pool_workers = int(workers.get("pool", 1)) sharded_workers = int(workers.get("sharded", 1)) embedding_widths = list(_mapping(experiment.get("embedding_pruning")).get("widths") or ()) - stages = { + stages: dict[str, dict[str, Any]] = { "convert": {"strategy": "single", "instances": 1, "parallel": single_gpu}, "tokenize_data": {"strategy": "single", "instances": 1}, "vllm_stats": { diff --git a/puzzletron_setup/v2/cli.py b/puzzletron_setup/v2/cli.py index bbce654d261..b3d7848739b 100644 --- a/puzzletron_setup/v2/cli.py +++ b/puzzletron_setup/v2/cli.py @@ -23,6 +23,9 @@ from puzzletron_setup import SetupError +from .presets import QUICK_SETUP_PRESETS +from .prompts import NonInteractiveBackend + if TYPE_CHECKING: from collections.abc import Sequence @@ -43,6 +46,22 @@ def _parser() -> argparse.ArgumentParser: type=Path, help="Explicit versioned defaults YAML; never discovered automatically.", ) + parser.add_argument( + "--campaign-dir", + type=Path, + help="Campaign directory for a new setup; bypasses that interactive prompt.", + ) + parser.add_argument( + "--profile", + choices=tuple(preset.name for preset in QUICK_SETUP_PRESETS), + default="balanced", + help="Guided setup profile for a new campaign (default: balanced).", + ) + parser.add_argument( + "--non-interactive", + action="store_true", + help="Accept resolved defaults and fail if any required answer has no default.", + ) parser.add_argument( "--full", action="store_true", @@ -56,7 +75,15 @@ def _parser() -> argparse.ArgumentParser: def main(argv: Sequence[str] | None = None) -> int: """Run the setup-v2 command-line interface.""" - args = _parser().parse_args(argv) + parser = _parser() + args = parser.parse_args(argv) + if args.resume is not None and args.campaign_dir is not None: + parser.error("--campaign-dir cannot be combined with --resume") + if args.non_interactive and args.resume is None: + if args.campaign_dir is None: + parser.error("--non-interactive requires --campaign-dir for a new campaign") + if args.defaults is None: + parser.error("--non-interactive requires --defaults for a new campaign") # Keep heavyweight model inspection out of --help and argument-error paths. from .wizard import run_wizard_v2 @@ -65,6 +92,9 @@ def main(argv: Sequence[str] | None = None) -> int: resume=args.resume, defaults_path=args.defaults, full=args.full, + campaign_dir=args.campaign_dir, + setup_profile=args.profile, + backend=NonInteractiveBackend() if args.non_interactive else None, ) except KeyboardInterrupt: target = args.resume or "" diff --git a/puzzletron_setup/v2/prompts.py b/puzzletron_setup/v2/prompts.py index 93e7386deca..b6bb55a3d06 100644 --- a/puzzletron_setup/v2/prompts.py +++ b/puzzletron_setup/v2/prompts.py @@ -29,6 +29,7 @@ __all__ = [ "BACK", "InteractiveBackend", + "NonInteractiveBackend", "PromptBackend", "PromptChoice", "ScriptedBackend", @@ -55,7 +56,7 @@ class PromptChoice: class PromptBackend(Protocol): """Minimal backend used by the navigable wizard session.""" - def text(self, message: str, default: str) -> Any: + def text(self, message: str, default: str | None) -> Any: """Request a text value.""" raise NotImplementedError @@ -127,10 +128,10 @@ class InteractiveBackend: _BACK_TITLE = "← Back" - def text(self, message: str, default: str) -> Any: + def text(self, message: str, default: str | None) -> Any: """Request text while supporting semantic Back navigation.""" print(" Press Esc to go back (or type :back).") - question = _bind_escape_back(_questionary().text(message, default=default)) + question = _bind_escape_back(_questionary().text(message, default=default or "")) value = _answer(question) if value is BACK: return BACK @@ -202,6 +203,60 @@ def checkbox( return list(values) +class NonInteractiveBackend: + """Accept resolved prompt defaults without depending on prompt ordering or labels.""" + + def text(self, message: str, default: str | None) -> Any: + """Return a resolved text default, including an explicit empty string.""" + if default is None: + raise SetupError(f"Non-interactive setup requires a default for {message!r}.") + return default + + def select( + self, + message: str, + choices: Sequence[PromptChoice], + default: Any, + ) -> Any: + """Return the enabled resolved choice default.""" + enabled = [choice for choice in choices if choice.disabled is None] + if default is None: + if len(enabled) == 1: + return enabled[0].value + raise SetupError(f"Non-interactive setup requires a default for {message!r}.") + selected = next((choice for choice in choices if choice.value == default), None) + if selected is None: + raise SetupError( + f"Non-interactive default {default!r} is not a choice for {message!r}." + ) + if selected.disabled is not None: + raise SetupError( + f"Non-interactive default {default!r} is unavailable for {message!r}: " + f"{selected.disabled}" + ) + return selected.value + + def checkbox( + self, + message: str, + choices: Sequence[PromptChoice], + defaults: Sequence[Any], + ) -> Any: + """Return the enabled resolved checkbox defaults.""" + for value in defaults: + choice = next((item for item in choices if item.value == value), None) + if choice is None: + raise SetupError( + f"Non-interactive default {value!r} is not a choice for {message!r}." + ) + if choice.disabled is not None: + raise SetupError( + f"Non-interactive default {value!r} is unavailable for {message!r}: " + f"{choice.disabled}" + ) + return list(defaults) + + class ScriptedBackend: """Deterministic non-interactive backend for embedding and automation.""" @@ -220,7 +275,7 @@ def _next(self) -> Any: value = self._answers.popleft() return BACK if value == ":back" else value - def text(self, message: str, default: str) -> Any: + def text(self, message: str, default: str | None) -> Any: """Return the next scripted text answer.""" del message, default return self._next() diff --git a/puzzletron_setup/v2/session.py b/puzzletron_setup/v2/session.py index a8f93772424..f05628e5ffe 100644 --- a/puzzletron_setup/v2/session.py +++ b/puzzletron_setup/v2/session.py @@ -19,7 +19,9 @@ from typing import TYPE_CHECKING, Any -from .prompts import BACK, InteractiveBackend, PromptBackend, PromptChoice +from puzzletron_setup import SetupError + +from .prompts import BACK, InteractiveBackend, NonInteractiveBackend, PromptBackend, PromptChoice from .state import PromptFrame, WizardState if TYPE_CHECKING: @@ -167,6 +169,8 @@ def text( if verdict is True: return rendered self.state.pop_frame() + if isinstance(self.backend, NonInteractiveBackend): + raise SetupError(str(verdict)) print(f" {verdict}") def integer( @@ -282,4 +286,6 @@ def disabled_verdict(values: Sequence[Any]) -> bool | str: if verdict is True: return rendered self.state.pop_frame() + if isinstance(self.backend, NonInteractiveBackend): + raise SetupError(str(verdict)) print(f" {verdict}") diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index 3c0a1935178..f0be4067f16 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -52,7 +52,7 @@ from .parallel_validation import validate_automodel_parallelism, validate_vllm_parallelism from .post_mip import FlowDraft, NodeDraft, PostMIPFlowEditor, recommended_flow from .presets import QUICK_SETUP_PRESETS, get_setup_preset -from .prompts import BACK, InteractiveBackend, PromptBackend, PromptChoice +from .prompts import BACK, InteractiveBackend, NonInteractiveBackend, PromptBackend, PromptChoice from .resources import ( ParallelProfile, ResourceProfileRegistry, @@ -134,6 +134,11 @@ ) +def _is_back(value: Any) -> bool: + """Return whether an untyped prompt result requests Back navigation.""" + return value is BACK + + def _model_family_choices(resolver: DefaultsResolver) -> list[PromptChoice]: choices = [] explicit = resolver.file_default("model.source") @@ -321,11 +326,16 @@ def _select_model_source( resolver: DefaultsResolver, ) -> Any: while True: + explicit_default = resolver.file_default("model.source") family = session.select( "model.source_family", "Model:", _model_family_choices(resolver), - default=_CUSTOM_MODEL_SOURCE, + default=( + _DEFAULT_MODEL_SOURCE + if explicit_default is not None and explicit_default.value + else _CUSTOM_MODEL_SOURCE + ), ) if family is BACK: return BACK @@ -363,6 +373,8 @@ def model_section(session: WizardSession, resolver: DefaultsResolver, context: d try: source = normalize_model_source(str(source)) except SetupError as error: + if isinstance(session.backend, NonInteractiveBackend): + raise print(f" {error}") continue session.state.set_field("model.source", source, source=source_kind) @@ -483,7 +495,7 @@ def data_section( modality_choices, default=suggested_modality, ) - if modality is BACK: + if _is_back(modality): return False modality_source = "user" else: @@ -544,7 +556,7 @@ def data_section( seed_default, minimum=0, ) - if seed is BACK: + if _is_back(seed): return False acquisition = { "adapter": adapter, @@ -564,6 +576,7 @@ def data_section( if selection is BACK: return False catalog, selected_subsets, weights = selection + assert selected_subsets is not None by_name = {item.name: item for item in catalog.subsets} subset_selection = { "source": catalog.source, @@ -645,7 +658,7 @@ def data_section( ], default=default_layout, ) - if layout is BACK: + if _is_back(layout): return False sequence = _integer_field( session, @@ -654,7 +667,7 @@ def data_section( "Sequence length used by width, depth, bypass, evaluation, and global KD:", 4096, ) - if sequence is BACK: + if _is_back(sequence): return False session.state.set_field("data.source", runtime_source, source=source_kind) session.state.set_field("data.selected_source", source, source=source_kind) @@ -1143,7 +1156,7 @@ def _profile_prompt( "Parallel configuration name:", default=stage_id, ) - if name is BACK: + if _is_back(name): return BACK values = {} for field_name, label, default in ( @@ -1312,7 +1325,7 @@ def _configure_stage_resource( default=int(defaults["instances"]), minimum=1, ) - if instances is BACK: + if _is_back(instances): return BACK batch_unit = profile.batch_unit requested_batch = session.integer( @@ -1440,7 +1453,7 @@ def depth_section(session: WizardSession, resolver: DefaultsResolver, context: d _depth_granularity_choices(inventory), default=granularity, ) - if granularity is BACK: + if _is_back(granularity): return False count = inventory.num_sublayers if granularity == "subblock" else inventory.num_layers remove = session.integer( @@ -1463,7 +1476,7 @@ def depth_section(session: WizardSession, resolver: DefaultsResolver, context: d default=samples, minimum=1, ) - if samples is BACK: + if _is_back(samples): return False pruning["depth_granularity"] = str(granularity) pruning["depth_remove"] = int(remove) @@ -1559,20 +1572,22 @@ def width_axes_section(session: WizardSession, resolver: DefaultsResolver, conte selected = defaults[axis.axis_id] if action == "customize": require_reduced = index == len(inventory.axes) - 1 and not has_reduced_axis + + def validate_selection(values: Any) -> bool | str: + return _axis_selection_validation( + axis, + values, + require_reduced=require_reduced, + ) + selected = session.checkbox( f"pruning.axes.{axis.axis_id}", f"Values for {axis.label}:", [(str(value), value) for value in axis.values], defaults=selected, - validate=lambda values, axis=axis, require_reduced=require_reduced: ( - _axis_selection_validation( - axis, - values, - require_reduced=require_reduced, - ) - ), + validate=validate_selection, ) - if selected is BACK: + if _is_back(selected): return False selected_values = _normalized_axis_values(axis, selected) enabled = any(value < int(axis.teacher_value) for value in selected_values) @@ -1633,7 +1648,7 @@ def width_importance_section( default=samples, minimum=1, ) - if samples is BACK: + if _is_back(samples): return False pruning["width_importance_samples"] = int(samples) session.state.set_collection("pruning", pruning) @@ -1684,7 +1699,7 @@ def sort_sanity_section(session: WizardSession, resolver: DefaultsResolver, cont "Run sorting sanity evaluation?", default=enabled, ) - if enabled is BACK: + if _is_back(enabled): return False if enabled: samples = session.integer( @@ -1693,7 +1708,7 @@ def sort_sanity_section(session: WizardSession, resolver: DefaultsResolver, cont default=samples, minimum=1, ) - if samples is BACK: + if _is_back(samples): return False pruning["sort_sanity"] = bool(enabled) pruning["sort_sanity_samples"] = int(samples) @@ -1781,7 +1796,7 @@ def width_sanity_section(session: WizardSession, resolver: DefaultsResolver, con "Run width sanity evaluation?", default=enabled, ) - if enabled is BACK: + if _is_back(enabled): return False if enabled: samples = session.integer( @@ -1806,7 +1821,7 @@ def width_sanity_section(session: WizardSession, resolver: DefaultsResolver, con minimum=1, maximum=max_targets_per_axis, ) - if BACK in (samples, layer_count, targets_per_axis): + if any(_is_back(value) for value in (samples, layer_count, targets_per_axis)): return False pruning["width_sanity"] = bool(enabled) pruning["width_sanity_samples"] = int(samples) @@ -1861,7 +1876,7 @@ def slicing_sanity_section( "Run slicing sanity evaluation?", default=enabled, ) - if enabled is BACK: + if _is_back(enabled): return False pruning["slicing_sanity"] = bool(enabled) session.state.set_collection("pruning", pruning) @@ -1916,7 +1931,7 @@ def bypass_section(session: WizardSession, resolver: DefaultsResolver, context: "Run local bypass distillation?", default=enabled, ) - if enabled is BACK: + if _is_back(enabled): return False if enabled: granularity = session.select( @@ -2128,7 +2143,7 @@ def pre_mip_stages_section( default=gpus_per_node, minimum=1, ) - if instances is BACK: + if _is_back(instances): return False requested_batch = session.integer( f"stages.{stage_id}.batch", @@ -2136,7 +2151,7 @@ def pre_mip_stages_section( default=resolve_batch(8, profile).effective, minimum=profile.batch_unit, ) - if requested_batch is BACK: + if _is_back(requested_batch): return False else: profile = ( @@ -3350,14 +3365,15 @@ def mip_section(session: WizardSession, resolver: DefaultsResolver, context: dic if homogeneous_mode == "all": homogeneous_keep = "all" else: - homogeneous_keep = session.integer( + selected_homogeneous_keep = session.integer( f"mip.search.{search_id}.homogeneous.keep", "Homogeneous solutions retained per solve:", default=5, minimum=1, ) - if homogeneous_keep is BACK: + if selected_homogeneous_keep is BACK: return False + homogeneous_keep = int(selected_homogeneous_keep) explicit_variants: OrderedDict[str, Any] = OrderedDict() if variants: @@ -3525,7 +3541,7 @@ def post_mip_section(session: WizardSession, resolver: DefaultsResolver, context ], default="recommended", ) - if flow_mode is BACK: + if _is_back(flow_mode): return False if flow_mode == "none": continue @@ -3836,7 +3852,7 @@ def _vllm_topology_prompt( ), default=bool(topology_defaults.get("enable_expert_parallel", False)), ) - if enable_expert_parallel is BACK: + if _is_back(enable_expert_parallel): return BACK topology.update( { @@ -3859,6 +3875,9 @@ def _vllm_topology_prompt( stage_id=stage_id, ) if issues: + if isinstance(session.backend, NonInteractiveBackend): + detail = "; ".join(f"{issue.path}: {issue.message}" for issue in issues) + raise SetupError(f"Non-interactive vLLM topology is incompatible: {detail}") _print_parallel_issues(issues) topology_defaults = topology continue @@ -4052,7 +4071,7 @@ def _configure_dynamic_resources( f"Customize resources and batch for {node_id}?", default=False, ) - if customize is BACK: + if _is_back(customize): return BACK strategy = _post_mip_strategy(node) instances = gpus_per_node @@ -4063,7 +4082,7 @@ def _configure_dynamic_resources( default=gpus_per_node, minimum=1, ) - if instances is BACK: + if _is_back(instances): return BACK entry = { "strategy": str(strategy), @@ -4142,7 +4161,7 @@ def _configure_dynamic_resources( default=profile.batch_unit, minimum=profile.batch_unit, ) - if requested is BACK: + if _is_back(requested): return BACK resolved = resolve_batch(int(requested), profile) config = dict(node.config) @@ -4482,10 +4501,26 @@ def _fresh_state( defaults_path: Path | None, *, full: bool, + campaign_dir: Path | None = None, + setup_profile: str = "balanced", ) -> WizardState: + if campaign_dir is not None: + if full: + return WizardState.start( + Path(campaign_dir).expanduser(), + defaults_path=defaults_path, + setup_mode="full", + ) + get_setup_preset(setup_profile) + return WizardState.start( + Path(campaign_dir).expanduser(), + defaults_path=defaults_path, + setup_mode="quick", + preset=setup_profile, + ) if full: while True: - value = backend.text("Campaign directory:", "") + value = backend.text("Campaign directory:", None) if value is BACK: continue if not str(value).strip(): @@ -4498,11 +4533,11 @@ def _fresh_state( ) while True: - preset = _select_setup_preset(backend) + preset = _select_setup_preset(backend, default=setup_profile) if preset is BACK: continue while True: - value = backend.text("Campaign directory:", "") + value = backend.text("Campaign directory:", None) if value is BACK: break if not str(value).strip(): @@ -4534,8 +4569,16 @@ def run_wizard_v2( defaults_path: Path | None, backend: PromptBackend | None = None, full: bool = False, + campaign_dir: Path | None = None, + setup_profile: str = "balanced", ) -> Path: - """Run setup v2, save every answer, validate bundles, and never launch jobs.""" + """Resolve setup answers, validate both bundles, and never launch jobs. + + ``campaign_dir`` and ``setup_profile`` let automation bypass only the new + campaign prompts while using the same wizard sections and bundle renderer. + """ + if resume is not None and campaign_dir is not None: + raise SetupError("campaign_dir cannot be combined with resume") backend = backend or InteractiveBackend() print("Welcome to Puzzletron setup v2.") if resume is None: @@ -4547,7 +4590,13 @@ def run_wizard_v2( "defaults from a profile." ) print(" Use --full only when you need every advanced control.") - state = _fresh_state(backend, defaults_path, full=full) + state = _fresh_state( + backend, + defaults_path, + full=full, + campaign_dir=campaign_dir, + setup_profile=setup_profile, + ) else: state = WizardState.resume(resume) if full and state.setup_mode != "full": diff --git a/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py b/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py index 616b4d24d9b..9e8986dccc6 100644 --- a/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py +++ b/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Tests for Puzzletron AIPerf context-capacity handling.""" + import json from pathlib import Path from types import SimpleNamespace @@ -20,6 +22,7 @@ import pytest from modelopt.torch.puzzletron.benchmarks.aiperf import ( + _aiperf_subprocess_environment, _canonical_topology, _clean_subprocess_environment, _exact_length_extra_inputs, @@ -28,6 +31,7 @@ _profile_command, _server_max_model_len, _topology_vllm_args, + _vllm_server_command, ) @@ -89,11 +93,27 @@ def test_prepare_vllm_checkpoint_refreshes_heterogeneous_metadata(monkeypatch, t observed = [] monkeypatch.setattr( "modelopt.torch.puzzletron.utils.vllm_adapter.refresh_realized_checkpoint_config", - lambda path: observed.append(path), + lambda path, **kwargs: observed.append((path, kwargs)), ) assert _prepare_vllm_checkpoint(tmp_path) is True - assert observed == [tmp_path] + assert observed == [(tmp_path, {"trust_remote_code": False})] + + +def test_prepare_vllm_checkpoint_preserves_explicit_remote_code_trust(monkeypatch, tmp_path): + config = { + "architectures": ["BaseModel"], + "text_config": {"per_layer_config": {"0": {"intermediate_size": 8}}}, + } + (tmp_path / "config.json").write_text(json.dumps(config)) + observed = [] + monkeypatch.setattr( + "modelopt.torch.puzzletron.utils.vllm_adapter.refresh_realized_checkpoint_config", + lambda path, **kwargs: observed.append((path, kwargs)), + ) + + assert _prepare_vllm_checkpoint(tmp_path, trust_remote_code=True) is True + assert observed == [(tmp_path, {"trust_remote_code": True})] def test_prepare_vllm_checkpoint_leaves_native_teacher_unchanged(tmp_path): @@ -103,6 +123,36 @@ def test_prepare_vllm_checkpoint_leaves_native_teacher_unchanged(tmp_path): assert _prepare_vllm_checkpoint(tmp_path) is False +def _offline_environment() -> dict[str, str]: + return { + "HF_HUB_OFFLINE": "1", + "TRANSFORMERS_OFFLINE": "1", + "HF_DATASETS_OFFLINE": "1", + "HF_HOME": "/cache/huggingface", + "UNCHANGED": "value", + } + + +def test_aiperf_online_tokenizer_relaxation_is_explicit_and_non_mutating(): + expected_source = _offline_environment() + source = dict(expected_source) + + default_environment = _aiperf_subprocess_environment(source) + resolved = _aiperf_subprocess_environment( + source, + allow_aiperf_v011_online_tokenizer_resolution=True, + ) + + assert default_environment == source + assert default_environment is not source + assert "HF_HUB_OFFLINE" not in resolved + assert "TRANSFORMERS_OFFLINE" not in resolved + assert resolved["HF_DATASETS_OFFLINE"] == "1" + assert resolved["HF_HOME"] == "/cache/huggingface" + assert resolved["UNCHANGED"] == "value" + assert source == expected_source + + def test_canonical_topology_covers_tp_pp_dp_effective_ep_and_context_parallel(): topology = _canonical_topology( { @@ -155,6 +205,67 @@ def test_vllm_topology_args_enable_dp_and_expert_parallel_only_when_requested(): assert "--expert-parallel-size" not in ep_args +def test_vllm_server_command_applies_explicit_remote_code_policy(monkeypatch, tmp_path): + monkeypatch.setattr( + "modelopt.torch.puzzletron.benchmarks.aiperf._descriptor_vllm_args", + lambda _checkpoint: [], + ) + + default_command = _vllm_server_command( + checkpoint_dir=tmp_path, + port=8000, + model_name="served-model", + input_tokens=32, + output_tokens=8, + topology={"gpu_group_size": 1}, + trust_remote_code=False, + ) + trusted_command = _vllm_server_command( + checkpoint_dir=tmp_path, + port=8000, + model_name="served-model", + input_tokens=32, + output_tokens=8, + topology={"gpu_group_size": 1}, + trust_remote_code=True, + ) + + assert "--trust-remote-code" not in default_command + assert "--trust-remote-code" in trusted_command + + +@pytest.mark.parametrize( + "extra_arg", + [ + "--trust-remote-code", + "--trust-remote-code=true", + "--trust_remote_code", + "--trust_remote_code=true", + "--trust-rem", + "--trust_rem", + "--config", + "--config=policy.yaml", + "--conf", + ], +) +def test_vllm_server_command_rejects_remote_code_policy_override(monkeypatch, tmp_path, extra_arg): + monkeypatch.setattr( + "modelopt.torch.puzzletron.benchmarks.aiperf._descriptor_vllm_args", + lambda _checkpoint: [], + ) + + with pytest.raises(ValueError, match="cannot set policy-owned vLLM options"): + _vllm_server_command( + checkpoint_dir=tmp_path, + port=8000, + model_name="served-model", + input_tokens=32, + output_tokens=8, + topology={"gpu_group_size": 1, "extra_vllm_args": [extra_arg]}, + trust_remote_code=False, + ) + + def test_profile_command_maps_each_workload_answer_to_aiperf_cli(tmp_path): command = _profile_command( executable=Path("/opt/aiperf/bin/aiperf"), @@ -179,6 +290,7 @@ def test_profile_command_maps_each_workload_answer_to_aiperf_cli(tmp_path): assert command[command.index("--synthetic-input-tokens-stddev") + 1] == "0" assert command[command.index("--output-tokens-mean") + 1] == "128" assert command[command.index("--output-tokens-stddev") + 1] == "0" + assert command[command.index("--tokenizer") + 1] == str(tmp_path / "tokenizer") assert "--use-server-token-count" in command diff --git a/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py b/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py index d1918c2c7e4..6d3ae235a16 100644 --- a/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py +++ b/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py @@ -289,6 +289,8 @@ def test_rpc_executor_scores_cumulative_depth_removals(monkeypatch): executor.params = {"micro_batch_size": 4} executor.sliced_teacher_baseline = {"lm_loss": {"avg": 0.0}} executor.latest_observability = None + executor.latest_score_device_type = "cpu" + executor.visible_cuda_device_count = 0 monkeypatch.setattr( executor, "_score", @@ -309,6 +311,48 @@ def test_rpc_executor_scores_cumulative_depth_removals(monkeypatch): assert result.metrics["lm_loss"]["avg"] == 1.0 assert [target["layer_idx"] for target in captured[0]] == [0, 1] assert result.provenance["micro_batch_size"] == 4 + assert result.provenance["score_device_type"] == "cpu" + assert result.provenance["visible_cuda_device_count"] == 0 + + +def test_rpc_executor_non_output_pipeline_rank_reaches_collective(monkeypatch): + from modelopt.torch.puzzletron.distributed_eval.automodel_executor import ( + AutoModelReplaceBlockExecutor, + ) + from modelopt.torch.utils import distributed as dist + + class NonOutputRecipe: + has_outputs = False + _groups = None + + @staticmethod + def tensor_parallel_group(): + return None + + @staticmethod + def iterate_captures(): + yield None, None + + @staticmethod + def context_parallel_group(): + return None + + @staticmethod + def observability_metadata(): + return {"pipeline_role": "non_output"} + + barriers = [] + executor = AutoModelReplaceBlockExecutor.__new__(AutoModelReplaceBlockExecutor) + executor.recipe = NonOutputRecipe() + executor.cache = object() + executor.params = {} + executor.bypass_checkpoint_dir = None + executor.latest_score_device_type = None + monkeypatch.setattr(dist, "barrier", lambda: barriers.append(True)) + + assert executor._score(None) is None + assert barriers == [True] + assert executor.latest_observability == {"pipeline_role": "non_output"} def test_runtime_fingerprint_ignores_distributed_compute_dtype_transition(): diff --git a/tests/unit/torch/puzzletron/test_ci_environment.py b/tests/unit/torch/puzzletron/test_ci_environment.py new file mode 100644 index 00000000000..8ede67110c4 --- /dev/null +++ b/tests/unit/torch/puzzletron/test_ci_environment.py @@ -0,0 +1,185 @@ +# 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 Puzzletron CI environment provenance checks.""" + +import json +import sys +from importlib import metadata + +import pytest + +import noxfile +from examples.puzzletron import ci_environment + + +class _Distribution: + def __init__(self, payload: dict): + self.payload = payload + + def read_text(self, filename: str) -> str | None: + assert filename == "direct_url.json" + return json.dumps(self.payload) + + +_EXPECTED_SOURCE = { + "repository": "https://github.com/Separius/Automodel.git", + "commit": "b22cd029d806197e249f2cc4a42c5de91713b772", +} + + +def _pep610_source(repository: str, commit: str) -> dict: + return { + "url": repository, + "vcs_info": {"vcs": "git", "commit_id": commit}, + } + + +def test_pep610_exact_source_is_accepted(monkeypatch): + monkeypatch.setattr( + ci_environment.metadata, + "distribution", + lambda _package: _Distribution( + _pep610_source(_EXPECTED_SOURCE["repository"], _EXPECTED_SOURCE["commit"]) + ), + ) + + ci_environment.verify_installed_vcs_source("nemo-automodel", _EXPECTED_SOURCE) + + +@pytest.mark.parametrize( + ("repository", "commit"), + [ + ("https://github.com/example/Automodel.git", _EXPECTED_SOURCE["commit"]), + (_EXPECTED_SOURCE["repository"], "0" * 40), + ], + ids=("repository", "commit"), +) +def test_pep610_vcs_source_mismatch_is_rejected(monkeypatch, repository, commit): + monkeypatch.setattr( + ci_environment.metadata, + "distribution", + lambda _package: _Distribution(_pep610_source(repository, commit)), + ) + + with pytest.raises(RuntimeError, match="source mismatch"): + ci_environment.verify_installed_vcs_source("nemo-automodel", _EXPECTED_SOURCE) + + +def test_editable_pinned_dependency_must_be_clean(monkeypatch): + monkeypatch.setattr( + ci_environment.metadata, + "distribution", + lambda _package: _Distribution( + {"url": "file:///src/automodel", "dir_info": {"editable": True}} + ), + ) + outputs = iter( + [ + "https://github.com/Separius/Automodel.git\n", + "b22cd029d806197e249f2cc4a42c5de91713b772\n", + " M nemo_automodel/model.py\n", + ] + ) + monkeypatch.setattr( + ci_environment.subprocess, + "check_output", + lambda *_args, **_kwargs: next(outputs), + ) + + with pytest.raises(RuntimeError, match="dependency 'nemo-automodel' is dirty"): + ci_environment.verify_installed_vcs_source( + "nemo-automodel", + _EXPECTED_SOURCE, + ) + + +def test_nox_verifier_executes_scalar_version_and_exact_vcs_checks(monkeypatch): + lmms_source = { + "base_version": "7.8.9", + "repository": "https://example.test/lmms-eval.git", + "commit": "1" * 40, + } + automodel_source = { + "base_version": "4.5.6", + "repository": "https://example.test/Automodel.git", + "commit": "2" * 40, + } + expected_versions = { + "python": f"{sys.version_info.major}.{sys.version_info.minor}", + "torch": "1.2.3", + "torchvision": "2.3.4", + "transformers": "3.4.5", + "lmms-eval": lmms_source["base_version"], + "nemo-automodel": automodel_source["base_version"], + } + monkeypatch.setattr( + noxfile, + "PUZZLETRON_V2_CI_ENVIRONMENT", + { + **expected_versions, + "lmms_eval": lmms_source, + "nemo_automodel": automodel_source, + }, + ) + monkeypatch.setattr(noxfile, "PUZZLETRON_V2_LMMS_SOURCE", lmms_source) + monkeypatch.setattr(noxfile, "PUZZLETRON_V2_AUTOMODEL_SOURCE", automodel_source) + monkeypatch.setattr(metadata, "version", lambda package: expected_versions[package]) + vcs_calls = [] + monkeypatch.setattr( + ci_environment, + "verify_installed_vcs_source", + lambda package, source: vcs_calls.append((package, source)), + ) + + class ExecutingSession: + def run(self, python, flag, script): + assert (python, flag) == ("python", "-c") + exec(compile(script, "", "exec"), {}) + + noxfile._verify_puzzletron_v2_environment(ExecutingSession()) + + assert vcs_calls == [ + ("lmms-eval", lmms_source), + ("nemo-automodel", automodel_source), + ] + + +def test_puzzletron_nox_session_verifies_environment_before_pytest(monkeypatch): + events = [] + + class RecordingSession: + def install(self, *args): + events.append(("install", args)) + + def run(self, *args): + events.append(("run", args)) + + monkeypatch.setattr( + noxfile, + "_verify_puzzletron_v2_environment", + lambda session: events.append(("verify", session)), + ) + session = RecordingSession() + + noxfile.puzzletron_v2.func(session) + + verify_index = events.index(("verify", session)) + pytest_index = next( + index + for index, event in enumerate(events) + if event[0] == "run" and event[1][:3] == ("python", "-m", "pytest") + ) + assert verify_index < pytest_index diff --git a/tests/unit/torch/puzzletron/test_future_stages.py b/tests/unit/torch/puzzletron/test_future_stages.py index 28a89e2fd92..ce11cf36e8d 100644 --- a/tests/unit/torch/puzzletron/test_future_stages.py +++ b/tests/unit/torch/puzzletron/test_future_stages.py @@ -1,5 +1,19 @@ # 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. + +"""Future-stage configuration and artifact-selection contracts.""" import json from pathlib import Path @@ -7,6 +21,35 @@ import pytest import torch +from modelopt.torch.puzzletron.security_policy import require_boolean_policy +from modelopt.torch.puzzletron.stages.future import evaluation_stage + + +def test_security_policy_rejects_non_boolean_values(): + with pytest.raises(ValueError, match="^policy must be a boolean$"): + require_boolean_policy("false", path="policy") + + +def test_security_policy_resolves_none_only_with_an_explicit_default(): + assert require_boolean_policy(None, path="policy", default=False) is False + assert require_boolean_policy(None, path="policy", default=True) is True + with pytest.raises(ValueError, match="^policy must be a boolean$"): + require_boolean_policy(None, path="policy") + + +def test_evaluation_stage_rejects_scalar_checkpoints(): + config = { + "zero_shot_evaluation": { + "enabled": True, + "checkpoints": "/checkpoint", + } + } + with pytest.raises( + ValueError, + match=r"^zero_shot_evaluation\.checkpoints must be a list or tuple$", + ): + evaluation_stage(config, object()) + def test_distillation_sanity_accepts_packed_cache_without_raw_dataset(tmp_path): from modelopt.torch.puzzletron.stages.future import _distillation_dataset_source @@ -72,9 +115,7 @@ def test_evaluation_descriptor_honors_explicit_legacy_override(monkeypatch, tmp_ ) -def test_scenario_grid_kd_builds_one_isolated_config_per_realized_checkpoint( - monkeypatch, tmp_path -): +def test_scenario_grid_kd_builds_one_isolated_config_per_realized_checkpoint(monkeypatch, tmp_path): from modelopt.torch.puzzletron.stages import future puzzle_dir = tmp_path / "model" @@ -226,12 +267,8 @@ def fail_first(value): observed.append(value) raise RuntimeError("stop") - try: + with pytest.raises(RuntimeError, match="^stop$"): future._bounded_map(fail_first, range(5), max_workers=1) - except RuntimeError as error: - assert str(error) == "stop" - else: - raise AssertionError("worker failure must propagate") assert observed == [0] diff --git a/tests/unit/torch/puzzletron/test_global_kd_canonical.py b/tests/unit/torch/puzzletron/test_global_kd_canonical.py index 7c9e0c12f76..2177f960c92 100644 --- a/tests/unit/torch/puzzletron/test_global_kd_canonical.py +++ b/tests/unit/torch/puzzletron/test_global_kd_canonical.py @@ -1,11 +1,26 @@ # 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 Puzzletron's canonical global-distillation behavior.""" import json from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace +import pytest import torch from modelopt.torch.puzzletron.distillation.global_automodel import ( @@ -87,9 +102,7 @@ def native_state_dict_adapter_context(cls, block_configs): assert observed[0][0].to_dict() == {"subblock_configs": []} -def test_distillation_overfit_stage_disables_mtp_objectives_by_default( - monkeypatch, tmp_path -): +def test_distillation_overfit_stage_disables_mtp_objectives_by_default(monkeypatch, tmp_path): """Nano-like checkpoints without MTP must not enable MTP loss implicitly.""" from modelopt.torch.puzzletron.manifest import StageManifest from modelopt.torch.puzzletron.stages import future @@ -146,9 +159,7 @@ def fake_run_global_kd(kd_config): }, } - future.distillation_overfit_stage( - config, StageManifest(stage="global_distillation_sanity") - ) + future.distillation_overfit_stage(config, StageManifest(stage="global_distillation_sanity")) assert captured["objective"] == { "main_ce": {"weight": 1.0}, @@ -171,19 +182,10 @@ def test_global_distillation_summary_publishes_canonical_training_records(tmp_pa training_log = output_dir / "checkpoints/training.jsonl" training_log.parent.mkdir(parents=True) training_log.write_text( - json.dumps({"step": 1, "loss": 2.0}) - + "\n" - + json.dumps({"step": 2, "loss": 1.0}) - + "\n" + json.dumps({"step": 1, "loss": 2.0}) + "\n" + json.dumps({"step": 2, "loss": 1.0}) + "\n" ) for step in (1, 2): - checkpoint = ( - output_dir - / "checkpoints" - / f"epoch_0_step_{step}" - / "model" - / "consolidated" - ) + checkpoint = output_dir / "checkpoints" / f"epoch_0_step_{step}" / "model" / "consolidated" checkpoint.mkdir(parents=True) (checkpoint / "config.json").write_text("{}") (checkpoint.parents[1] / "saving_completed").touch() @@ -209,9 +211,7 @@ def test_global_distillation_summary_publishes_canonical_training_records(tmp_pa assert payload["max_steps"] == 256 assert payload["sequence_length"] == 16384 assert payload["records"][-1] == {"step": 2, "loss": 1.0} - assert payload["post_kd_checkpoint"].endswith( - "checkpoints/epoch_0_step_2/model/consolidated" - ) + assert payload["post_kd_checkpoint"].endswith("checkpoints/epoch_0_step_2/model/consolidated") def test_global_kd_metric_logger_flushes_every_optimizer_step(): @@ -227,9 +227,7 @@ def test_global_kd_metric_logger_flushes_every_optimizer_step(): assert logger.flush is True -def test_global_kd_uses_memory_bounded_1f1b_by_default_and_allows_override( - tmp_path, monkeypatch -): +def test_global_kd_uses_memory_bounded_1f1b_by_default_and_allows_override(tmp_path, monkeypatch): monkeypatch.setattr( "modelopt.torch.puzzletron.plugins.automodel.config.inject_descriptor_pipeline_config", lambda *args, **kwargs: None, @@ -262,10 +260,7 @@ def test_global_kd_uses_memory_bounded_1f1b_by_default_and_allows_override( assert default_recipe["distributed"]["pipeline"]["pp_microbatch_size"] == 1 assert default_recipe["distributed"]["pipeline"]["pp_batch_size"] == 8 assert default_recipe["dataloader"]["batch_size"] == 8 - assert ( - override_recipe["distributed"]["pipeline"]["pp_schedule"] - == "interleaved1f1b" - ) + assert override_recipe["distributed"]["pipeline"]["pp_schedule"] == "interleaved1f1b" def test_global_kd_auto_domain_uses_canonical_text_dataset(tmp_path, monkeypatch): @@ -310,10 +305,7 @@ def test_global_kd_auto_domain_uses_canonical_text_dataset(tmp_path, monkeypatch assert recipe["recipe"] == "KnowledgeDistillationRecipeForNextTokenPrediction" assert recipe["dataset"] == { - "_target_": ( - "modelopt.torch.puzzletron.distillation.dataset." - "make_puzzletron_llm_dataset" - ), + "_target_": ("modelopt.torch.puzzletron.distillation.dataset.make_puzzletron_llm_dataset"), "dataset_path": str(tmp_path / "dataset"), "split": "train", "num_samples": 4096, @@ -326,9 +318,7 @@ def test_global_kd_auto_domain_uses_canonical_text_dataset(tmp_path, monkeypatch ) -def test_global_kd_packed_text_uses_native_chat_data_and_canonical_pack_size( - tmp_path, monkeypatch -): +def test_global_kd_packed_text_uses_native_chat_data_and_canonical_pack_size(tmp_path, monkeypatch): monkeypatch.setattr( "modelopt.torch.puzzletron.plugins.automodel.config.inject_descriptor_model_kwargs", lambda *args, **kwargs: None, @@ -361,9 +351,7 @@ def test_global_kd_packed_text_uses_native_chat_data_and_canonical_pack_size( recipe = build_automodel_global_kd_recipe(config) - assert recipe["dataset"]["_target_"].endswith( - "make_puzzletron_chat_dataset" - ) + assert recipe["dataset"]["_target_"].endswith("make_puzzletron_chat_dataset") assert recipe["dataloader"]["collate_fn"].endswith("default_collater") assert recipe["packed_sequence"] == { "packed_sequence_size": 256, @@ -395,12 +383,16 @@ def test_global_kd_recipe_publishes_explicit_resume_policy(tmp_path, monkeypatch "pp": 1, } - assert build_automodel_global_kd_recipe(GlobalKDConfig(**common, resume=True))[ - "puzzletron_resume" - ] is True - assert build_automodel_global_kd_recipe(GlobalKDConfig(**common, resume=False))[ - "puzzletron_resume" - ] is False + assert ( + build_automodel_global_kd_recipe(GlobalKDConfig(**common, resume=True))["puzzletron_resume"] + is True + ) + assert ( + build_automodel_global_kd_recipe(GlobalKDConfig(**common, resume=False))[ + "puzzletron_resume" + ] + is False + ) def test_global_kd_config_preserves_per_model_dtype_overrides(tmp_path): @@ -430,12 +422,12 @@ def test_global_kd_config_preserves_per_model_dtype_overrides(tmp_path): kd = build_global_kd_config(config) - assert global_automodel._model_recipe(kd, teacher=False, domain="llm")[ - "torch_dtype" - ] == "float32" - assert global_automodel._model_recipe(kd, teacher=True, domain="llm")[ - "torch_dtype" - ] == "bfloat16" + assert ( + global_automodel._model_recipe(kd, teacher=False, domain="llm")["torch_dtype"] == "float32" + ) + assert ( + global_automodel._model_recipe(kd, teacher=True, domain="llm")["torch_dtype"] == "bfloat16" + ) def test_global_kd_load_checkpoint_honors_resume_policy(): @@ -491,9 +483,7 @@ def test_global_distillation_stage_promotes_canonical_namespace(tmp_path): assert (kd_config.pp, kd_config.cp, kd_config.dp) == (2, 4, 8) -def test_global_kd_preserves_physical_dp_mesh_when_ep_overlays_shards( - tmp_path, monkeypatch -): +def test_global_kd_preserves_physical_dp_mesh_when_ep_overlays_shards(tmp_path, monkeypatch): monkeypatch.setattr( "modelopt.torch.puzzletron.plugins.automodel.config.inject_descriptor_pipeline_config", lambda *args, **kwargs: None, @@ -563,18 +553,14 @@ def from_local(cls, local, *, device_mesh, placements, run_check): source_mesh = Mesh() head_mesh = Mesh() hidden = FakeDTensor(torch.ones(2, 3), source_mesh) - base_layer = type( - "BaseLayer", (), {"weight": FakeDTensor(torch.ones(4, 3), head_mesh)} - )() + base_layer = type("BaseLayer", (), {"weight": FakeDTensor(torch.ones(4, 3), head_mesh)})() head = type("WrappedHead", (), {"base_layer": base_layer})() aligned = global_kd_recipe._align_dtensor_to_module_mesh(hidden, head) assert aligned.device_mesh is head_mesh assert aligned.placements == hidden.placements - assert FakeDTensor.calls == [ - (hidden.local, head_mesh, hidden.placements, False) - ] + assert FakeDTensor.calls == [(hidden.local, head_mesh, hidden.placements, False)] def test_teacher_mtp_projection_uses_local_head_and_student_logit_mesh(monkeypatch): @@ -603,15 +589,13 @@ def from_local(cls, local, *, device_mesh, placements, run_check): teacher_mesh = Mesh() student_mesh = Mesh() hidden = FakeDTensor(torch.tensor([[1.0, 2.0]]), teacher_mesh, ("replicate",)) - weight = FakeDTensor( - torch.tensor([[1.0, 0.0], [0.0, 2.0]]), teacher_mesh, ("shard0",) - ) - head = type("WrappedHead", (), {"base_layer": type("Base", (), {"weight": weight, "bias": None})()})() + weight = FakeDTensor(torch.tensor([[1.0, 0.0], [0.0, 2.0]]), teacher_mesh, ("shard0",)) + head = type( + "WrappedHead", (), {"base_layer": type("Base", (), {"weight": weight, "bias": None})()} + )() reference = FakeDTensor(torch.empty(1, 2), student_mesh, ("shard_vocab",)) - projected = global_kd_recipe._project_teacher_hidden_on_reference_mesh( - hidden, head, reference - ) + projected = global_kd_recipe._project_teacher_hidden_on_reference_mesh(hidden, head, reference) assert projected.device_mesh is student_mesh assert projected.placements == reference.placements @@ -740,10 +724,12 @@ def __init__(self): assert all(value.item() > 0 for value in recipe._gradient_squared.values()) -def test_global_kd_checkpoint_forwards_best_metric_key(tmp_path): +def test_global_kd_checkpoint_forwards_best_metric_key(tmp_path, monkeypatch): + # Lazy import keeps the optional NeMo AutoModel runtime out of test collection. from modelopt.torch.puzzletron.distillation.global_kd_recipe import _WeightedObjectiveMixin calls = [] + refreshes = [] class BaseRecipe: def save_checkpoint( @@ -755,7 +741,12 @@ def save_checkpoint( best_metric_key="default", ): calls.append((epoch, step, train_loss, val_loss, best_metric_key)) - (tmp_path / f"epoch_{epoch}_step_{step}").mkdir() + checkpoint = tmp_path / f"epoch_{epoch}_step_{step}" + consolidated = checkpoint / "model/consolidated" + consolidated.mkdir(parents=True) + (consolidated / "config.json").write_text( + json.dumps({"block_configs": [{"subblock_configs": []}]}) + ) return "saved" class Recipe(_WeightedObjectiveMixin, BaseRecipe): @@ -768,6 +759,13 @@ class Recipe(_WeightedObjectiveMixin, BaseRecipe): {"config": type("Config", (), {"checkpoint_dir": tmp_path})()}, )() recipe.dist_env = type("DistEnv", (), {"is_main": True})() + recipe.cfg = {"model": {"trust_remote_code": True}} + monkeypatch.setattr( + "modelopt.torch.puzzletron.utils.vllm_adapter.refresh_realized_checkpoint_config", + lambda path, **kwargs: refreshes.append( + (path, kwargs, (tmp_path / "epoch_2_step_17/saving_completed").exists()) + ), + ) result = recipe.save_checkpoint( 2, @@ -779,9 +777,81 @@ class Recipe(_WeightedObjectiveMixin, BaseRecipe): assert result == "saved" assert calls == [(2, 17, 0.5, {"lm_loss": 0.25}, "lm_loss")] + assert refreshes == [ + ( + tmp_path / "epoch_2_step_17/model/consolidated", + {"trust_remote_code": True}, + False, + ) + ] assert (tmp_path / "epoch_2_step_17" / "saving_completed").is_file() +def test_global_kd_checkpoint_publication_failure_reaches_all_ranks(tmp_path, monkeypatch): + # Import the dynamic mixin at test runtime to preserve lightweight module collection. + from modelopt.torch.puzzletron.distillation.global_kd_recipe import _WeightedObjectiveMixin + + publication = {"error": None, "broadcasts": 0} + + class BaseRecipe: + def save_checkpoint( + self, + epoch, + step, + train_loss, + val_loss, + best_metric_key="default", + ): + del train_loss, val_loss, best_metric_key + consolidated = tmp_path / f"epoch_{epoch}_step_{step}/model/consolidated" + consolidated.mkdir(parents=True, exist_ok=True) + (consolidated / "config.json").write_text( + json.dumps({"block_configs": [{"subblock_configs": []}]}) + ) + return "saved" + + class Recipe(_WeightedObjectiveMixin, BaseRecipe): + pass + + def broadcast_object_list(payload, *, src): + assert src == 0 + publication["broadcasts"] += 1 + if payload[0] is not None: + publication["error"] = payload[0] + else: + payload[0] = publication["error"] + + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list) + monkeypatch.setattr( + "modelopt.torch.puzzletron.utils.vllm_adapter.refresh_realized_checkpoint_config", + lambda *args, **kwargs: (_ for _ in ()).throw(ValueError("refresh failed")), + ) + + def recipe(*, is_main): + instance = Recipe() + instance.checkpointer = type( + "Checkpointer", + (), + {"config": type("Config", (), {"checkpoint_dir": tmp_path})()}, + )() + instance.dist_env = type("DistEnv", (), {"is_main": is_main})() + instance.cfg = {"model": {"trust_remote_code": False}} + return instance + + with pytest.raises(ValueError, match="refresh failed"): + recipe(is_main=True).save_checkpoint(2, 19, 0.5, {"lm_loss": 0.25}) + assert publication == {"error": "ValueError: refresh failed", "broadcasts": 1} + assert not (tmp_path / "epoch_2_step_19/saving_completed").exists() + + with pytest.raises( + RuntimeError, + match="global KD checkpoint publication failed on rank 0: ValueError: refresh failed", + ): + recipe(is_main=False).save_checkpoint(2, 19, 0.5, {"lm_loss": 0.25}) + assert publication["broadcasts"] == 2 + + def test_global_kd_optimizer_save_uses_the_actual_pipeline_model_parts(): import torch @@ -800,9 +870,7 @@ def save_model(self, model, path): def save_optimizer(self, saved_optimizer, model, path, scheduler): del path, scheduler saved_models.append(model) - assert saved_optimizer.param_groups[0]["params"] == list( - original.parameters() - ) + assert saved_optimizer.param_groups[0]["params"] == list(original.parameters()) recipe = object.__new__(_WeightedObjectiveMixin) recipe.model_parts = [original] @@ -854,9 +922,7 @@ def __init__(self): recipe._remove_text_inactive_optimizer_parameters() optimized = { - id(parameter) - for group in recipe.optimizer.param_groups - for parameter in group["params"] + id(parameter) for group in recipe.optimizer.param_groups for parameter in group["params"] } assert all(id(parameter) not in optimized for parameter in model.visual.parameters()) assert all(id(parameter) not in optimized for parameter in model.mm_projector.parameters()) @@ -950,9 +1016,7 @@ def test_llm_pp_optimizer_step_publishes_every_weighted_objective(monkeypatch): lambda self, batches, max_grad_norm: log_data, ) - recipe = object.__new__( - global_kd_recipe.KnowledgeDistillationRecipeForNextTokenPrediction - ) + recipe = object.__new__(global_kd_recipe.KnowledgeDistillationRecipeForNextTokenPrediction) recipe.needs_teacher = True recipe.pp_enabled = True recipe.device_mesh = type( @@ -973,8 +1037,7 @@ def test_llm_pp_optimizer_step_publishes_every_weighted_objective(monkeypatch): } recipe._objective_step_cursor = dict.fromkeys(recipe._objective_buffers, 0) recipe._gradient_squared = { - name: torch.tensor(0.0) - for name in ("vision", "projector", "language", "mtp") + name: torch.tensor(0.0) for name in ("vision", "projector", "language", "mtp") } recipe._dp_allreduce = lambda value, include_cp: value @@ -1032,9 +1095,7 @@ def test_global_kd_uses_canonical_multimodal_packing_and_train_all(tmp_path, mon assert kd.freeze_policy == "train_all" assert kd.teacher_descriptor == "qwen3_5" assert kd.student_descriptor == "qwen3_5" - assert recipe["dataset"]["_target_"].endswith( - "load_materialized_conversation_dataset" - ) + assert recipe["dataset"]["_target_"].endswith("load_materialized_conversation_dataset") assert recipe["dataset"]["path_or_dataset"] == str(tmp_path / "intersyn") assert recipe["packed_sequence"]["pack_size"] == 2048 assert recipe["packed_sequence"]["max_packs"] == 128 diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index 617b53fb66f..e31845e6fd4 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -19,6 +19,7 @@ import time from pathlib import Path +import pytest import yaml import puzzletron_orchestrator.adapters.sharded as sharded_module @@ -28,6 +29,7 @@ load_execution_config, load_runner_config, ) +from puzzletron_orchestrator.config import load_experiment_config from puzzletron_orchestrator.executors.baremetal import BareMetalSSHExecutor from puzzletron_orchestrator.executors.local import LocalExecutor from puzzletron_orchestrator.executors.slurm import SlurmExecutor, render_sbatch_script @@ -206,6 +208,86 @@ def test_render_sbatch_script_omits_gpu_requests_for_cpu_stage(): assert "--gpu-bind" not in srun +def _aiperf_attempt(tmp_path: Path, experiment_config: dict): + runner = RunnerEnvironment( + kind="slurm", + contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), + slurm=SlurmRunnerConfig(account="acct", partition_batch="batch"), + ) + node = StagePlanNode( + stage_id="aiperf", + strategy=ExecutionStrategy.SHARDED, + instances=2, + failure_policy=FailurePolicy.STRICT, + mesh={}, + gpus_per_instance=8, + gpus_per_node=8, + nodes=2, + total_gpus=16, + exclusive=False, + parents=("mip",), + distributed=False, + ) + plan = CampaignPlan( + experiment_config_path=str(tmp_path / "experiment.yaml"), + puzzle_dir=tmp_path / "run", + experiment_config=experiment_config, + runner=runner, + execution_defaults={"gpus_per_node": 8}, + stages=(node,), + contract_hash="contract", + ) + adapter = adapter_for_stage(node) + work_plan = adapter.plan(plan, node) + return adapter.command( + plan=plan, + node=node, + item=work_plan.items[0], + attempt_id="a1", + runner=runner, + ) + + +@pytest.mark.parametrize( + ("experiment_config", "trust_remote_code", "online_tokenizer"), + [ + ({}, False, False), + ({"model": {"trust_remote_code": True}}, True, False), + ( + { + "model": {"trust_remote_code": True}, + "aiperf": { + "trust_remote_code": False, + "allow_aiperf_v011_online_tokenizer_resolution": True, + }, + }, + False, + True, + ), + ], +) +def test_sharded_aiperf_command_applies_explicit_security_policy( + tmp_path, experiment_config, trust_remote_code, online_tokenizer +): + argv = _aiperf_attempt(tmp_path, experiment_config).command.argv + + assert ("--trust-remote-code" in argv) is trust_remote_code + assert ("--allow-aiperf-v011-online-tokenizer-resolution" in argv) is online_tokenizer + + +@pytest.mark.parametrize( + "experiment_config", + [ + {"model": {"trust_remote_code": "true"}}, + {"aiperf": {"trust_remote_code": 1}}, + {"aiperf": {"allow_aiperf_v011_online_tokenizer_resolution": "false"}}, + ], +) +def test_sharded_aiperf_command_rejects_non_boolean_security_policy(tmp_path, experiment_config): + with pytest.raises(ValueError, match="must be a boolean"): + _aiperf_attempt(tmp_path, experiment_config) + + def test_vllm_aggregation_uses_slurm_execution_contract(tmp_path: Path, monkeypatch): """Controller-side merges must run in the same container/venv as workers.""" @@ -272,10 +354,12 @@ def poll(self, handles): monkeypatch.setattr(sharded_module, "SlurmExecutor", FakeSlurmExecutor, raising=False) - def fail_local_subprocess(*_args, **_kwargs): - raise AssertionError("aggregation escaped the Slurm execution contract") + class FailLocalExecutor: + def __init__(self, *_args, **_kwargs): + del _args, _kwargs + raise AssertionError("aggregation escaped the Slurm execution contract") - monkeypatch.setattr(sharded_module.subprocess, "run", fail_local_subprocess) + monkeypatch.setattr(sharded_module, "LocalExecutor", FailLocalExecutor) adapter = adapter_for_stage(node) result = adapter.aggregate( plan=plan, @@ -298,6 +382,77 @@ def fail_local_subprocess(*_args, **_kwargs): assert result.summary["merge_handle"] == "slurm-123" +def test_aiperf_aggregation_uses_reviewed_local_executor(tmp_path: Path, monkeypatch): + runner = RunnerEnvironment( + kind="local", + contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), + ) + node = StagePlanNode( + stage_id="aiperf", + strategy=ExecutionStrategy.SHARDED, + instances=1, + failure_policy=FailurePolicy.STRICT, + mesh={}, + gpus_per_instance=8, + gpus_per_node=8, + nodes=1, + total_gpus=8, + exclusive=False, + parents=("mip",), + distributed=False, + ) + plan = CampaignPlan( + experiment_config_path=str(tmp_path / "experiment.yaml"), + puzzle_dir=tmp_path / "run", + experiment_config={"aiperf": {"profile_id": "runtime-075"}}, + runner=runner, + execution_defaults={"gpus_per_node": 8}, + stages=(node,), + contract_hash="contract", + ) + submitted = [] + + class FakeLocalExecutor: + def submit(self, attempt): + submitted.append(attempt) + return JobHandle( + backend="local", + handle_id="local-1", + attempt_id=attempt.attempt_id, + metadata={"log_paths": (attempt.command.log_path,)}, + ) + + def poll(self, handles): + return [ + JobStatus( + handle=handles[0], + state=JobState.COMPLETED, + exit_code=0, + log_paths=tuple(handles[0].metadata["log_paths"]), + ) + ] + + monkeypatch.setattr(sharded_module, "LocalExecutor", FakeLocalExecutor) + + result = adapter_for_stage(node).aggregate( + plan=plan, + node=node, + work_plan=WorkPlan( + stage_id="aiperf", + strategy=ExecutionStrategy.SHARDED, + items=(), + aggregate_required=True, + ), + ) + + assert result is not None + assert len(submitted) == 1 + attempt = submitted[0] + assert attempt.allocation_gpus == 0 + assert attempt.command.cwd == str(tmp_path) + assert attempt.command.argv[-1] == "--merge" + + def test_render_sbatch_script_never_requests_exclusive(): """Exclusive attempt metadata must not request exclusive Slurm nodes.""" @@ -680,6 +835,54 @@ def test_depth_pool_packs_four_two_gpu_workers_per_node(tmp_path: Path): assert "--gpus-per-task=2" in script +def test_depth_pool_splits_one_sixteen_gpu_worker_across_two_nodes(tmp_path: Path): + runner = RunnerEnvironment( + kind="slurm", + contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), + slurm=SlurmRunnerConfig(account="acct", partition_batch="batch"), + ) + node = StagePlanNode( + stage_id="depth_importance", + strategy=ExecutionStrategy.PERSISTENT_POOL, + instances=1, + failure_policy=FailurePolicy.STRICT, + mesh={"tp": 2, "cp": 1, "pp": 2, "ep": 2, "dp_shard": 2, "dp_replicate": 1}, + gpus_per_instance=16, + gpus_per_node=8, + nodes=2, + total_gpus=16, + exclusive=True, + parents=("tokenize_data",), + distributed=True, + partition="batch", + ) + plan = CampaignPlan( + experiment_config_path=str(tmp_path / "experiment.yaml"), + puzzle_dir=tmp_path / "run", + experiment_config={"depth_importance": {"output_dir": str(tmp_path / "depth")}}, + runner=runner, + execution_defaults={"gpus_per_node": 8}, + stages=(node,), + contract_hash="contract", + ) + + adapter = adapter_for_stage(node) + item = adapter.plan(plan, node).items[0] + attempt = adapter.command( + plan=plan, + node=node, + item=item, + attempt_id="a1", + runner=runner, + ) + + assert attempt.allocation_nodes == 2 + assert attempt.task_topology.task_count == 2 + assert attempt.task_topology.tasks_per_group == 2 + assert attempt.task_topology.gpus_per_task == 8 + assert attempt.command.env["NPROC_PER_NODE"] == "8" + + def test_post_mip_workers_share_one_packed_allocation(tmp_path: Path): runner = RunnerEnvironment( kind="slurm", @@ -904,6 +1107,21 @@ def test_replacement_pool_splits_workers_across_embedding_widths(tmp_path: Path) assert attempts[1].command.env["PUZZLE_DIR"].endswith("scenarios/width-1792/depth-00") +def test_replacement_width_overrides_compose_with_base_config(tmp_path: Path): + _, _, attempts = _replacement_width_attempts(tmp_path, [2048]) + attempt = attempts[0] + + config = load_experiment_config( + Path(__file__).parents[4] / "examples/puzzletron/configs/base.yaml", + overrides=attempt.command.env["DISTRIBUTED_EVAL_OVERRIDES"].splitlines(), + ) + + teacher = str(Path(attempt.command.env["PUZZLE_DIR"]) / "ckpts" / "sorted_teacher") + assert config["build_replacement_library"]["source_checkpoint_dir"] == teacher + assert config["replacement_scoring"]["source_checkpoint_dir"] == teacher + assert config["replacement_scoring"]["target_teacher_dir"] == teacher + + def test_replacement_pool_completion_identity_changes_with_embedding_widths(tmp_path: Path): _, _, baseline_attempts = _replacement_width_attempts(tmp_path, [2048, 1792]) _, _, changed_attempts = _replacement_width_attempts(tmp_path, [2048, 1792, 1536, 1280]) diff --git a/tests/unit/torch/puzzletron/test_orchestration_task_topology.py b/tests/unit/torch/puzzletron/test_orchestration_task_topology.py index 22791749112..7c14f1236fe 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_task_topology.py +++ b/tests/unit/torch/puzzletron/test_orchestration_task_topology.py @@ -1,8 +1,25 @@ # 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 explicit orchestration task topology.""" +import os +import shutil +import subprocess +from pathlib import Path + import pytest from puzzletron_orchestrator import task_launcher @@ -107,18 +124,27 @@ def test_resolve_task_topology_accepts_one_cpu_task() -> None: @pytest.mark.parametrize( - ("task_count", "local_task_index", "gpus_per_task", "expected"), + ( + "task_count", + "local_task_index", + "gpus_per_task", + "expected_gpus", + "wait_status", + "expected_exit_code", + ), [ - (8, 3, 1, "3"), - (4, 2, 2, "4,5"), + (8, 3, 1, "3", 0, 0), + (4, 2, 2, "4,5", 7 << 8, 7), ], ) -def test_task_launcher_slices_full_node_visibility_for_packed_container_tasks( +def test_task_launcher_slices_visibility_and_propagates_payload_exit_status( monkeypatch, task_count: int, local_task_index: int, gpus_per_task: int, - expected: str, + expected_gpus: str, + wait_status: int, + expected_exit_code: int, ) -> None: captured: dict[str, object] = {} monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1,2,3,4,5,6,7") @@ -126,10 +152,12 @@ def test_task_launcher_slices_full_node_visibility_for_packed_container_tasks( monkeypatch.setenv("PUZZLETRON_LOCAL_TASK_INDEX", str(local_task_index)) monkeypatch.setenv("PUZZLETRON_TASK_HOSTS", "node-a") - def fake_execvpe(executable, command, env) -> None: + def fake_posix_spawnp(executable, command, env) -> int: captured.update(executable=executable, command=command, env=env) + return 123 - monkeypatch.setattr(task_launcher.os, "execvpe", fake_execvpe) + monkeypatch.setattr(task_launcher.os, "posix_spawnp", fake_posix_spawnp) + monkeypatch.setattr(task_launcher.os, "waitpid", lambda pid, _options: (pid, wait_status)) result = task_launcher.main( [ @@ -153,9 +181,117 @@ def fake_execvpe(executable, command, env) -> None: ] ) - assert result == 0 - assert captured["env"]["CUDA_VISIBLE_DEVICES"] == expected + assert result == expected_exit_code + assert captured["env"]["CUDA_VISIBLE_DEVICES"] == expected_gpus assert captured["env"]["PUZZLETRON_TASK_LAUNCHER"] == "direct" + assert captured["env"]["PUZZLETRON_RENDEZVOUS_ENDPOINT"] == "localhost:0" + + +def test_task_launcher_exports_shared_multi_node_rendezvous(monkeypatch) -> None: + captured: dict[str, object] = {} + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1,2,3,4,5,6,7") + monkeypatch.setenv("PUZZLETRON_TASK_INDEX", "1") + monkeypatch.setenv("PUZZLETRON_LOCAL_TASK_INDEX", "0") + monkeypatch.setenv("PUZZLETRON_TASK_HOSTS", "node-a,node-b") + + def fake_posix_spawnp(executable, command, env) -> int: + captured.update(executable=executable, command=command, env=env) + return 123 + + monkeypatch.setattr(task_launcher.os, "posix_spawnp", fake_posix_spawnp) + monkeypatch.setattr(task_launcher.os, "waitpid", lambda pid, _options: (pid, 0)) + + assert ( + task_launcher.main( + [ + "--attempt-id", + "attempt-a", + "--nodes", + "2", + "--gpus-per-node", + "8", + "--task-count", + "2", + "--gpus-per-task", + "8", + "--tasks-per-group", + "2", + "--launcher", + "direct", + "--", + "python", + "worker.py", + ] + ) + == 0 + ) + + env = captured["env"] + assert isinstance(env, dict) + assert env["PUZZLETRON_GROUP_SIZE"] == "2" + assert env["PUZZLETRON_GROUP_RANK"] == "1" + assert env["PUZZLETRON_RENDEZVOUS_ENDPOINT"].startswith("node-a:") + assert env["PUZZLETRON_RENDEZVOUS_ID"] == "attempt-a-group-0" + + +def test_run_worker_consumes_multi_node_task_launcher_identity(tmp_path: Path) -> None: + script = Path(__file__).parents[4] / "examples/puzzletron/distributed_eval/run_worker.sh" + env = { + **os.environ, + "CAMPAIGN_DIR": str(tmp_path / "campaign"), + "CONFIG_PATH": str(tmp_path / "experiment.yaml"), + "TORCHRUN": "/bin/echo", + "NPROC_PER_NODE": "4", + "PUZZLETRON_GROUP_SIZE": "2", + "PUZZLETRON_GROUP_RANK": "1", + "PUZZLETRON_RENDEZVOUS_ENDPOINT": "node-a:23456", + "PUZZLETRON_RENDEZVOUS_ID": "attempt-a-group-0", + } + + result = subprocess.run( + ["bash", str(script)], + env=env, + check=True, + capture_output=True, + text=True, + timeout=10, + ) + + assert "--nnodes 2" in result.stdout + assert "--node-rank 1" in result.stdout + assert "--rdzv-endpoint node-a:23456" in result.stdout + assert "--rdzv-id attempt-a-group-0" in result.stdout + + +@pytest.mark.parametrize("script_name", ["run_replacement_pool.sh", "run_depth_pool.sh"]) +def test_nonzero_group_rank_does_not_own_pool_control_path( + tmp_path: Path, script_name: str +) -> None: + true_bin = shutil.which("true") + false_bin = shutil.which("false") + assert true_bin is not None + assert false_bin is not None + campaign_dir = tmp_path / "campaign" + campaign_dir.mkdir() + (campaign_dir / "manifest.json").write_text("{}\n") + script = Path(__file__).parents[4] / "examples/puzzletron/distributed_eval" / script_name + env = { + **os.environ, + "CAMPAIGN_DIR": str(campaign_dir), + "CONFIG_PATH": str(tmp_path / "experiment.yaml"), + "WORLD_SIZE": "2", + "WORKER_COUNT": "1", + "NPROC_PER_NODE": "1", + "TORCHRUN": true_bin, + "PYTHON_BIN": false_bin, + "PUZZLETRON_GROUP_INDEX": "0", + "PUZZLETRON_GROUP_RANK": "1", + "PUZZLETRON_GROUP_SIZE": "2", + "PUZZLETRON_RENDEZVOUS_ENDPOINT": "node-a:23456", + "PUZZLETRON_RENDEZVOUS_ID": "attempt-a-group-0", + } + + subprocess.run(["bash", str(script)], env=env, check=True, timeout=10) def _task_binding(*, group_size: int) -> task_launcher.TaskBinding: @@ -172,7 +308,7 @@ def _task_binding(*, group_size: int) -> task_launcher.TaskBinding: ) -def test_single_node_torchrun_uses_localhost_for_rendezvous() -> None: +def test_single_node_torchrun_lets_c10d_choose_a_free_local_port() -> None: command = task_launcher.build_task_command( payload=("python", "worker.py"), launcher=TaskLauncher.TORCHRUN, @@ -180,7 +316,7 @@ def test_single_node_torchrun_uses_localhost_for_rendezvous() -> None: gpus_per_task=4, ) - assert "--rdzv-endpoint=localhost:23456" in command + assert "--rdzv-endpoint=localhost:0" in command def test_multi_node_torchrun_uses_master_hostname_for_rendezvous() -> None: diff --git a/tests/unit/torch/puzzletron/test_post_mip_runner.py b/tests/unit/torch/puzzletron/test_post_mip_runner.py index 4754d8dfe72..555c6243c56 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -214,6 +214,7 @@ def fake_run_aiperf_sweep(checkpoint, **settings): "minimum_request_count": 4, "requests_per_concurrency": 2, "best_selection_mode": "individual_best", + "allow_aiperf_v011_online_tokenizer_resolution": True, "input_tokens": 1024, "output_tokens": 128, "topology": {"gpu_group_size": 1}, @@ -226,7 +227,7 @@ def fake_run_aiperf_sweep(checkpoint, **settings): ) result = runner._aiperf( - {"puzzle_dir": str(tmp_path)}, + {"puzzle_dir": str(tmp_path), "model": {"trust_remote_code": True}}, node, source, "execution", @@ -235,6 +236,8 @@ def fake_run_aiperf_sweep(checkpoint, **settings): assert captured["checkpoint"] == str(tmp_path / "checkpoint") assert captured["concurrencies"] == (8,) assert captured["request_counts"] == {8: 23} + assert captured["trust_remote_code"] is True + assert captured["allow_aiperf_v011_online_tokenizer_resolution"] is True assert "request_count" not in captured assert "minimum_request_count" not in captured assert "requests_per_concurrency" not in captured diff --git a/tests/unit/torch/puzzletron/test_profile_aiperf_worker.py b/tests/unit/torch/puzzletron/test_profile_aiperf_worker.py index 11a226c614f..b955b29247c 100644 --- a/tests/unit/torch/puzzletron/test_profile_aiperf_worker.py +++ b/tests/unit/torch/puzzletron/test_profile_aiperf_worker.py @@ -16,6 +16,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import pytest + def test_profile_aiperf_filters_registry_to_explicit_solutions(): from examples.puzzletron.run_profile_aiperf_worker import select_registry_solutions @@ -136,3 +138,81 @@ def test_profile_aiperf_merge_honors_explicit_concurrency_subset(tmp_path): payload = json.loads(output.read_text()) assert payload["concurrencies"] == [1] assert len(payload["results"]) == 12 + + +def test_profile_aiperf_cli_forwards_explicit_security_flags(tmp_path, monkeypatch): + import sys + + from examples.puzzletron import run_profile_aiperf_worker as worker_module + + captured = {} + + def run_worker(puzzle_dir, **kwargs): + captured["puzzle_dir"] = puzzle_dir + captured.update(kwargs) + return tmp_path / "worker.json" + + monkeypatch.setattr(worker_module, "run_worker", run_worker) + monkeypatch.setattr( + sys, + "argv", + [ + "run_profile_aiperf_worker.py", + "--puzzle-dir", + str(tmp_path), + "--trust-remote-code", + "--allow-aiperf-v011-online-tokenizer-resolution", + ], + ) + + worker_module.main() + + assert captured["puzzle_dir"] == tmp_path + assert captured["trust_remote_code"] is True + assert captured["allow_aiperf_v011_online_tokenizer_resolution"] is True + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_profile_aiperf_worker_forwards_security_policy_to_real_sweep( + tmp_path, monkeypatch, enabled +): + import json + + from examples.puzzletron.run_profile_aiperf_worker import run_worker + from modelopt.torch.puzzletron import benchmarks + + profile_id = "runtime-075" + registry_path = tmp_path / "mip/profiles" / profile_id / "selected_solutions.json" + registry_path.parent.mkdir(parents=True) + registry_path.write_text( + json.dumps( + { + "profile_id": profile_id, + "solutions": [{"solution_id": "teacher", "checkpoint": "/teacher"}], + } + ) + ) + observed = [] + + def run_aiperf_sweep(*args, **kwargs): + observed.append((args, kwargs)) + return [] + + monkeypatch.setattr(benchmarks, "run_aiperf_sweep", run_aiperf_sweep) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1,2,3,4,5,6,7") + + run_worker( + tmp_path, + profile_id=profile_id, + worker_index=0, + worker_count=6, + input_tokens=32, + output_tokens=8, + trust_remote_code=enabled, + allow_aiperf_v011_online_tokenizer_resolution=enabled, + ) + + assert len(observed) == 1 + _, kwargs = observed[0] + assert kwargs["trust_remote_code"] is enabled + assert kwargs["allow_aiperf_v011_online_tokenizer_resolution"] is enabled diff --git a/tests/unit/torch/puzzletron/test_setup_bundle.py b/tests/unit/torch/puzzletron/test_setup_bundle.py index 5f4f264ea65..48c47c43385 100644 --- a/tests/unit/torch/puzzletron/test_setup_bundle.py +++ b/tests/unit/torch/puzzletron/test_setup_bundle.py @@ -258,6 +258,13 @@ def test_custom_dataset_rendering_does_not_add_acquisition_fields() -> None: assert "acquisition" not in experiment["data"] +def test_rendered_data_keeps_controller_and_worker_sequence_length_in_sync() -> None: + experiment = render_experiment(_nemotron_render_state(latent_moe=False), "production") + + assert experiment["data"]["sequence_length"] == 2048 + assert experiment["data"]["sequence_length"] == experiment["data"]["max_sample_length"] + + def test_packed_text_uses_native_automodel_data_instead_of_fixed_token_memmaps() -> None: state = _nemotron_render_state(latent_moe=False) state["answers"]["data"]["layout"] = "packed_varlen" diff --git a/tests/unit/torch/puzzletron/test_setup_v2_quick.py b/tests/unit/torch/puzzletron/test_setup_v2_quick.py index 55819c319eb..5eea1243d83 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_quick.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_quick.py @@ -36,6 +36,7 @@ from puzzletron_setup.v2.prompts import ( BACK, InteractiveBackend, + NonInteractiveBackend, PromptChoice, ScriptedBackend, _bind_escape_back, @@ -44,7 +45,6 @@ from puzzletron_setup.v2.state import WizardState from puzzletron_setup.v2.wizard import ( _CUSTOM_DATA_SOURCE, - _CUSTOM_MODEL_SOURCE, _PUZZLE_KD_DATA_SOURCE, _acquisition_sample_requirements, _fresh_state, @@ -59,6 +59,50 @@ _NEMOTRON_FAMILY_CONFIG = "examples/puzzletron/configs/families/nemotron3/family.yaml" +def _qwen_inspected_model(model_path) -> InspectedModel: + inventory = ModelInventory( + family="qwen3_5", + descriptor="qwen3_5_text", + family_config=_QWEN_FAMILY_CONFIG, + model_type="qwen3_5_text", + architectures=("Qwen3_5ForCausalLM",), + multimodal=False, + moe=False, + num_layers=24, + num_sublayers=48, + layer_counts={"full_attention": 6, "linear_attention": 18}, + facts={ + "hidden_size": 1024, + "num_attention_heads": 8, + "num_key_value_heads": 2, + "intermediate_size": 3584, + }, + axes=( + AxisInventory( + axis_id="hidden_width", + label="Hidden width", + teacher_value=1024, + values=(1024, 768), + alignment=256, + ), + ), + ) + return InspectedModel( + source=str(model_path), + requested_revision=None, + resolved_revision=None, + is_local=True, + config={ + "model_type": "qwen3_5_text", + "text_config": { + "num_hidden_layers": 24, + "layer_types": ["linear_attention"] * 18 + ["full_attention"] * 6, + }, + }, + inventory=inventory, + ) + + def test_common_helpers_remain_available_from_wizard_facade(): names = ( "BUILTINS", @@ -364,17 +408,140 @@ def test_fresh_guided_state_records_profile_and_cli_full_is_explicit(tmp_path): assert _parser().parse_args(["--full"]).full is True -def test_cli_forwards_full_to_the_wizard(monkeypatch, tmp_path): +def test_non_interactive_backend_uses_semantic_defaults() -> None: + backend = NonInteractiveBackend() + choices = [PromptChoice("First", "first"), PromptChoice("Second", "second")] + + assert backend.text("Path:", "/resolved/path") == "/resolved/path" + assert backend.text("Optional commands:", "") == "" + assert backend.select("Choice:", choices, "second") == "second" + assert backend.checkbox("Choices:", choices, ["first"]) == ["first"] + with pytest.raises(SetupError, match="requires a default"): + backend.text("Path:", None) + + +def test_cli_forwards_noninteractive_setup_contract_to_the_wizard(monkeypatch, tmp_path): captured = {} + defaults = tmp_path / "defaults.yaml" + defaults.write_text("schema_version: 1\n") + campaign = tmp_path / "campaign" def run_wizard_v2(**kwargs): captured.update(kwargs) - return tmp_path / "campaign" + return campaign monkeypatch.setattr(wizard_module, "run_wizard_v2", run_wizard_v2) - assert cli_module.main(["--full"]) == 0 + assert ( + cli_module.main( + [ + "--non-interactive", + "--defaults", + str(defaults), + "--campaign-dir", + str(campaign), + "--profile", + "smoke", + "--full", + ] + ) + == 0 + ) + assert captured["resume"] is None + assert captured["defaults_path"] == defaults + assert captured["campaign_dir"] == campaign + assert captured["setup_profile"] == "smoke" assert captured["full"] is True + assert isinstance(captured["backend"], NonInteractiveBackend) + + +@pytest.mark.parametrize( + "argv", + [ + ["--resume", "existing", "--campaign-dir", "new"], + ["--non-interactive", "--defaults", "defaults.yaml"], + ["--non-interactive", "--campaign-dir", "campaign"], + ], +) +def test_cli_rejects_invalid_automation_argument_combinations(argv): + with pytest.raises(SystemExit) as error: + cli_module.main(argv) + + assert error.value.code == 2 + + +def test_cli_incomplete_noninteractive_defaults_fail_fast(tmp_path, capsys): + defaults = tmp_path / "defaults.yaml" + defaults.write_text("schema_version: 1\n") + + assert ( + cli_module.main( + [ + "--non-interactive", + "--defaults", + str(defaults), + "--campaign-dir", + str(tmp_path / "campaign"), + ] + ) + == 2 + ) + assert "Setup stopped: Enter a model path or Hugging Face URL." in capsys.readouterr().out + + +def test_cli_invalid_noninteractive_vllm_topology_fails_fast(tmp_path, monkeypatch, capsys): + campaign = tmp_path / "campaign" + model_path = tmp_path / "model" + dataset = tmp_path / "dataset" + model_path.mkdir() + dataset.mkdir() + inspected = _qwen_inspected_model(model_path) + monkeypatch.setattr(wizard_module, "inspect_model", lambda source: inspected) + monkeypatch.setattr( + wizard_module, + "infer_dataset_modality", + lambda source: SimpleNamespace(modality="text", evidence="local fixture"), + ) + defaults = tmp_path / "defaults.yaml" + defaults.write_text( + yaml.safe_dump( + { + "schema_version": 1, + "model": {"source": str(model_path)}, + "data": {"source": str(dataset), "modality": "text"}, + "infrastructure": { + "execution_contract": { + "repository": "/worker/modelopt", + "venv": "/worker/venv", + } + }, + "vllm": { + "enabled": True, + "topology": {"tensor_parallel_size": 3}, + }, + }, + sort_keys=False, + ) + ) + + assert ( + cli_module.main( + [ + "--non-interactive", + "--defaults", + str(defaults), + "--campaign-dir", + str(campaign), + "--profile", + "smoke", + ] + ) + == 2 + ) + output = capsys.readouterr().out + assert "Setup stopped: Non-interactive vLLM topology is incompatible" in output + assert "TP=3 is incompatible" in output + assert "valid choices [1, 2, 4, 8]" in output @pytest.mark.parametrize("full", [False, True]) @@ -825,76 +992,45 @@ def test_guided_wizard_runs_real_sections_and_generates_valid_bundles( dataset = tmp_path / "dataset" model_path.mkdir() dataset.mkdir() - inventory = ModelInventory( - family="qwen3_5", - descriptor="qwen3_5_text", - family_config="examples/puzzletron/configs/families/qwen3_5/family.yaml", - model_type="qwen3_5_text", - architectures=("Qwen3_5ForCausalLM",), - multimodal=False, - moe=False, - num_layers=24, - num_sublayers=48, - layer_counts={"full_attention": 6, "linear_attention": 18}, - facts={ - "hidden_size": 1024, - "num_attention_heads": 8, - "num_key_value_heads": 2, - "intermediate_size": 3584, - }, - axes=( - AxisInventory( - axis_id="hidden_width", - label="Hidden width", - teacher_value=1024, - values=(1024, 768), - alignment=256, - ), - ), - ) - inspected = InspectedModel( - source=str(model_path), - requested_revision=None, - resolved_revision=None, - is_local=True, - config={ - "model_type": "qwen3_5_text", - "text_config": { - "num_hidden_layers": 24, - "layer_types": ["linear_attention"] * 18 + ["full_attention"] * 6, - }, - }, - inventory=inventory, - ) + inspected = _qwen_inspected_model(model_path) monkeypatch.setattr(wizard_module, "inspect_model", lambda source: inspected) monkeypatch.setattr( wizard_module, "infer_dataset_modality", lambda source: SimpleNamespace(modality="text", evidence="local fixture"), ) - backend = ScriptedBackend( - [ - "smoke", - str(campaign), - _CUSTOM_MODEL_SOURCE, - str(model_path), - _CUSTOM_DATA_SOURCE, - str(dataset), - "defaults", - "/worker/modelopt", - "/worker/venv", - True, - ] + defaults = tmp_path / "defaults.yaml" + defaults.write_text( + yaml.safe_dump( + { + "schema_version": 1, + "model": {"source": str(model_path)}, + "data": { + "source": str(dataset), + "modality": "text", + "layout": "fixed", + "sequence_length": 32, + }, + "infrastructure": { + "execution_contract": { + "repository": "/worker/modelopt", + "venv": "/worker/venv", + } + }, + }, + sort_keys=False, + ) ) result = wizard_module.run_wizard_v2( resume=None, - defaults_path=None, - backend=backend, + defaults_path=defaults, + backend=NonInteractiveBackend(), + campaign_dir=campaign, + setup_profile="smoke", ) assert result == campaign.resolve() - assert backend.remaining == 0 assert (campaign / "smoke" / "experiment.yaml").is_file() assert (campaign / "production" / "experiment.yaml").is_file() assert (campaign / "resolved_defaults.yaml").is_file() diff --git a/tests/unit/torch/puzzletron/test_vllm_axis_contract.py b/tests/unit/torch/puzzletron/test_vllm_axis_contract.py index 08acb053a15..efa5cc77d3c 100644 --- a/tests/unit/torch/puzzletron/test_vllm_axis_contract.py +++ b/tests/unit/torch/puzzletron/test_vllm_axis_contract.py @@ -1,3 +1,20 @@ +# 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 Puzzletron vLLM geometry and checkpoint interchange contracts.""" + from types import SimpleNamespace from modelopt.torch.puzzletron.block_config import ( @@ -8,9 +25,38 @@ MoEConfig, ) from modelopt.torch.puzzletron.candidates import build_candidate_library +from modelopt.torch.puzzletron.utils import vllm_adapter from modelopt.torch.puzzletron.utils.vllm_adapter import convert_block_configs_to_per_layer_config +def test_checkpoint_config_refresh_does_not_trust_remote_code_by_default( + tmp_path, monkeypatch +) -> None: + observed = [] + config = SimpleNamespace(to_json_string=lambda **_kwargs: "{}\n") + + def load_config(path, **kwargs): + observed.append((path, kwargs)) + return config + + monkeypatch.setattr("transformers.AutoConfig.from_pretrained", load_config) + monkeypatch.setattr( + "modelopt.torch.puzzletron.anymodel.registry.resolve_descriptor", + lambda _config: SimpleNamespace(descriptor=object()), + ) + monkeypatch.setattr(vllm_adapter, "configure_anymodel_metadata", lambda *_args: True) + monkeypatch.setattr( + vllm_adapter, + "convert_block_configs_to_per_layer_config", + lambda *_args, **_kwargs: True, + ) + + config_path = vllm_adapter.refresh_realized_checkpoint_config(tmp_path) + + assert observed == [(tmp_path, {"trust_remote_code": False})] + assert config_path.read_text() == "{}\n" + + def test_mla_search_axes_create_cartesian_typed_candidates() -> None: teacher = BlockConfig( subblock_configs=(MLAConfig(num_heads=16, q_lora_rank=768, kv_lora_rank=512),) @@ -285,12 +331,10 @@ def test_qwen35_moe_descriptor_exposes_bounded_runtime_benchmark_contract() -> N flat_text_config = SimpleNamespace(hidden_size=2048) nested_vlm_config = SimpleNamespace(text_config=flat_text_config) assert ( - Qwen3P5MoeVLModelDescriptor.get_language_model_config(flat_text_config) - is flat_text_config + Qwen3P5MoeVLModelDescriptor.get_language_model_config(flat_text_config) is flat_text_config ) assert ( - Qwen3P5MoeVLModelDescriptor.get_language_model_config(nested_vlm_config) - is flat_text_config + Qwen3P5MoeVLModelDescriptor.get_language_model_config(nested_vlm_config) is flat_text_config ) base = Qwen3P5MoeTextModelDescriptor.runtime_benchmark_base_block_config(runtime) assert base.require_subblock("attention").num_query_heads == 16 From d0d090865ce43aea70b19e9b4464bbcbea811fdc Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Tue, 25 Aug 2026 15:19:28 +0200 Subject: [PATCH 2/4] Fix Puzzletron runtime integration issues Signed-off-by: Johannes Rausch --- modelopt/torch/puzzletron/__init__.py | 1 + .../distillation/global_kd_recipe.py | 55 +++++++++++----- .../orchestration/adapters/sharded.py | 2 +- puzzletron_setup/v2/wizard.py | 3 +- .../test_automodel_solution_scoring.py | 12 ++-- .../puzzletron/test_global_kd_canonical.py | 66 +++++++++++++++++-- .../test_orchestration_executors.py | 3 + .../puzzletron/test_profile_aiperf_worker.py | 53 ++++++++------- .../torch/puzzletron/test_setup_v2_quick.py | 19 ++++++ 9 files changed, 156 insertions(+), 58 deletions(-) diff --git a/modelopt/torch/puzzletron/__init__.py b/modelopt/torch/puzzletron/__init__.py index 5e4c5f7489f..ce3765a1ff1 100644 --- a/modelopt/torch/puzzletron/__init__.py +++ b/modelopt/torch/puzzletron/__init__.py @@ -39,3 +39,4 @@ tools, utils, ) +from .security_policy import * diff --git a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py index 9aceb6a9441..74daa4cffa0 100644 --- a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py +++ b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py @@ -653,25 +653,48 @@ def save_checkpoint( ): """Publish a completion marker only after model and optimizer DCP succeed.""" - result = super().save_checkpoint( # type: ignore[misc] - epoch, - step, - train_loss, - val_loss, - best_metric_key=best_metric_key, - ) + result = None + publication_error: Exception | None = None + publication_error_text: str | None = None + try: + result = super().save_checkpoint( # type: ignore[misc] + epoch, + step, + train_loss, + val_loss, + best_metric_key=best_metric_key, + ) + except Exception as error: # noqa: BLE001 - all ranks must reach the collective + publication_error = error + publication_error_text = f"{type(error).__name__}: {error}" + distributed = torch.distributed.is_initialized() + if distributed: + parent_save_errors: list[str | None] = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(parent_save_errors, publication_error_text) + parent_save_failure = next( + ( + (rank, error) + for rank, error in enumerate(parent_save_errors) + if error is not None + ), + None, + ) + if parent_save_failure is not None: + failing_rank, error_text = parent_save_failure + publication_error_text = f"parent save failed on rank {failing_rank}: {error_text}" checkpoint_path = os.path.join( str(self.checkpointer.config.checkpoint_dir), f"epoch_{epoch}_step_{step}", ) - publication_error: Exception | None = None - publication_error_text: str | None = None - if self.dist_env.is_main: + if publication_error_text is None and self.dist_env.is_main: try: consolidated = Path(checkpoint_path, "model", "consolidated") config_path = consolidated / "config.json" config = json.loads(config_path.read_text()) if config_path.is_file() else {} - if config.get("block_configs"): + text_config = config.get("text_config") + if config.get("block_configs") or ( + isinstance(text_config, dict) and text_config.get("block_configs") + ): from ..utils.vllm_adapter import refresh_realized_checkpoint_config model_config = _config_value(getattr(self, "cfg", None), "model") @@ -687,17 +710,17 @@ def save_checkpoint( Path(checkpoint_path, "saving_completed").touch() except Exception as error: # noqa: BLE001 - all ranks must reach the collective publication_error = error - publication_error_text = f"{type(error).__name__}: {error}" - if torch.distributed.is_initialized(): + publication_error_text = ( + f"publication failed on rank 0: {type(error).__name__}: {error}" + ) + if distributed: publication_status = [publication_error_text] torch.distributed.broadcast_object_list(publication_status, src=0) publication_error_text = publication_status[0] if publication_error is not None: raise publication_error if publication_error_text is not None: - raise RuntimeError( - f"global KD checkpoint publication failed on rank 0: {publication_error_text}" - ) + raise RuntimeError(f"global KD checkpoint {publication_error_text}") return result def _install_vision_observers(self, parts, *, role: str): diff --git a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py index 692859cbace..a4047ca48b2 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py @@ -154,7 +154,7 @@ def _run_local_aggregate( metadata={"gpus_per_node": 0}, task_topology=TaskTopology(task_count=1, gpus_per_task=0), ) - executor = LocalExecutor() + executor = LocalExecutor(plan.runner) handle = executor.submit(attempt) while True: status = executor.poll([handle])[0] diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index f0be4067f16..79b32be80a4 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -4505,13 +4505,14 @@ def _fresh_state( setup_profile: str = "balanced", ) -> WizardState: if campaign_dir is not None: + get_setup_preset(setup_profile) if full: return WizardState.start( Path(campaign_dir).expanduser(), defaults_path=defaults_path, setup_mode="full", + preset=setup_profile, ) - get_setup_preset(setup_profile) return WizardState.start( Path(campaign_dir).expanduser(), defaults_path=defaults_path, diff --git a/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py b/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py index 6d3ae235a16..9ffa77001a8 100644 --- a/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py +++ b/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py @@ -34,6 +34,9 @@ MLAConfig, MoEConfig, ) +from modelopt.torch.puzzletron.distributed_eval.automodel_executor import ( + AutoModelReplaceBlockExecutor, +) from modelopt.torch.puzzletron.plugins.automodel import solution_recipe from modelopt.torch.puzzletron.plugins.automodel.solution_launch import ( _can_skip_parent_model_load, @@ -62,6 +65,7 @@ from modelopt.torch.puzzletron.pruning.gated_delta_net import GDNShape from modelopt.torch.puzzletron.pruning.runtime_candidate import apply_runtime_candidate from modelopt.torch.puzzletron.stages.diagnostics import _annotate_solution_selections +from modelopt.torch.utils import distributed as dist def test_baseline_only_scoring_does_not_require_candidate_solutions(tmp_path): @@ -254,9 +258,6 @@ def test_rpc_executor_infers_descriptor_when_config_has_no_override(monkeypatch, def test_rpc_executor_scores_cumulative_depth_removals(monkeypatch): - from modelopt.torch.puzzletron.distributed_eval.automodel_executor import ( - AutoModelReplaceBlockExecutor, - ) from modelopt.torch.puzzletron.distributed_eval.schema import EvaluationRequest teacher_blocks = [ @@ -316,11 +317,6 @@ def test_rpc_executor_scores_cumulative_depth_removals(monkeypatch): def test_rpc_executor_non_output_pipeline_rank_reaches_collective(monkeypatch): - from modelopt.torch.puzzletron.distributed_eval.automodel_executor import ( - AutoModelReplaceBlockExecutor, - ) - from modelopt.torch.utils import distributed as dist - class NonOutputRecipe: has_outputs = False _groups = None diff --git a/tests/unit/torch/puzzletron/test_global_kd_canonical.py b/tests/unit/torch/puzzletron/test_global_kd_canonical.py index 2177f960c92..d1eb0868ad9 100644 --- a/tests/unit/torch/puzzletron/test_global_kd_canonical.py +++ b/tests/unit/torch/puzzletron/test_global_kd_canonical.py @@ -724,7 +724,14 @@ def __init__(self): assert all(value.item() > 0 for value in recipe._gradient_squared.values()) -def test_global_kd_checkpoint_forwards_best_metric_key(tmp_path, monkeypatch): +@pytest.mark.parametrize( + "checkpoint_config", + [ + {"block_configs": [{"subblock_configs": []}]}, + {"text_config": {"block_configs": [{"subblock_configs": []}]}}, + ], +) +def test_global_kd_checkpoint_forwards_best_metric_key(tmp_path, monkeypatch, checkpoint_config): # Lazy import keeps the optional NeMo AutoModel runtime out of test collection. from modelopt.torch.puzzletron.distillation.global_kd_recipe import _WeightedObjectiveMixin @@ -744,9 +751,7 @@ def save_checkpoint( checkpoint = tmp_path / f"epoch_{epoch}_step_{step}" consolidated = checkpoint / "model/consolidated" consolidated.mkdir(parents=True) - (consolidated / "config.json").write_text( - json.dumps({"block_configs": [{"subblock_configs": []}]}) - ) + (consolidated / "config.json").write_text(json.dumps(checkpoint_config)) return "saved" class Recipe(_WeightedObjectiveMixin, BaseRecipe): @@ -792,8 +797,11 @@ def test_global_kd_checkpoint_publication_failure_reaches_all_ranks(tmp_path, mo from modelopt.torch.puzzletron.distillation.global_kd_recipe import _WeightedObjectiveMixin publication = {"error": None, "broadcasts": 0} + parent_save_errors = [None, None] class BaseRecipe: + fail_rank = None + def save_checkpoint( self, epoch, @@ -803,6 +811,9 @@ def save_checkpoint( best_metric_key="default", ): del train_loss, val_loss, best_metric_key + rank = 0 if self.dist_env.is_main else 1 + if self.fail_rank == rank: + raise OSError("parent save failed") consolidated = tmp_path / f"epoch_{epoch}_step_{step}/model/consolidated" consolidated.mkdir(parents=True, exist_ok=True) (consolidated / "config.json").write_text( @@ -822,6 +833,12 @@ def broadcast_object_list(payload, *, src): payload[0] = publication["error"] monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2) + monkeypatch.setattr( + torch.distributed, + "all_gather_object", + lambda output, _value: output.__setitem__(slice(None), parent_save_errors), + ) monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list) monkeypatch.setattr( "modelopt.torch.puzzletron.utils.vllm_adapter.refresh_realized_checkpoint_config", @@ -841,7 +858,10 @@ def recipe(*, is_main): with pytest.raises(ValueError, match="refresh failed"): recipe(is_main=True).save_checkpoint(2, 19, 0.5, {"lm_loss": 0.25}) - assert publication == {"error": "ValueError: refresh failed", "broadcasts": 1} + assert publication == { + "error": "publication failed on rank 0: ValueError: refresh failed", + "broadcasts": 1, + } assert not (tmp_path / "epoch_2_step_19/saving_completed").exists() with pytest.raises( @@ -851,6 +871,42 @@ def recipe(*, is_main): recipe(is_main=False).save_checkpoint(2, 19, 0.5, {"lm_loss": 0.25}) assert publication["broadcasts"] == 2 + publication.update(error=None, broadcasts=0) + parent_save_errors[:] = ["OSError: parent save failed", None] + BaseRecipe.fail_rank = 0 + with pytest.raises(OSError, match="parent save failed"): + recipe(is_main=True).save_checkpoint(2, 23, 0.5, {"lm_loss": 0.25}) + assert publication == { + "error": "parent save failed on rank 0: OSError: parent save failed", + "broadcasts": 1, + } + assert not (tmp_path / "epoch_2_step_23/saving_completed").exists() + + with pytest.raises( + RuntimeError, + match="global KD checkpoint parent save failed on rank 0: OSError: parent save failed", + ): + recipe(is_main=False).save_checkpoint(2, 23, 0.5, {"lm_loss": 0.25}) + assert publication["broadcasts"] == 2 + + publication.update(error=None, broadcasts=0) + parent_save_errors[:] = [None, "OSError: parent save failed"] + BaseRecipe.fail_rank = 1 + with pytest.raises( + RuntimeError, + match="global KD checkpoint parent save failed on rank 1: OSError: parent save failed", + ): + recipe(is_main=True).save_checkpoint(2, 29, 0.5, {"lm_loss": 0.25}) + assert publication == { + "error": "parent save failed on rank 1: OSError: parent save failed", + "broadcasts": 1, + } + assert not (tmp_path / "epoch_2_step_29/saving_completed").exists() + + with pytest.raises(OSError, match="parent save failed"): + recipe(is_main=False).save_checkpoint(2, 29, 0.5, {"lm_loss": 0.25}) + assert publication["broadcasts"] == 2 + def test_global_kd_optimizer_save_uses_the_actual_pipeline_model_parts(): import torch diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index e31845e6fd4..eff5c957100 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -413,6 +413,9 @@ def test_aiperf_aggregation_uses_reviewed_local_executor(tmp_path: Path, monkeyp submitted = [] class FakeLocalExecutor: + def __init__(self, configured_runner): + assert configured_runner is runner + def submit(self, attempt): submitted.append(attempt) return JobHandle( diff --git a/tests/unit/torch/puzzletron/test_profile_aiperf_worker.py b/tests/unit/torch/puzzletron/test_profile_aiperf_worker.py index b955b29247c..e9b4fb0c433 100644 --- a/tests/unit/torch/puzzletron/test_profile_aiperf_worker.py +++ b/tests/unit/torch/puzzletron/test_profile_aiperf_worker.py @@ -16,12 +16,27 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +"""Tests for sharded AIPerf worker selection, merging, and policy forwarding.""" + +import json +import sys + import pytest +from examples.puzzletron import run_profile_aiperf_worker as worker_module +from examples.puzzletron.run_profile_aiperf_worker import ( + TOPOLOGIES, + build_work_items, + expected_result_count, + merge_results, + run_worker, + select_registry_solutions, + shard_work, +) +from modelopt.torch.puzzletron import benchmarks -def test_profile_aiperf_filters_registry_to_explicit_solutions(): - from examples.puzzletron.run_profile_aiperf_worker import select_registry_solutions +def test_profile_aiperf_filters_registry_to_explicit_solutions(): registry = { "profile_id": "latency-095", "solutions": [ @@ -39,8 +54,6 @@ def test_profile_aiperf_filters_registry_to_explicit_solutions(): def test_profile_aiperf_work_matrix_uses_six_all_eight_gpu_topologies(): - from examples.puzzletron.run_profile_aiperf_worker import build_work_items - registry = { "profile_id": "params-080", "solutions": [ @@ -58,8 +71,6 @@ def test_profile_aiperf_work_matrix_uses_six_all_eight_gpu_topologies(): def test_profile_aiperf_work_shards_cover_every_item_once(): - from examples.puzzletron.run_profile_aiperf_worker import shard_work - items = [{"id": value} for value in range(15)] shards = [shard_work(items, worker_index=index, worker_count=8) for index in range(8)] @@ -68,8 +79,6 @@ def test_profile_aiperf_work_shards_cover_every_item_once(): def test_profile_aiperf_expected_results_follow_registry_size(): - from examples.puzzletron.run_profile_aiperf_worker import expected_result_count - registry = { "solutions": [ {"solution_id": "best-loss"}, @@ -81,10 +90,6 @@ def test_profile_aiperf_expected_results_follow_registry_size(): def test_profile_aiperf_merge_honors_explicit_concurrency_subset(tmp_path): - import json - - from examples.puzzletron.run_profile_aiperf_worker import TOPOLOGIES, merge_results - profile_id = "latency-095" solutions = ("teacher", "h4096-d4") registry_path = tmp_path / "mip/profiles" / profile_id / "selected_solutions.json" @@ -141,10 +146,6 @@ def test_profile_aiperf_merge_honors_explicit_concurrency_subset(tmp_path): def test_profile_aiperf_cli_forwards_explicit_security_flags(tmp_path, monkeypatch): - import sys - - from examples.puzzletron import run_profile_aiperf_worker as worker_module - captured = {} def run_worker(puzzle_dir, **kwargs): @@ -172,15 +173,13 @@ def run_worker(puzzle_dir, **kwargs): assert captured["allow_aiperf_v011_online_tokenizer_resolution"] is True -@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize( + ("trust_remote_code", "online_tokenizer"), + [(False, False), (True, False), (False, True), (True, True)], +) def test_profile_aiperf_worker_forwards_security_policy_to_real_sweep( - tmp_path, monkeypatch, enabled + tmp_path, monkeypatch, trust_remote_code, online_tokenizer ): - import json - - from examples.puzzletron.run_profile_aiperf_worker import run_worker - from modelopt.torch.puzzletron import benchmarks - profile_id = "runtime-075" registry_path = tmp_path / "mip/profiles" / profile_id / "selected_solutions.json" registry_path.parent.mkdir(parents=True) @@ -208,11 +207,11 @@ def run_aiperf_sweep(*args, **kwargs): worker_count=6, input_tokens=32, output_tokens=8, - trust_remote_code=enabled, - allow_aiperf_v011_online_tokenizer_resolution=enabled, + trust_remote_code=trust_remote_code, + allow_aiperf_v011_online_tokenizer_resolution=online_tokenizer, ) assert len(observed) == 1 _, kwargs = observed[0] - assert kwargs["trust_remote_code"] is enabled - assert kwargs["allow_aiperf_v011_online_tokenizer_resolution"] is enabled + assert kwargs["trust_remote_code"] is trust_remote_code + assert kwargs["allow_aiperf_v011_online_tokenizer_resolution"] is online_tokenizer diff --git a/tests/unit/torch/puzzletron/test_setup_v2_quick.py b/tests/unit/torch/puzzletron/test_setup_v2_quick.py index 5eea1243d83..4e1af98ce92 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_quick.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_quick.py @@ -407,6 +407,25 @@ def test_fresh_guided_state_records_profile_and_cli_full_is_explicit(tmp_path): assert _parser().parse_args([]).full is False assert _parser().parse_args(["--full"]).full is True + full_state = _fresh_state( + ScriptedBackend([]), + None, + full=True, + campaign_dir=tmp_path / "full-campaign", + setup_profile="smoke", + ) + assert full_state.setup_mode == "full" + assert full_state.preset == "smoke" + + with pytest.raises(SetupError, match="Unknown setup preset"): + _fresh_state( + ScriptedBackend([]), + None, + full=True, + campaign_dir=tmp_path / "invalid-campaign", + setup_profile="invalid", + ) + def test_non_interactive_backend_uses_semantic_defaults() -> None: backend = NonInteractiveBackend() From e39e158c7b052837c317658879856ec91d18914f Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Tue, 25 Aug 2026 16:21:15 +0200 Subject: [PATCH 3/4] Define Puzzletron package exports Signed-off-by: Johannes Rausch --- modelopt/torch/puzzletron/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/modelopt/torch/puzzletron/__init__.py b/modelopt/torch/puzzletron/__init__.py index ce3765a1ff1..12b15da474f 100644 --- a/modelopt/torch/puzzletron/__init__.py +++ b/modelopt/torch/puzzletron/__init__.py @@ -40,3 +40,4 @@ utils, ) from .security_policy import * +from .security_policy import __all__ as __all__ From 901e7a123e4681bddf92c6ab65d8090116d24285 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Tue, 25 Aug 2026 16:35:05 +0200 Subject: [PATCH 4/4] Keep Puzzletron security policy internal Signed-off-by: Johannes Rausch --- modelopt/torch/puzzletron/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/modelopt/torch/puzzletron/__init__.py b/modelopt/torch/puzzletron/__init__.py index 12b15da474f..5e4c5f7489f 100644 --- a/modelopt/torch/puzzletron/__init__.py +++ b/modelopt/torch/puzzletron/__init__.py @@ -39,5 +39,3 @@ tools, utils, ) -from .security_policy import * -from .security_policy import __all__ as __all__