From 9e6c9aa094cc925f1816581ae0d17b8698b4dada Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Mon, 7 Sep 2026 15:01:19 +0200 Subject: [PATCH 1/4] Make Puzzletron first runs reliable Signed-off-by: Johannes Rausch --- examples/puzzletron/README.md | 243 ++++++++-------- .../qwen3p5_0p8b/runs/full_vlm_smoke.yaml | 2 + .../qwen3p5_0p8b/runs/vlm_campaign.yaml | 5 +- .../qwen3_5/qwen3p5_0p8b/vlm_base.yaml | 6 +- .../qwen3p5_0p8b/vlm_quality_evaluation.yaml | 1 + .../families/qwen3_5/setup_v2_defaults.yaml | 25 ++ .../qwen3p5_0p8b/execution.vlm_campaign.yaml | 4 +- examples/puzzletron/docs/environment_setup.md | 20 +- .../docs/orchestration_operations.md | 9 +- .../puzzletron/docs/qwen3p5_0p8b_smoke.md | 11 +- .../puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md | 70 +++-- examples/puzzletron/docs/setup_wizard.md | 3 +- .../puzzletron/docs/slurm_configuration.md | 8 + .../puzzletron/evaluation/vlm/evaluator.py | 31 +- .../puzzletron/evaluation/vlm/post_mip.py | 1 + .../vlm/preparation/benchmark_data.py | 28 +- modelopt/torch/puzzletron/evaluation/lmms.py | 270 +++++++++++++++++- .../puzzletron/orchestration/controller.py | 63 +++- .../puzzletron/orchestration/dashboard.py | 15 +- .../puzzletron/orchestration/progress.py | 28 ++ .../puzzletron/orchestration/reporting.py | 17 +- .../plugins/automodel/solution_launch.py | 138 +++++---- .../torch/puzzletron/post_mip/reporting.py | 143 +++++++++- .../puzzletron/pruning/compact_runtime.py | 152 ++++++++-- modelopt/torch/puzzletron/scoring.py | 41 ++- .../torch/puzzletron/stages/diagnostics.py | 34 ++- .../puzzletron/tools/validation_utils.py | 54 +++- puzzletron_setup/inspection.py | 62 +++- puzzletron_setup/v2/bundle.py | 65 ++++- puzzletron_setup/v2/post_mip.py | 3 +- puzzletron_setup/v2/wizard.py | 4 + .../vlm/preparation/test_benchmark_data.py | 46 ++- .../evaluation/vlm/test_post_mip.py | 3 + .../test_automodel_solution_scoring.py | 60 +++- .../torch/puzzletron/test_compact_runtime.py | 15 +- .../test_hidden_width_diagnostic.py | 19 ++ .../torch/puzzletron/test_lmms_evaluation.py | 84 +++++- .../test_orchestration_reporting.py | 14 + .../test_orchestration_shutdown_progress.py | 65 +++++ .../puzzletron/test_post_mip_reporting.py | 101 +++++++ .../test_qwen3p5_0p8b_full_vlm_smoke_plan.py | 21 +- .../torch/puzzletron/test_setup_inspection.py | 93 +++++- .../torch/puzzletron/test_setup_v2_data.py | 3 + .../puzzletron/test_setup_v2_post_mip.py | 7 +- .../torch/puzzletron/test_setup_v2_quick.py | 108 ++++++- 45 files changed, 1852 insertions(+), 343 deletions(-) create mode 100644 tests/unit/torch/puzzletron/test_post_mip_reporting.py diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 6e589c526dd..2f14ff0bac9 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -7,164 +7,123 @@ evaluate, benchmark, materialize, or distill those candidates. ## Table of contents -- [Start here: first campaign](#start-here-first-campaign) +- [Start here: lifecycle smoke](#start-here-lifecycle-smoke) - [Understand the campaign stages](#understand-the-campaign-stages) - [Evaluate a checkpoint](#evaluate-a-checkpoint) - [Configure a campaign](#configure-a-campaign) - [Operate and recover a campaign](#operate-and-recover-a-campaign) - [Extend Puzzletron](#extend-puzzletron) -## Start here: first campaign +## Start here: lifecycle smoke -A first Puzzletron run consists of preparing the environments, generating a -campaign, preparing its data, inspecting a dry run, and launching it. The same -launch command resumes compatible work. +Start with the maintained Qwen 3.5 0.8B VLM smoke. It exercises the complete +image-text lifecycle with bounded workloads and at most one GPU per stage. Once +it succeeds, the optional longer example campaign uses the same environment, +runner, orchestrator, progress display, report, and resume command. That larger +campaign is a scheduled, multi-hour example, not a routine smoke or presubmit. -For a maintained image-text example using Qwen 3.5 0.8B and Nemotron-VLM, -follow the -[Qwen VLM pruning smoke](docs/qwen3p5_0p8b_vlm_smoke.md). -For a larger, FFN-only example that keeps evaluation and distillation opt-in, -see the [Qwen 3.5 4B VLM example](docs/qwen3p5_4b_vlm_example.md). +### 1. Create the controller environment -### 1. Prepare the environments - -Create one lightweight Python environment for the setup wizard and the command -that launches campaigns: +From the ModelOpt checkout, use Python 3.10 through 3.14 to create a lightweight +virtual environment for the command that plans, launches, and resumes the +smoke: ```bash +python3 --version # must report Python 3.10 through 3.14 python3 -m venv .venv-puzzletron source .venv-puzzletron/bin/activate +python -m pip install --upgrade pip python -m pip install -r examples/puzzletron/requirements-setup.txt ``` -This environment creates campaign files and runs `orchestrate.py`. Model -conversion, training, evaluation, and benchmarking run in the worker -environment or container selected during setup. Prepare the -[worker environment](docs/environment_setup.md) before launching a campaign. - -### Worker image +This controller environment does not need PyTorch or model weights. Model +conversion, evaluation, serving, and distillation run in the worker environment +selected by the runner. -The repository includes a Dockerfile and build command for Puzzletron workers. -Follow the [image guide](docs/worker_image.md) to build, check, or export an -Enroot/Pyxis SquashFS image. +### 2. Configure the worker runner -### 2. Generate a campaign - -Start the guided setup with the repository defaults: +Use a reviewed Puzzletron worker image supplied by your site. Create one runner +file from the template if your site does not already provide one: ```bash -python examples/puzzletron/puzzletron_setup_v2.py \ - --defaults examples/puzzletron/configs/setup/defaults.example.yaml +cp examples/puzzletron/configs/orchestration/runner.slurm.example.yaml \ + runner.slurm.yaml ``` -Choose **Balanced pruning** for a first-generated text campaign. For the -maintained Qwen text route, select `Qwen/Qwen3.5-0.8B` and the Puzzle-KD v2 -text dataset or an existing worker-visible dataset. For VLM, select the same -model and the Nemotron-VLM v2 image-text dataset. -The generated post-MIP flow follows the detected modality, including image -serving, VLM distillation, and the pinned RealWorldQA/MMMU comparison. -Review the detected model, data modality, worker and scheduler settings, and -output directory. - -The wizard reads model configuration, not model weights, and does not submit -jobs. It writes a smoke bundle under `smoke/`, a campaign bundle -under `production/`, and a generated `README.md` with the commands for both. -Users do not need to assemble or edit a separate smoke configuration. Named MIP -bundles derive the teacher embedding width and layer depth from the inspected -model configuration. Setup stops with -a model-inspection error instead of emitting values that require hand edits. -For Qwen 3.5 0.8B, both bundles include -a final pinned student-versus-teacher downstream comparison that records -results without enforcing a minimum score. See the -[setup wizard guide](docs/setup_wizard.md) for profiles, hosted datasets, full -configuration mode, generated files, and setup resume. - -The checked-in Qwen 3.5 0.8B recipes provide one complete integration check per -modality and one illustrative VLM campaign: - -| Scope | Recipe | -| --- | --- | -| Text lifecycle smoke | `full_smoke.yaml` | -| VLM lifecycle smoke | `full_vlm_smoke.yaml` | -| VLM campaign | `vlm_campaign.yaml` | - -The smoke recipes use the shared `execution.single_gpu.yaml` profile. The VLM -campaign uses its model-specific execution profile. See the -[Qwen VLM example](docs/qwen3p5_0p8b_vlm_smoke.md) for the runnable commands. - -Unit tests compile every listed recipe and verify its stage and resource -contract. This is plan-level validation: it detects configuration and -orchestration drift, but it does not prove that the current model, data, -evaluator, and GPU runtime complete successfully together. Treat a recipe as -runtime-validated only when a recent end-to-end run is available for the same -dependencies. - -### 3. Prepare datasets and caches - -Run the dataset-preparation command in the generated `README.md` from the -[worker environment](docs/environment_setup.md#worker-environment). The command -writes into the dataset and cache paths selected during setup and can be rerun -to validate or resume preparation. For VLM evaluation caches prepared outside a -campaign, follow [cache benchmark data](docs/vlm_checkpoint_evaluation.md#cache-benchmark-data). - -### 4. Validate the generated setup - -Activate `.venv-puzzletron` and inspect the generated smoke plan: +Set the Slurm account and partitions, worker image, and any shared-storage +mounts in `runner.slurm.yaml`. Cache, dataset, and run-root paths must be +worker-visible at their configured locations. The runner's repository and +Python paths must exist inside the worker image. See [environment +setup](docs/environment_setup.md) for the worker contract and [Slurm +configuration](docs/slurm_configuration.md) for each runner field. -```bash -PUZZLETRON_BUNDLE=/path/to/generated/campaign/smoke +The three campaign inputs have separate jobs: -python examples/puzzletron/orchestrate.py \ - --experiment "$PUZZLETRON_BUNDLE/experiment.yaml" \ - --runner "$PUZZLETRON_BUNDLE/runner.yaml" \ - --execution "$PUZZLETRON_BUNDLE/execution.yaml" \ - --stage full --dry-run +- The experiment file defines the model, data, and bounded smoke workload. Do + not edit it for this first run. +- The runner file contains the site-specific scheduler and worker settings. +- The execution file maps the smoke stages onto one GPU or a CPU worker. + +### 3. Select shared paths + +Set a Hugging Face cache visible to workers and a new run root on shared +storage: + +```bash +export HF_HOME=/path/to/shared/huggingface-cache +export PUZZLETRON_SOURCE_REVISION="$(git rev-parse HEAD)" +export PUZZLETRON_DATASET_REVISION=51f4f4d219315c3283950994d4eb3d7fc30aa87b +export PUZZLETRON_RUN_ROOT=/path/to/qwen3p5_0p8b_vlm_smoke +export PUZZLETRON_RUN_ROOT="$(python -c 'import os; print(os.path.realpath(os.environ["PUZZLETRON_RUN_ROOT"]))')" ``` -Review the stage list, worker paths, resources, and output directory. -`--stage full` runs every stage enabled by the generated validation experiment -in dependency order. Remove `--dry-run` to launch it. This smoke run is the -default setup check; detailed smoke limits and manually maintained smoke recipes -are documented for advanced use and reproducibility in the -[text](docs/qwen3p5_0p8b_smoke.md) and -[VLM](docs/qwen3p5_0p8b_vlm_smoke.md) guides. +The smoke's first stage downloads, validates, and caches its eight Nemotron-VLM +image-conversation samples and the pinned evaluation rows. No PyTorch data +command or separate cache script runs in the controller venv. Offline workers +must use a cache prepared as described in the +[VLM evaluation guide](docs/vlm_checkpoint_evaluation.md#cache-benchmark-data). -### 5. Launch the campaign +### 4. Inspect and run the smoke -After validation succeeds, inspect and launch the production bundle: +Activate `.venv-puzzletron`, select the maintained files, and inspect the exact +plan before it requests resources: ```bash -PUZZLETRON_BUNDLE=/path/to/generated/campaign/production +EXPERIMENT=examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/full_vlm_smoke.yaml +EXECUTION=examples/puzzletron/configs/orchestration/execution.single_gpu.yaml +RUNNER=runner.slurm.yaml python examples/puzzletron/orchestrate.py \ - --experiment "$PUZZLETRON_BUNDLE/experiment.yaml" \ - --runner "$PUZZLETRON_BUNDLE/runner.yaml" \ - --execution "$PUZZLETRON_BUNDLE/execution.yaml" \ + --experiment "$EXPERIMENT" \ + --runner "$RUNNER" \ + --execution "$EXECUTION" \ --stage full --dry-run ``` -Launch the inspected plan: +Review the resolved paths, CPU and GPU requests, worker image, mounts, and log +locations. Then run the same plan without `--dry-run`: ```bash python examples/puzzletron/orchestrate.py \ - --experiment "$PUZZLETRON_BUNDLE/experiment.yaml" \ - --runner "$PUZZLETRON_BUNDLE/runner.yaml" \ - --execution "$PUZZLETRON_BUNDLE/execution.yaml" \ + --experiment "$EXPERIMENT" \ + --runner "$RUNNER" \ + --execution "$EXECUTION" \ --stage full ``` -The experiment file defines the model and algorithm choices, the runner file -defines the worker environment, and the execution file says how each stage -runs. Keep all three together when launching or resuming a generated campaign. +### 5. Resume and inspect results -### 6. Resume and inspect results +Run the launch command above again to recover an interrupted smoke or verify a +completed one. Compatible completed stages are not submitted again. Puzzletron +stores its structured runtime state under +`$PUZZLETRON_RUN_ROOT/orchestration/`. -The experiment file calls the campaign output directory `puzzle_dir`. -`orchestrate.py` shows live progress and stores resume information under -`/orchestration/`. Run the same command with the same three files to -recover a detached or interrupted campaign; completed compatible stages are not -submitted again. Use the production command above without `--dry-run` for both -the first launch and every resume. +While the command is running, it shows completed/total stages, queued and +running stages, elapsed time, 30-second heartbeats, and the active log path. +Stages expose their native inner unit when available, including evaluation +samples and measured rate. An ETA appears only after progress provides a +reliable denominator and the controller observes throughput; otherwise it says +`ETA unavailable`. After the selected plan completes cleanly, `orchestrate.py` attempts to write the final report to @@ -176,6 +135,33 @@ stages, `--once`, logging controls, security options, and recovery details, or report. For a failed or interrupted run, follow the actionable checks in [run and recovery options](docs/orchestration_operations.md#progress-and-interruption). +The [Qwen VLM example guide](docs/qwen3p5_0p8b_vlm_smoke.md) lists the smoke's +expected lifecycle checks and explains how to interpret its bounded results. + +### 6. Optional: run the longer multi-hour example + +After the smoke succeeds, keep the controller venv and runner and select a new +run root plus the longer example files: + +```bash +EXPERIMENT=examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/vlm_campaign.yaml +EXECUTION=examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_campaign.yaml +export PUZZLETRON_RUN_ROOT=/path/to/qwen3p5_0p8b_vlm_campaign +export PUZZLETRON_RUN_ROOT="$(python -c 'import os; print(os.path.realpath(os.environ["PUZZLETRON_RUN_ROOT"]))')" + +python examples/puzzletron/orchestrate.py \ + --experiment "$EXPERIMENT" \ + --runner "$RUNNER" \ + --execution "$EXECUTION" \ + --stage full --dry-run +``` + +Inspect the larger plan, then remove `--dry-run` to launch it. Use that same +launch command for every resume. The example increases sample counts, +candidate coverage, and distillation work, but does not introduce another +operational path. Use focused tests and the lifecycle smoke for routine +development; reserve this campaign for scheduled integration validation. + ## Understand the campaign stages `--stage full` runs every stage enabled by the experiment in dependency order. @@ -184,8 +170,9 @@ The complete pipeline is organized into these steps: 1. **Prepare inputs.** Convert the source checkpoint and, when configured, tokenize the campaign dataset. -2. **Measure pruning choices.** Collect width importance and optional depth - importance or vLLM runtime statistics, then sort the teacher checkpoint. +2. **Measure pruning choices.** Collect importance for hidden width (the + model's transformer hidden size) and optional depth importance or vLLM + runtime statistics, then sort the teacher checkpoint. 3. **Validate and score.** Run the enabled sorting, width, slicing, and bypass sanity checks; collect bypass observations; build the replacement library; and score individual replacements. @@ -229,9 +216,21 @@ runtime constraints, results, and troubleshooting for their respective routes. ## Configure a campaign -Generated bundles contain an experiment, runner, and execution file. The -default route above is sufficient when their generated values match your model -and infrastructure. +Use the maintained smoke and longer example above to learn the workflow before +creating a different campaign. The setup wizard is the customization path for +another model, dataset, search profile, or execution environment: + +```bash +python examples/puzzletron/puzzletron_setup_v2.py \ + --defaults examples/puzzletron/configs/setup/defaults.example.yaml +``` + +The wizard is implemented by `puzzletron_setup_v2.py`. It inspects model +configuration metadata and writes validated smoke and production bundles. Each +bundle contains the same three inputs used above: experiment, runner, and +execution. For named width and depth search, setup derives the teacher hidden +size and layer count and stops on missing metadata instead of asking users to +edit generated YAML. - Use [configuration and overrides](docs/configuration_overrides.md) to find the built-in configuration files, choose where campaign outputs are stored, @@ -241,8 +240,8 @@ and infrastructure. at their teacher values. - Use [Slurm configuration](docs/slurm_configuration.md) to change partitions, CPU-routed stages, log locations, and accepted compatibility fields. -- Use the [setup wizard guide](docs/setup_wizard.md) to change profiles, - datasets, generated files, or setup automation. +- Use the [setup wizard guide](docs/setup_wizard.md) for profiles, generated + files, non-interactive setup, and setup resume. Run `--dry-run` after every configuration change. It resolves and validates the experiment, runner, and execution files before any job is submitted. diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/full_vlm_smoke.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/full_vlm_smoke.yaml index 9653f5ebb13..cc7f4d1fc56 100644 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/full_vlm_smoke.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/full_vlm_smoke.yaml @@ -74,6 +74,8 @@ post_mip: extra_vllm_args: - -cc.cudagraph_mode=NONE - --no-enable-flashinfer-autotune + - --gdn-prefill-backend + - triton - --gpu-memory-utilization - "0.5" - --reasoning-parser diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/vlm_campaign.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/vlm_campaign.yaml index e7e128bad8a..6f127c6a095 100644 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/vlm_campaign.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/vlm_campaign.yaml @@ -52,7 +52,8 @@ depth_importance: max_subblocks_to_remove: 2 replacement_scoring: - eval_samples: 64 + # Covers every replacement candidate while fitting the four-hour single-GPU stage contract. + eval_samples: 16 search_space: axes: @@ -94,6 +95,8 @@ vlm_campaign_aiperf: extra_vllm_args: - -cc.cudagraph_mode=NONE - --no-enable-flashinfer-autotune + - --gdn-prefill-backend + - triton - --gpu-memory-utilization - "0.5" - --reasoning-parser diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/vlm_base.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/vlm_base.yaml index 86bb249b27e..9165a10cdb5 100644 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/vlm_base.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/vlm_base.yaml @@ -113,9 +113,9 @@ sort: sort_sanity: packed_token_cache_path: # Exact FFN channel permutations change BF16 reduction order and produced - # sub-0.002 loss drift on the VLM smoke batch; keep the gate symmetric. - max_abs_lm_loss_delta: 0.002 - max_abs_reverse_lm_loss_delta: 0.002 + # sub-0.003 loss drift on the maintained VLM batches; keep the gate symmetric. + max_abs_lm_loss_delta: 0.003 + max_abs_reverse_lm_loss_delta: 0.003 automodel: # Keep the portable smoke independent of the optional FLA fused-KL backend. lm_head_backend: streaming diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/vlm_quality_evaluation.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/vlm_quality_evaluation.yaml index 5793153df3a..6a646737f4e 100644 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/vlm_quality_evaluation.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/vlm_quality_evaluation.yaml @@ -64,6 +64,7 @@ vlm_full_smoke_evaluation: batch_size: ${vlm_smoke_evaluation.batch_size} timeout_seconds: ${vlm_smoke_evaluation.timeout_seconds} dtype: bfloat16 + gdn_prefill_backend: triton gpu_memory_utilization: 0.5 max_model_len: 16384 limit_mm_per_prompt: {image: 32} diff --git a/examples/puzzletron/configs/families/qwen3_5/setup_v2_defaults.yaml b/examples/puzzletron/configs/families/qwen3_5/setup_v2_defaults.yaml index b3a96810ec4..bd2c0f79df6 100644 --- a/examples/puzzletron/configs/families/qwen3_5/setup_v2_defaults.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/setup_v2_defaults.yaml @@ -87,6 +87,25 @@ model_overrides: bypass: enabled: false post_mip: + serving: + # Qwen3.5 GDN FlashInfer kernels may JIT for longer than the server + # readiness window on a fresh worker. The vLLM-supported Triton + # backend starts deterministically in the reviewed worker image. + topology: + # Reserve room for visual tokens when AIPerf derives the vLLM + # context limit for generated multimodal workloads. + server_context_overhead_tokens: 16384 + extra_vllm_args: + - -cc.cudagraph_mode=NONE + - --no-enable-flashinfer-autotune + - --gdn-prefill-backend + - triton + - --gpu-memory-utilization + - "0.5" + - --reasoning-parser + - qwen3 + - --default-chat-template-kwargs + - '{"enable_thinking": false}' quality_comparison: by_modality: text: @@ -156,6 +175,12 @@ model_overrides: instances: 1 mip: goal_value: 90% + balanced: + # With the maintained 0.8B search space, 75% is infeasible when the + # teacher embedding and depth are fixed. Keep the first-run profile at + # the validated 90% target used by the checked-in lifecycle recipes. + mip: + goal_value: 90% qwen3p5_9b: match: diff --git a/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_campaign.yaml b/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_campaign.yaml index 51ef7bc52ff..2cc5c0a7d75 100644 --- a/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_campaign.yaml +++ b/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_campaign.yaml @@ -13,7 +13,9 @@ execution: mip: {strategy: single, resource: cpu} depth_importance: {strategy: single} replacement_scoring: {strategy: single} - post.candidates.image_eval: &candidates {instances: 2} + # The documented default runner reserves one GPU at a time. One shard still + # evaluates every candidate sequentially and can resume between candidates. + post.candidates.image_eval: &candidates {instances: 1} post.candidates.best_image_loss: {strategy: single, resource: cpu} post.candidates.materialized: *candidates post.candidates.pre_kd_eval: *candidates diff --git a/examples/puzzletron/docs/environment_setup.md b/examples/puzzletron/docs/environment_setup.md index 567740d0a80..e8b566a2b8d 100644 --- a/examples/puzzletron/docs/environment_setup.md +++ b/examples/puzzletron/docs/environment_setup.md @@ -14,14 +14,22 @@ an optional Slurm container. ## Local Puzzletron environment The setup wizard and `orchestrate.py` do not import PyTorch or initialize CUDA. -Create one environment for both: +Use Python 3.10 through 3.14 to create one environment for both. Upgrade the +venv's bundled `pip` before resolving the controller dependencies: ```bash +python3 --version # must report Python 3.10 through 3.14 python3 -m venv .venv-puzzletron source .venv-puzzletron/bin/activate +python -m pip install --upgrade pip python -m pip install -r examples/puzzletron/requirements-setup.txt ``` +Setup resolves a Hugging Face model name to a commit before writing campaign +files. If that lookup is temporarily unavailable, setup can reuse a sole cached +model snapshot and still records its commit. With multiple cached snapshots it +keeps the network error instead of guessing which revision to use. + Only one local virtual environment is needed for a first campaign. `requirements-setup.txt` includes the packages required to generate, launch, and monitor a campaign. @@ -33,16 +41,18 @@ ModelOpt, CUDA, the worker container, or the worker virtual environment. ## Worker environment -The repository [`Dockerfile`](../Dockerfile) builds the worker image. It +Use your cluster's reviewed Puzzletron worker image when one is available. The +repository [`Dockerfile`](../Dockerfile) is the source for that image. It installs ModelOpt, the pinned vLLM and AutoModel sources, AIPerf, LMMS-Eval, the required CUDA extensions, and the teacher-evaluation resources. Its LMMS-Eval install includes the pinned native Qwen 3.5 image and video backend. Do not maintain a second set of worker installation commands or evaluator overlays outside the Dockerfile. -Build the Linux amd64 image from the repository root by following the -[image build and validation guide](worker_image.md). That guide provides the -build command and the revision-specific image tag. +Ask your cluster administrator for its registry reference or cluster-readable +path and enter it during setup. Build a Linux amd64 image only when your site +does not provide one; the [image build and validation guide](worker_image.md) +provides the build command and revision-specific image tag. The amd64 platform is required because the current CUDA extension set and Linux `eva-decord 0.6.1` dependency do not have a validated ARM build path. diff --git a/examples/puzzletron/docs/orchestration_operations.md b/examples/puzzletron/docs/orchestration_operations.md index 78ee8d331bc..98303f00375 100644 --- a/examples/puzzletron/docs/orchestration_operations.md +++ b/examples/puzzletron/docs/orchestration_operations.md @@ -47,9 +47,12 @@ online even when the surrounding campaign is configured for offline loading. ## Progress and interruption Interactive terminals show a live stage table with status, resources, elapsed -time, and a best-effort ETA when a stage reports progress. Completed stages, -dependency waits, failures, and descendants blocked by failures remain visible. -Redirected output uses timestamped one-line updates instead. +time, the active log path, and a best-effort ETA after the controller measures +item throughput. Completed stages, dependency waits, failures, and descendants +blocked by failures remain visible. Redirected output emits a heartbeat every +30 seconds with completed/total stages, queued/running jobs, elapsed time, each +active stage's state and progress, its log path, and a measured ETA. It says +`ETA unavailable` until it has enough progress evidence to estimate one. Press `q` or Ctrl-C in an interactive terminal to cancel active jobs and quit, detach while leaving jobs running, or continue. Non-interactive Ctrl-C and diff --git a/examples/puzzletron/docs/qwen3p5_0p8b_smoke.md b/examples/puzzletron/docs/qwen3p5_0p8b_smoke.md index d2dd8ea8204..cb27cbf5db4 100644 --- a/examples/puzzletron/docs/qwen3p5_0p8b_smoke.md +++ b/examples/puzzletron/docs/qwen3p5_0p8b_smoke.md @@ -16,6 +16,14 @@ conversion. These small budgets check that the complete workflow runs and resumes correctly. They do not establish model quality or production throughput. +## Before you start + +Prepare the controller venv and worker environment described in +[environment setup](environment_setup.md). The worker environment must provide +the [pinned evaluator installation](checkpoint_evaluation.md#quick-start). +IFEval task data must be fetchable from each worker or already present in its +Hugging Face cache. + ## Generate a complete bundle with the setup wizard For a new run, start with the [setup wizard](setup_wizard.md) and select Qwen @@ -74,9 +82,6 @@ The flow deliberately uses two candidate-evaluation samples, two IFEval samples, four AIPerf requests per serving candidate, and two distillation steps. These budgets validate workflow correctness, comparative serving selection, and resumability; they are not quality or throughput claims. -The worker environment must provide the [pinned evaluator -installation](checkpoint_evaluation.md#quick-start). IFEval task data must be -fetchable from each worker or already present in its Hugging Face cache. After completion, inspect the `checkpoint_eval` and `post_kd_checkpoint_eval` nodes under `artifacts/post_mip/nodes`. Their summaries must name the corresponding diff --git a/examples/puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md b/examples/puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md index f4dbc8c9c6b..a1db391e493 100644 --- a/examples/puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md +++ b/examples/puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md @@ -5,7 +5,7 @@ The Qwen 3.5 0.8B VLM example has two experiment files with distinct jobs: | Experiment | Purpose | Execution profile | | --- | --- | --- | | `full_vlm_smoke.yaml` | Check the complete lifecycle on one GPU | `execution.single_gpu.yaml` | -| `vlm_campaign.yaml` | Run an illustrative multi-axis campaign | `qwen3p5_0p8b/execution.vlm_campaign.yaml` | +| `vlm_campaign.yaml` | Run a scheduled, multi-hour multi-axis example | `qwen3p5_0p8b/execution.vlm_campaign.yaml` | Start with `full_vlm_smoke.yaml`. It uses small workloads to check dataset preparation, pruning, MIP, materialization, checkpoint evaluation, serving, @@ -13,8 +13,9 @@ two-step VLM distillation, final evaluation, and resume. Its scores and throughput are integration observations, not model-quality or production performance results. -The campaign is a larger illustrative experiment. It is not a recommended -pruning recipe, training duration, or candidate-selection policy. +The campaign is a larger illustrative experiment intended for scheduled +integration validation, not routine development or presubmit use. It is not a +recommended pruning recipe, training duration, or candidate-selection policy. Unit tests compile these recipes and verify their stage and resource contracts. They do not replace an end-to-end GPU run against the current model, data, @@ -32,30 +33,41 @@ Prepare the setup and worker environments described in - a shared campaign output directory. If workers cannot access the network, populate those caches before launch and -mount them through the runner. Do not replace the configured model identity -with a machine-specific path. +mount them through the runner. Keep the experiment's public model repository +and revision unchanged; a local cache is only where workers obtain those files. ## Run the lifecycle smoke -Set paths visible to every worker: +Set the shared cache, source identity, and run paths: ```bash -export PUZZLETRON_DATASET_PATH=/path/to/qwen3p5-vlm-smoke-data +export HF_HOME=/path/to/huggingface-cache +export PUZZLETRON_SOURCE_REVISION="$(git rev-parse HEAD)" export PUZZLETRON_DATASET_REVISION=51f4f4d219315c3283950994d4eb3d7fc30aa87b export PUZZLETRON_RUN_ROOT=/path/to/qwen3p5_0p8b_vlm_smoke +export PUZZLETRON_RUN_ROOT="$(python -c 'import os; print(os.path.realpath(os.environ["PUZZLETRON_RUN_ROOT"]))')" +DATASET_PATH="$PUZZLETRON_RUN_ROOT/datasets/nemotron_vlm_v2" ``` -Prepare the eight image-conversation samples. This command is safe to rerun -when the existing manifest matches the same request. - -```bash -python examples/puzzletron/materialize_dataset.py nemotron_vlm_v2 \ - --output "$PUZZLETRON_DATASET_PATH" \ - --revision "$PUZZLETRON_DATASET_REVISION" \ - --subsets sparsetables plotqa_cot wiki_en \ - --num-samples 8 \ - --max-shards-per-subset 1 -``` +`HF_HOME` is the shared destination for the evaluator cache and must be visible +to workers. With network access, the campaign's `prepare_dataset` stage fills +and validates that cache. For offline workers, populate it before launch using +the cache command in [VLM checkpoint evaluation](vlm_checkpoint_evaluation.md#cache-benchmark-data). +Keep `PUZZLETRON_SOURCE_REVISION` exported for both dry-run and launch; submitted +workers inherit it so controller and worker artifact identities stay identical. +The run-root normalization keeps provenance checks reproducible on systems +where a shared-storage alias traverses a symbolic link. Both canonical paths +must be visible to workers through the runner's mounts. + +The launch includes a `prepare_dataset` worker stage that creates and validates +the eight image-conversation samples at `DATASET_PATH`, inside the campaign +root. It also prepares the configured evaluation cache. You do not need dataset +or model dependencies in the lightweight controller venv, and there is no +separate data-preparation command for the online first-run path. + +Unlike the text-only example, this VLM route does not publish a separate +tokenized-dataset artifact. It keeps the image conversations in their native +format so the model processor can construct text and image inputs together. Use the maintained experiment with the shared single-GPU execution profile and a site-specific runner: @@ -87,9 +99,10 @@ After completion, inspect the campaign report and verify that: The smoke deliberately uses tiny workloads. Run a separate benchmark with representative requests before drawing performance conclusions. -## Inspect the illustrative campaign +## Run the longer multi-hour illustrative campaign -The campaign enables hidden width, heterogeneous FFN width, depth, +The campaign enables hidden width (the model's shared transformer hidden size), +heterogeneous FFN width (per-block feed-forward intermediate sizes), depth, grouped-attention geometry, and GDN geometry. All retained candidates use the same frozen evaluator rows, teacher checkpoint, and 128-step example KD budget. @@ -106,13 +119,17 @@ The 128-step value demonstrates the integration. It is not a convergence criterion or recommended training duration. The aggregate-rank rule is also an example policy rather than a general definition of the best model. -Compile the campaign before allocating resources: +After the lifecycle smoke succeeds, keep its controller venv and runner, choose +a new run root, and inspect the larger plan before allocating resources: ```bash EXPERIMENT=examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/vlm_campaign.yaml EXECUTION=examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_campaign.yaml RUNNER=/path/to/site-specific/runner.slurm.yaml +export HF_HOME=/path/to/huggingface-cache +export PUZZLETRON_SOURCE_REVISION="$(git rev-parse HEAD)" export PUZZLETRON_RUN_ROOT=/path/to/qwen3p5_0p8b_vlm_campaign +export PUZZLETRON_RUN_ROOT="$(python -c 'import os; print(os.path.realpath(os.environ["PUZZLETRON_RUN_ROOT"]))')" python examples/puzzletron/orchestrate.py \ --experiment "$EXPERIMENT" \ @@ -121,6 +138,14 @@ python examples/puzzletron/orchestrate.py \ --stage full --dry-run ``` +Remove `--dry-run` to launch the inspected plan. Rerun that identical launch +command to resume it. + +The campaign's `prepare_dataset` worker stage creates its 512-sample dataset +under this new run root. Do not point both recipes at the same run root: the +smoke manifest records an eight-sample request and is intentionally not +rewritten in place. + The campaign has no checked-in expected-result baseline. Evaluate its output against decision thresholds chosen for the target workload. @@ -130,7 +155,8 @@ Keep omitted architecture dimensions at their teacher values. When adding an axis, verify measurement and physical slicing on the target checkpoint before expanding the campaign. Use [MIP profiles](mip_profiles.md#search-space) for the search-space syntax and [configuration overrides](configuration_overrides.md) -for temporary changes. +for temporary changes. Change campaign inputs before launch; do not patch +resolved or generated artifacts inside an existing run root. Keep the model revision, frozen evaluator profile, KD exposure, and serving workload fixed when comparing candidates. Changing any of them creates a diff --git a/examples/puzzletron/docs/setup_wizard.md b/examples/puzzletron/docs/setup_wizard.md index 31e877e9cf9..79f41e5aee0 100644 --- a/examples/puzzletron/docs/setup_wizard.md +++ b/examples/puzzletron/docs/setup_wizard.md @@ -10,7 +10,8 @@ does not submit jobs. The guided flow offers three profiles: - **Quick smoke** creates the smallest campaign for checking campaign shape. -- **Balanced pruning** provides the recommended defaults for a first campaign. +- **Balanced pruning** provides the recommended defaults for a generated + campaign after the maintained lifecycle smoke succeeds. - **High-confidence search** spends more runtime on scoring and sanity checks. The selected profile supplies pruning and search defaults from the detected diff --git a/examples/puzzletron/docs/slurm_configuration.md b/examples/puzzletron/docs/slurm_configuration.md index 1c6c16af810..16efecbc010 100644 --- a/examples/puzzletron/docs/slurm_configuration.md +++ b/examples/puzzletron/docs/slurm_configuration.md @@ -117,3 +117,11 @@ not persisted. Inherit credentials from the launch environment, require an existing variable such as `${API_KEY:?set API_KEY}`, retrieve it from a secret command, or source a permission-protected `setup_env` file. This check catches common mistakes but is not a shell parser or a complete credential scanner. + +The setup defaults keep `TMPDIR` at the short worker-local `/tmp` path and put +the vLLM, FlashInfer, Triton, and PyTorch kernel caches in explicit writable +directories there. +Preserve those commands for containerized workers: vLLM uses Unix-domain +sockets with a platform path limit, and a read-only container home prevents +the runtime caches from being initialized. Another short, worker-local writable +directory is also valid. diff --git a/examples/puzzletron/evaluation/vlm/evaluator.py b/examples/puzzletron/evaluation/vlm/evaluator.py index 5de80282a15..f4ab33a97b6 100644 --- a/examples/puzzletron/evaluation/vlm/evaluator.py +++ b/examples/puzzletron/evaluation/vlm/evaluator.py @@ -83,6 +83,34 @@ def _completion_identity( } +def _progress_tasks(report: Mapping[str, object]) -> list[dict[str, object]]: + """Return exact task denominators when the selected profile provides them.""" + + source_tasks = report.get("source_tasks") + if not isinstance(source_tasks, list): + return [] + denominators = report.get("quick_task_denominators") + field = "selected_rows" + if not isinstance(denominators, Mapping): + denominators = report.get("profile_population_rows") + field = "" + tasks = [] + if isinstance(denominators, Mapping): + for task in source_tasks: + if not isinstance(task, str): + return [] + entry = denominators.get(task) + total = entry.get(field) if field and isinstance(entry, Mapping) else entry + if not isinstance(total, int) or isinstance(total, bool) or total <= 0: + return [] + tasks.append({"name": task, "total": total}) + return tasks + limit = report.get("sample_limit") + if len(source_tasks) == 1 and isinstance(limit, int) and limit > 0: + return [{"name": source_tasks[0], "total": limit}] + return [] + + def _file_identity(path: Path, *, root: Path) -> dict[str, object]: """Return a content identity for one checkpoint or evaluation artifact.""" digest = sha256() @@ -408,10 +436,11 @@ def evaluate( ) run_result = _load_completed_run(output_root, identity=identity) if run_result is None: + runtime_settings = {**settings, "progress_tasks": _progress_tasks(report)} run_result = checkpoint.run_lmms_eval_checkpoint( args.checkpoint, output_root=output_root, - settings=settings, + settings=runtime_settings, ) if "mmmu_val" in prepared.source_tasks: _attach_mmmu_parser_audit(run_result) diff --git a/examples/puzzletron/evaluation/vlm/post_mip.py b/examples/puzzletron/evaluation/vlm/post_mip.py index 6c6a698c4d7..d0a8cefc02e 100644 --- a/examples/puzzletron/evaluation/vlm/post_mip.py +++ b/examples/puzzletron/evaluation/vlm/post_mip.py @@ -50,6 +50,7 @@ _RUNNER_OVERRIDES = frozenset( { "dtype", + "gdn_prefill_backend", "gpu_memory_utilization", "limit_mm_per_prompt", "max_model_len", diff --git a/examples/puzzletron/evaluation/vlm/preparation/benchmark_data.py b/examples/puzzletron/evaluation/vlm/preparation/benchmark_data.py index 4c93678d3df..e39c7967b50 100644 --- a/examples/puzzletron/evaluation/vlm/preparation/benchmark_data.py +++ b/examples/puzzletron/evaluation/vlm/preparation/benchmark_data.py @@ -250,16 +250,24 @@ def _inventory_is_current( else: return False stat_result = inspected.stat() - if ( - stat_result.st_size != entry.get("bytes") - or stat_result.st_mtime_ns != entry.get("mtime_ns") - or stat_result.st_ctime_ns != entry.get("ctime_ns") + if stat_result.st_size != entry.get("bytes") or stat_result.st_ctime_ns != entry.get( + "ctime_ns" ): return False expected_sha256 = entry.get("sha256") - if not isinstance(expected_sha256, str) or ( - verify_content and _sha256(inspected) != expected_sha256 - ): + if not isinstance(expected_sha256, str): + return False + mtime_changed = stat_result.st_mtime_ns != entry.get("mtime_ns") + if mtime_changed: + if repository_cache is not None: + return False + # A distributed filesystem can expose a metadata-server mtime after + # preparation that differs from the client-observed value recorded + # immediately after fsync. Confirm just that file by content so a + # benign timestamp reconciliation does not invalidate a large cache. + if _sha256(inspected) != expected_sha256: + return False + elif verify_content and _sha256(inspected) != expected_sha256: return False observed_paths = sorted( path.relative_to(root).as_posix() @@ -963,7 +971,11 @@ def _prepare( with _task_lock(hf_home, task): target, complete = _inspect_prepare_target(hf_home, task, verify_content=verify_content) if complete is not None: - return complete + return { + **complete, + "snapshot": str(snapshot), + "media_root": str(target), + } staging = Path( tempfile.mkdtemp(prefix=f".{target.name}.modelopt-staging.", dir=target.parent) ) diff --git a/modelopt/torch/puzzletron/evaluation/lmms.py b/modelopt/torch/puzzletron/evaluation/lmms.py index 30d94bf8a53..d021f60e05a 100644 --- a/modelopt/torch/puzzletron/evaluation/lmms.py +++ b/modelopt/torch/puzzletron/evaluation/lmms.py @@ -22,10 +22,12 @@ import json import math import os +import re import shlex import signal import sys import tempfile +import time import uuid from dataclasses import dataclass from pathlib import Path @@ -42,6 +44,7 @@ _MODEL_ARG_FIELDS = frozenset( { "dtype", + "gdn_prefill_backend", "gpu_memory_utilization", "attention_config", "chat_template", @@ -90,6 +93,18 @@ _PROCESS_CLEANUP_TIMEOUT_SECONDS = 10.0 _PROCESS_GROUP_POLL_INTERVAL_SECONDS = 0.1 _TIMEOUT_ERRORS = (TimeoutError, asyncio.TimeoutError) +_PROGRESS_PATH_ENV = "PUZZLETRON_EVALUATION_PROGRESS_PATH" +_PROGRESS_TASKS_ENV = "PUZZLETRON_EVALUATION_PROGRESS_TASKS" +_MODEL_RESPONDING = re.compile( + r"Model Responding:.*?(?P\d+)/(?P\d+)\s*" + r"\[(?P[^<\],]+)<(?P[^,\]]+),\s*" + r"(?P\d+(?:\.\d+)?)(?Pit/s|s/it)\]" +) +_MODEL_RESPONDING_WITHOUT_TOTAL = re.compile( + r"Model Responding:\s*(?P\d+)it\s*" + r"\[(?P[^,\]]+),\s*(?P\d+(?:\.\d+)?)" + r"(?Pit/s|s/it)\]" +) _COMPATIBILITY_TASKS: dict[str, dict[str, Any]] = { "gsm8k": { "alias": "modelopt_gsm8k", @@ -623,6 +638,166 @@ def _stream_text(value: str | bytes | None) -> str: return value.decode(errors="replace") if isinstance(value, bytes) else value or "" +def _live_stream_bytes(stream: Any, data: bytes) -> None: + try: + target = getattr(stream, "buffer", None) + if target is not None: + target.write(data) + target.flush() + return + stream.write(data.decode(errors="replace")) + stream.flush() + except (OSError, ValueError): + pass + + +def _sample_rate_per_second(value: str, unit: str) -> float | None: + rate = float(value) + if rate <= 0: + return None + return rate if unit == "it/s" else 1.0 / rate + + +def _task_progress( + current: int, + total: int | None, + tasks: object, +) -> dict[str, object] | None: + if total is None or not isinstance(tasks, list): + return None + parsed: list[tuple[str, int]] = [] + for entry in tasks: + if not isinstance(entry, Mapping): + return None + name = entry.get("name") + task_total = entry.get("total") + if not isinstance(name, str) or not isinstance(task_total, int) or task_total <= 0: + return None + parsed.append((name, task_total)) + if not parsed or sum(task_total for _name, task_total in parsed) != total: + return None + offset = current + for name, task_total in parsed: + if offset <= task_total: + return {"name": name, "current": offset, "total": task_total} + offset -= task_total + name, task_total = parsed[-1] + return {"name": name, "current": task_total, "total": task_total} + + +def _set_evaluation_progress_status( + path: Path, + status: str, + *, + sample_counts: Mapping[str, float] | None = None, + tasks: object = None, +) -> None: + """Best-effort terminal update for the evaluator progress sidecar.""" + + try: + payload = json.loads(path.read_text()) + except (OSError, ValueError): + return + if not isinstance(payload, Mapping): + return + updated = dict(payload) + updated["status"] = status + updated["updated_at"] = time.time() + if status == "completed" and sample_counts: + sample_total = float(sum(sample_counts.values())) + if math.isfinite(sample_total) and sample_total > 0 and sample_total.is_integer(): + total = int(sample_total) + updated["current"] = total + updated["total"] = total + task = _task_progress(total, total, tasks) + if task is None: + updated.pop("task", None) + else: + updated["task"] = task + try: + _atomic_json(path, updated) + except OSError: + pass + + +def _evaluation_progress_payload(text: str, tasks: object) -> dict[str, object] | None: + matches = list(_MODEL_RESPONDING.finditer(text)) + total: int | None + if matches: + match = matches[-1] + total = int(match.group("total")) + else: + fallback = list(_MODEL_RESPONDING_WITHOUT_TOTAL.finditer(text)) + if not fallback: + return None + match = fallback[-1] + total = None + current = int(match.group("current")) + payload: dict[str, object] = { + "schema": "modelopt.puzzletron.evaluation-progress/v1", + "status": "running", + "unit": "samples", + "current": current, + "total": total, + "rate_per_second": _sample_rate_per_second(match.group("rate"), match.group("rate_unit")), + "updated_at": time.time(), + } + task = _task_progress(current, total, tasks) + if task is not None: + payload["task"] = task + return payload + + +async def _pump_process_stream( + reader: asyncio.StreamReader, + capture: Any, + live_stream: Any, + *, + progress_path: Path | None = None, + progress_tasks: object = None, +) -> None: + progress_text = "" + while data := await reader.read(65536): + capture.write(data) + capture.flush() + _live_stream_bytes(live_stream, data) + if progress_path is None: + continue + progress_text = (progress_text + data.decode(errors="replace"))[-131072:] + payload = _evaluation_progress_payload(progress_text, progress_tasks) + if payload is not None: + try: + _atomic_json(progress_path, payload) + except OSError: + pass + + +async def _drain_process_pumps( + process: asyncio.subprocess.Process, + pumps: tuple[asyncio.Task[None], ...], + *, + suppress_errors: bool, +) -> None: + done, pending = await asyncio.wait(pumps, timeout=_PROCESS_CLEANUP_TIMEOUT_SECONDS) + for pump in pending: + pump.cancel() + if pending: + # asyncio Process has no public stream-close API. Closing its transport + # releases inherited pipe readers after a descendant outlives the parent. + process_transport = getattr(process, "_transport", None) + if process_transport is not None: + try: + process_transport.close() + except (OSError, RuntimeError): + pass + results = await asyncio.gather(*pumps, return_exceptions=True) + if suppress_errors: + return + for pump, result in zip(pumps, results): + if pump in done and isinstance(result, BaseException): + raise result + + def _output_tail(result: _ProcessResult, *, max_lines: int = 20) -> str: sections = [] for stream_name, text in (("stderr", result.stderr), ("stdout", result.stdout)): @@ -666,6 +841,18 @@ async def _wait_for_process_group_exit( await asyncio.sleep(min(_PROCESS_GROUP_POLL_INTERVAL_SECONDS, remaining)) +async def _wait_for_process_returncode( + process: asyncio.subprocess.Process, *, timeout: float +) -> None: + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while process.returncode is None: + remaining = deadline - loop.time() + if remaining <= 0: + raise TimeoutError + await asyncio.sleep(min(_PROCESS_GROUP_POLL_INTERVAL_SECONDS, remaining)) + + async def _run_process_async( argv: list[str], *, @@ -675,17 +862,39 @@ async def _run_process_async( ) -> _ProcessResult: # lmms-eval needs process isolation for bounded GPU-worker cleanup. The argument # vector is passed directly; no shell interprets checkpoint or configuration values. + child_env = dict(env) + raw_progress_path = child_env.pop(_PROGRESS_PATH_ENV, None) + raw_progress_tasks = child_env.pop(_PROGRESS_TASKS_ENV, None) + progress_path = Path(raw_progress_path) if raw_progress_path else None + try: + progress_tasks = json.loads(raw_progress_tasks) if raw_progress_tasks else None + except json.JSONDecodeError: + progress_tasks = None with tempfile.TemporaryFile() as stdout_file, tempfile.TemporaryFile() as stderr_file: process = await asyncio.create_subprocess_exec( *argv, cwd=cwd, - env=env, - stdout=stdout_file, - stderr=stderr_file, + env=child_env, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, start_new_session=os.name == "posix", ) + assert process.stdout is not None + assert process.stderr is not None + pumps = ( + asyncio.create_task(_pump_process_stream(process.stdout, stdout_file, sys.stdout)), + asyncio.create_task( + _pump_process_stream( + process.stderr, + stderr_file, + sys.stderr, + progress_path=progress_path, + progress_tasks=progress_tasks, + ) + ), + ) try: - await asyncio.wait_for(process.wait(), timeout) + await _wait_for_process_returncode(process, timeout=timeout) except _TIMEOUT_ERRORS as error: _signal_process_group(process, signal.SIGTERM) try: @@ -702,6 +911,7 @@ async def _run_process_async( process, deadline=asyncio.get_running_loop().time() + _PROCESS_CLEANUP_TIMEOUT_SECONDS, ) + await _drain_process_pumps(process, pumps, suppress_errors=True) stdout_file.seek(0) stderr_file.seek(0) raise LmmsEvalTimeoutError( @@ -710,6 +920,7 @@ async def _run_process_async( output=_stream_text(stdout_file.read()), stderr=_stream_text(stderr_file.read()), ) from error + await _drain_process_pumps(process, pumps, suppress_errors=False) stdout_file.seek(0) stderr_file.seek(0) return _ProcessResult( @@ -777,6 +988,38 @@ def run_lmms_eval_checkpoint( "timeout": timeout, } command_path = _atomic_json(output / "command.json", command_payload) + progress_tasks = settings.get("progress_tasks") + if not isinstance(progress_tasks, list) or not all( + isinstance(task, Mapping) + and isinstance(task.get("name"), str) + and isinstance(task.get("total"), int) + and not isinstance(task.get("total"), bool) + and int(task["total"]) > 0 + for task in progress_tasks + ): + progress_tasks = [] + try: + progress_tasks_json = json.dumps(progress_tasks, separators=(",", ":")) + except (TypeError, ValueError): + progress_tasks = [] + progress_tasks_json = "[]" + progress_path = output / "progress.json" + progress_total = sum(int(task["total"]) for task in progress_tasks) + initial_progress: dict[str, object] = { + "schema": "modelopt.puzzletron.evaluation-progress/v1", + "status": "starting", + "unit": "samples", + "current": 0, + "total": progress_total or None, + "rate_per_second": None, + "updated_at": time.time(), + } + initial_task = _task_progress(0, progress_total or None, progress_tasks) + if initial_task is not None: + initial_progress["task"] = initial_task + _atomic_json(progress_path, initial_progress) + env[_PROGRESS_PATH_ENV] = str(progress_path) + env[_PROGRESS_TASKS_ENV] = progress_tasks_json try: result = _run_process(argv, cwd=str(output), env=env, timeout=timeout) except LmmsEvalTimeoutError as error: @@ -785,6 +1028,10 @@ def run_lmms_eval_checkpoint( captured = _ProcessResult(argv, -1, error.output, error.stderr) stream_paths = _write_streams(output, captured) _annotate_error(error, command_path=command_path, stream_paths=stream_paths) + _set_evaluation_progress_status(progress_path, "timed_out") + raise + except OSError: + _set_evaluation_progress_status(progress_path, "failed") raise stream_paths = _write_streams(output, result) @@ -794,6 +1041,7 @@ def run_lmms_eval_checkpoint( f"lmms-eval failed with exit code {result.returncode}" + (f": {tail}" if tail else "") ) _annotate_error(failure, command_path=command_path, stream_paths=stream_paths) + _set_evaluation_progress_status(progress_path, "failed") raise failure try: @@ -806,9 +1054,11 @@ def run_lmms_eval_checkpoint( tail = _output_tail(result) failure = FileNotFoundError(f"{error}: {tail}" if tail else str(error)) _annotate_error(failure, command_path=command_path, stream_paths=stream_paths) + _set_evaluation_progress_status(progress_path, "failed") raise failure from error except RuntimeError as error: _annotate_error(error, command_path=command_path, stream_paths=stream_paths) + _set_evaluation_progress_status(progress_path, "failed") raise summary = { @@ -817,7 +1067,17 @@ def run_lmms_eval_checkpoint( "raw_result_path": str(result_path), "sample_counts": sample_counts, } - summary_path = _atomic_json(output / "summary.json", summary) + try: + summary_path = _atomic_json(output / "summary.json", summary) + except OSError: + _set_evaluation_progress_status(progress_path, "failed") + raise + _set_evaluation_progress_status( + progress_path, + "completed", + sample_counts=sample_counts, + tasks=progress_tasks, + ) return { "metrics": metrics, "result_path": str(summary_path), diff --git a/modelopt/torch/puzzletron/orchestration/controller.py b/modelopt/torch/puzzletron/orchestration/controller.py index f87d2477121..91ac1796be9 100644 --- a/modelopt/torch/puzzletron/orchestration/controller.py +++ b/modelopt/torch/puzzletron/orchestration/controller.py @@ -33,7 +33,7 @@ from .adapters.registry import adapter_for_stage from .adapters.stage_compat import stage_is_complete from .compiler import _resolve_artifact_settling_timeout_seconds, plan_to_dict -from .dashboard import StageView, format_duration, progress_eta, progress_fraction +from .dashboard import StageView, format_duration, format_eta, progress_fraction from .executors import BareMetalSSHExecutor, Executor, LocalExecutor, SlurmExecutor from .executors.slurm import render_slurm_attempt_script from .identity import stable_hash @@ -237,6 +237,7 @@ def __init__( self._active: dict[str, tuple[JobHandle, str, str]] = {} self._last_states: dict[str, JobState] = {} self._last_heartbeat = 0.0 + self._progress_samples: dict[str, tuple[float, float, float]] = {} self._campaign_started_monotonic = time.monotonic() self._shutdown_requested = False self._shutdown_signal: int | None = None @@ -1018,17 +1019,24 @@ def _emit_progress_heartbeat(self) -> None: running = sum(state is JobState.RUNNING for state in self._last_states.values()) self.logger.wait( f"progress {completed}/{len(self.plan.stages)} stages; " - f"jobs: {running} running, {pending} pending" + f"jobs: {running} running, {pending} pending; " + f"elapsed={format_duration(time.monotonic() - self._campaign_started_monotonic)}" ) for view in self._stage_views(): if self._shutdown_requested: return if view.status not in {"running", "pending"}: continue + log_path = "unavailable" + if view.log_paths: + log_path = view.log_paths[0] + if len(view.log_paths) > 1: + log_path += f" (+{len(view.log_paths) - 1} more)" self.logger.progress( f"{view.stage_id} {view.nodes}n/{view.tasks}t/{view.gpus}g " - f"{view.progress}; elapsed={format_duration(view.elapsed_seconds)}, " - f"eta={format_duration(view.eta_seconds, approximate=True)}" + f"state={view.status}; {view.progress}; " + f"elapsed={format_duration(view.elapsed_seconds)}, " + f"eta={format_eta(view.eta_seconds)}; log={log_path}" ) def _failed_ancestor_ids(self, stage_id: str) -> list[str]: @@ -1053,8 +1061,20 @@ def visit(current: str) -> None: visit(stage_id) return [node.stage_id for node in self.plan.stages if node.stage_id in failed] - def _stage_elapsed(self, stage_id: str, *, active: bool) -> float | None: + def _stage_elapsed( + self, + stage_id: str, + *, + active: bool, + active_attempt_ids: set[str] | None = None, + ) -> float | None: attempts = self.store.list_attempts(stage_id) + if active and active_attempt_ids: + attempts = [ + attempt + for attempt in attempts + if str(attempt.get("attempt_id")) in active_attempt_ids + ] starts = [ float(attempt["submitted_at"]) for attempt in attempts @@ -1072,6 +1092,29 @@ def _stage_elapsed(self, stage_id: str, *, active: bool) -> float | None: ] return max(0.0, max(ends) - start) if ends else None + def _measured_progress_eta( + self, + stage_id: str, + current: float | None, + total: float | None, + ) -> float | None: + """Estimate remaining time only after observing progress in this controller run.""" + + if current is None or total is None or total <= 0 or current >= total: + self._progress_samples.pop(stage_id, None) + return None + now = time.monotonic() + sample = self._progress_samples.get(stage_id) + if sample is None or sample[2] != total or current < sample[1]: + self._progress_samples[stage_id] = (now, current, total) + return None + started_at, started_current, _sample_total = sample + completed = current - started_current + elapsed = now - started_at + if completed <= 0 or elapsed <= 0: + return None + return elapsed * (total - current) / completed + def _stage_views(self) -> list[StageView]: """Build dashboard rows from durable state, DAG state, and progress artifacts.""" @@ -1097,7 +1140,12 @@ def _stage_views(self) -> list[StageView]: self._last_states.get(handle_id, JobState.UNKNOWN) for handle_id, _entry in stage_entries ] - elapsed = self._stage_elapsed(node.stage_id, active=stage_active) + active_attempt_ids = {entry[2] for _handle_id, entry in stage_entries} + elapsed = self._stage_elapsed( + node.stage_id, + active=stage_active, + active_attempt_ids=active_attempt_ids, + ) progress = "" current = total = None if completed: @@ -1162,9 +1210,10 @@ def _stage_views(self) -> list[StageView]: gpus=node.total_gpus, progress=progress, elapsed_seconds=elapsed, - eta_seconds=progress_eta(elapsed, current, total), + eta_seconds=self._measured_progress_eta(node.stage_id, current, total), current=current, total=total, + log_paths=tuple(stage_logs) if stage_active else (), ) ) return views diff --git a/modelopt/torch/puzzletron/orchestration/dashboard.py b/modelopt/torch/puzzletron/orchestration/dashboard.py index ba7932c96f3..83208691e9b 100644 --- a/modelopt/torch/puzzletron/orchestration/dashboard.py +++ b/modelopt/torch/puzzletron/orchestration/dashboard.py @@ -28,6 +28,7 @@ "StageView", "TerminalDashboard", "format_duration", + "format_eta", "progress_fraction", "progress_eta", ] @@ -51,6 +52,7 @@ class StageView: eta_seconds: float | None = None current: float | None = None total: float | None = None + log_paths: tuple[str, ...] = () def progress_fraction(detail: str | None) -> tuple[float, float] | None: @@ -99,6 +101,14 @@ def format_duration(seconds: float | None, *, approximate: bool = False) -> str: return f"~{value}" if approximate else value +def format_eta(seconds: float | None) -> str: + """Format a measured ETA without implying one when no estimate is available.""" + + if seconds is None: + return "unavailable" + return format_duration(seconds, approximate=True) + + def _progress_bar(current: float | None, total: float | None, *, width: int = 14) -> str: if current is None or total is None or total <= 0: return "" @@ -205,6 +215,9 @@ def _render( bar = _progress_bar(stage.current, stage.total) if bar: progress.append(f" {bar}", style="blue") + if stage.log_paths: + suffix = f" (+{len(stage.log_paths) - 1} more)" if len(stage.log_paths) > 1 else "" + progress.append(f"\nlog: {stage.log_paths[0]}{suffix}", style="dim") allocation = f"{stage.nodes}n · {stage.tasks}t · {stage.gpus}g" table.add_row( status_renderable, @@ -212,7 +225,7 @@ def _render( allocation, progress, format_duration(stage.elapsed_seconds), - format_duration(stage.eta_seconds, approximate=True), + format_eta(stage.eta_seconds), ) footer = Text("Last event", style="dim") footer.append(f" {self._last_event}") diff --git a/modelopt/torch/puzzletron/orchestration/progress.py b/modelopt/torch/puzzletron/orchestration/progress.py index 338589a46ef..61725d02e16 100644 --- a/modelopt/torch/puzzletron/orchestration/progress.py +++ b/modelopt/torch/puzzletron/orchestration/progress.py @@ -422,6 +422,34 @@ def _post_mip_progress( return None executions_root = puzzle_dir / "artifacts" / "post_mip" / "nodes" / node_id / "executions" + if node_type in {"evaluation", "downstream_evaluation"}: + progress_paths = list(executions_root.rglob("lmms_eval/attempt_*/progress.json")) + if progress_paths: + progress_path = max(progress_paths, key=lambda path: path.stat().st_mtime_ns) + payload = _read_json(progress_path) + if ( + isinstance(payload, Mapping) + and payload.get("schema") == "modelopt.puzzletron.evaluation-progress/v1" + and payload.get("unit") == "samples" + ): + current = payload.get("current") + total = payload.get("total") + rate = payload.get("rate_per_second") + task = payload.get("task") + prefix = "evaluation" + if isinstance(task, Mapping) and isinstance(task.get("name"), str): + prefix += f" {task['name']}" + current = task.get("current", current) + total = task.get("total", total) + if isinstance(current, int) and isinstance(total, int) and total > 0: + detail = f"{prefix} {current}/{total} samples" + elif isinstance(current, int): + detail = f"{prefix} {current} samples (total unavailable)" + else: + detail = f"{prefix} starting (total unavailable)" + if isinstance(rate, (int, float)) and rate > 0: + detail += f" at {rate:.2f} samples/s" + return detail executions = [path for path in executions_root.glob("post_mip_execution_*") if path.is_dir()] rows_by_revision: dict[str, Mapping[str, Any]] = {} if executions: diff --git a/modelopt/torch/puzzletron/orchestration/reporting.py b/modelopt/torch/puzzletron/orchestration/reporting.py index 791fa7e79b3..1316fc5b769 100644 --- a/modelopt/torch/puzzletron/orchestration/reporting.py +++ b/modelopt/torch/puzzletron/orchestration/reporting.py @@ -97,6 +97,17 @@ def _completion_path(plan: CampaignPlan) -> Path: return report_path.parent / "completion.json" +def _stage_state_sha256(plan: CampaignPlan) -> str: + """Fingerprint durable stage records consumed by the progress report.""" + + digest = hashlib.sha256() + stage_root = plan.puzzle_dir / "orchestration" / "stages" + for path in sorted(stage_root.glob("*.json")): + digest.update(path.name.encode()) + digest.update(_sha256(path).encode()) + return digest.hexdigest() + + def completed_final_report(plan: CampaignPlan) -> FinalReportResult | None: """Return a sealed final report when its contract and artifact hashes still match.""" @@ -109,8 +120,9 @@ def completed_final_report(plan: CampaignPlan) -> FinalReportResult | None: payload = json.loads(record_bytes) log_paths = payload["log_paths"] if ( - payload["schema_version"] != 1 + payload["schema_version"] != 2 or payload["contract_hash"] != plan.contract_hash + or payload["stage_state_sha256"] != _stage_state_sha256(plan) or payload["report_sha256"] != _sha256(report_path) or payload["manifest_sha256"] != _sha256(manifest_path) or not isinstance(log_paths, list) @@ -134,8 +146,9 @@ def record_completed_final_report( report_path, manifest_path = final_report_paths(plan) payload = { - "schema_version": 1, + "schema_version": 2, "contract_hash": plan.contract_hash, + "stage_state_sha256": _stage_state_sha256(plan), "report_sha256": _sha256(report_path), "manifest_sha256": _sha256(manifest_path), "log_paths": list(log_paths), diff --git a/modelopt/torch/puzzletron/plugins/automodel/solution_launch.py b/modelopt/torch/puzzletron/plugins/automodel/solution_launch.py index 831d0767f0e..e10ce6e64e8 100644 --- a/modelopt/torch/puzzletron/plugins/automodel/solution_launch.py +++ b/modelopt/torch/puzzletron/plugins/automodel/solution_launch.py @@ -47,7 +47,7 @@ from ...replacement_library.replacement_utils import parse_layer_replacement from ...tools.logger import mprint from ...tools.validate_puzzle_with_multi_replacements import load_puzzle_solutions -from ...tools.validation_utils import write_results +from ...tools.validation_utils import scoring_result_matches, write_results from .config import build_solution_recipe_config, solution_scoring_params from .launch import _free_scoring_memory from .module_trace import synchronized_module_trace @@ -81,11 +81,7 @@ def _can_skip_parent_model_load( needs_parent_evaluation: bool, ) -> bool: """Skip completed candidates, but always rebuild the run-local teacher cache.""" - return ( - role != "original" - and not pending_ids - and not needs_parent_evaluation - ) + return role != "original" and not pending_ids and not needs_parent_evaluation def _solution_output_location(scoring, output_dir: Path, solution_id: int) -> tuple[Path, str]: @@ -100,9 +96,7 @@ def _solution_output_location(scoring, output_dir: Path, solution_id: int) -> tu def _solution_result_path(scoring, output_dir: Path, solution_id: int) -> Path: - solution_output, solution_name = _solution_output_location( - scoring, output_dir, solution_id - ) + solution_output, solution_name = _solution_output_location(scoring, output_dir, solution_id) return solution_output / f"{solution_name}.json" @@ -132,7 +126,11 @@ def _load_solution_work(scoring, output_dir: Path) -> tuple[list[dict], list[int ids = [ i for i in ids - if not _solution_result_path(scoring, output_dir, i).exists() + if not scoring_result_matches( + _solution_result_path(scoring, output_dir, i), + scoring, + expected_payload={"puzzle_solution": solutions[i]}, + ) ] return solutions, ids @@ -204,7 +202,9 @@ def _run_recipe( return recipe -def _extract_teacher_targets(recipe, cache: TeacherTargetCache, params: dict | None = None) -> dict | None: +def _extract_teacher_targets( + recipe, cache: TeacherTargetCache, params: dict | None = None +) -> dict | None: """Phase 1: fill the cache and optionally compute the teacher baseline metrics.""" if recipe.has_outputs: cache.set_lm_head_weight(recipe.lm_head_weight()) @@ -345,7 +345,9 @@ def _score_candidate( synchronized_module_trace(recipe), ): for batch_idx, (hidden, targets) in enumerate(recipe.iterate_captures()): - _trace_batch("candidate_capture_yield", batch_idx, has_hidden=hidden is not None, name=name) + _trace_batch( + "candidate_capture_yield", batch_idx, has_hidden=hidden is not None, name=name + ) if hidden is not None: torch.cuda.synchronize(hidden.device) forward_seconds = time.perf_counter() - batch_started @@ -379,8 +381,13 @@ def _score_candidate( metric_started = time.perf_counter() per_batch.append( score_batch( - candidate_hidden, candidate_w, teacher_hidden, teacher_w, targets, - temperature=params["temperature"], chunk_size=params["chunk_size"], + candidate_hidden, + candidate_w, + teacher_hidden, + teacher_w, + targets, + temperature=params["temperature"], + chunk_size=params["chunk_size"], lm_head_backend=params["lm_head_backend"], tp_group=tp_group, flash_kld_token_chunk_size=params["flash_kld_token_chunk_size"], @@ -461,12 +468,8 @@ def _validate_parent_equivalence( "normalized_mse_loss_hidden_states": float( tolerances.get("max_normalized_mse_loss_hidden_states", 2.0e-2) ), - "mse_loss_hidden_states": float( - tolerances.get("max_mse_loss_hidden_states", 1.0e-1) - ), - "mae_loss_hidden_states": float( - tolerances.get("max_mae_loss_hidden_states", 5.0e-1) - ), + "mse_loss_hidden_states": float(tolerances.get("max_mse_loss_hidden_states", 1.0e-1)), + "mae_loss_hidden_states": float(tolerances.get("max_mae_loss_hidden_states", 5.0e-1)), } for metric, limit in metric_limits.items(): value = _metric_average(parent, metric) @@ -544,7 +547,9 @@ def _solution_hidden_width(solution: dict) -> int | None: return None if value is None else int(value) -def _solution_prune_target(layer_replacements, teacher_block_configs, num_q_heads, head_dim) -> dict | None: +def _solution_prune_target( + layer_replacements, teacher_block_configs, num_q_heads, head_dim +) -> dict | None: """Resolve a single-block replacement into prune_block_context kwargs (orig + target dims). Attention targets are interpreted as sorted-prefix removal: reducing KV groups also removes @@ -568,25 +573,11 @@ def _solution_prune_target(layer_replacements, teacher_block_configs, num_q_head teacher_ffn = teacher.get_subblock("ffn") teacher_attn = teacher.get_subblock("attention") - t_ffn = ( - child_ffn.intermediate_size - if child_ffn is not None and not child_ffn.no_op - else None - ) - t_kv = ( - child_attn.num_kv_heads - if child_attn is not None and not child_attn.no_op - else None - ) - t_q = ( - child_attn.num_query_heads - if child_attn is not None and not child_attn.no_op - else None - ) + t_ffn = child_ffn.intermediate_size if child_ffn is not None and not child_ffn.no_op else None + t_kv = child_attn.num_kv_heads if child_attn is not None and not child_attn.no_op else None + t_q = child_attn.num_query_heads if child_attn is not None and not child_attn.no_op else None orig_kv = ( - teacher_attn.num_kv_heads - if teacher_attn is not None and not teacher_attn.no_op - else None + teacher_attn.num_kv_heads if teacher_attn is not None and not teacher_attn.no_op else None ) target_num_q = None if t_kv is not None: @@ -731,9 +722,7 @@ def launch_score_solutions_automodel(hydra_cfg, num_nodes: int = 1, node_index: if dist.is_master(): mprint(f"[solution/automodel] building sorted teacher -> {default_sorted_dir}") sort_cfg = hydra_cfg.get("sort", {}) - embedding_widths = tuple( - hydra_cfg.get("embedding_pruning", {}).get("widths", ()) or () - ) + embedding_widths = tuple(hydra_cfg.get("embedding_pruning", {}).get("widths", ()) or ()) build_sorted_teacher( teacher_dir, activations_log_dir, @@ -807,10 +796,8 @@ def launch_score_solutions_automodel(hydra_cfg, num_nodes: int = 1, node_index: return # ---- Cache teacher targets, then score candidates from the requested source. ---- - mprint( - "[solution/automodel] checkpoint roles | " - f"target={target_dir} source={source_dir}" - ) + mprint(f"[solution/automodel] checkpoint roles | target={target_dir} source={source_dir}") + def score_pending(recipe, cache) -> None: sliced_teacher_baseline = None if baseline_only or bool(scoring.get("score_source_baseline", True)): @@ -847,7 +834,11 @@ def score_pending(recipe, cache) -> None: f"as {solution_output / solution_name} {prune_target}" ) _score_candidate( - recipe, cache, params, solution_output, scoring, + recipe, + cache, + params, + solution_output, + scoring, name=solution_name, payload={ "i_solution": 0 if solution_name == "solution_0" else i_solution, @@ -864,7 +855,10 @@ def score_pending(recipe, cache) -> None: recipe = _run_recipe( build_solution_recipe_config(hydra_cfg, target_dir), - scoring, params["eval_iters"], params["use_puzzletron_dataloader"], params["data_cfg"], + scoring, + params["eval_iters"], + params["use_puzzletron_dataloader"], + params["data_cfg"], ) try: mprint("[solution/automodel] Phase 1: caching teacher targets") @@ -886,7 +880,10 @@ def score_pending(recipe, cache) -> None: recipe = _run_recipe( build_solution_recipe_config(hydra_cfg, source_dir), - scoring, params["eval_iters"], params["use_puzzletron_dataloader"], params["data_cfg"], + scoring, + params["eval_iters"], + params["use_puzzletron_dataloader"], + params["data_cfg"], ) try: score_pending(recipe, cache) @@ -971,13 +968,32 @@ def write_manifest() -> None: pending_ids = [ idx for idx in range(len(solutions)) - if force_rescore or not (output_dir / f"solution_{idx}.json").is_file() + if force_rescore + or not scoring_result_matches( + output_dir / f"solution_{idx}.json", + scoring, + expected_payload={ + "puzzle_solution": solutions[idx], + "parent_role": role, + "checkpoint_dir": str(checkpoint_dir), + }, + ) ] parent_result_path = output_dir / "parent.json" needs_parent_evaluation = ( evaluation_mode != "realized_baseline" and not skip_parent_equivalence - and (force_rescore or not parent_result_path.is_file()) + and ( + force_rescore + or not scoring_result_matches( + parent_result_path, + scoring, + expected_payload={ + "parent_role": role, + "checkpoint_dir": str(checkpoint_dir), + }, + ) + ) ) if _can_skip_parent_model_load( role, @@ -1000,17 +1016,13 @@ def write_manifest() -> None: ) try: if not (checkpoint_dir / "config.json").is_file(): - raise FileNotFoundError( - f"{role} parent missing config.json: {checkpoint_dir}" - ) + raise FileNotFoundError(f"{role} parent missing config.json: {checkpoint_dir}") parent_config = _load_model_config_distributed( checkpoint_dir, descriptor, loader=load_model_config, ) - parent_width = int( - descriptor.get_language_model_config(parent_config).hidden_size - ) + parent_width = int(descriptor.get_language_model_config(parent_config).hidden_size) retained_hidden_indices = _source_hidden_channel_indices( checkpoint_dir, parent_width, @@ -1035,7 +1047,9 @@ def write_manifest() -> None: raise manifest["checkpoint_loads"][role] += 1 if manifest["checkpoint_loads"][role] != 1: - raise RuntimeError(f"parent {role} loaded more than once: {manifest['checkpoint_loads']}") + raise RuntimeError( + f"parent {role} loaded more than once: {manifest['checkpoint_loads']}" + ) write_manifest() parent_summary = None @@ -1065,7 +1079,9 @@ def write_manifest() -> None: parent_summary = {"passed": True, "reference": True} elif evaluation_mode == "runtime_slice" and not skip_parent_equivalence: if original_result_path is None or len(cache) == 0 and recipe.has_outputs: - raise RuntimeError("original teacher cache is unavailable for sorted parent") + raise RuntimeError( + "original teacher cache is unavailable for sorted parent" + ) if needs_parent_evaluation: mprint( "[solution/automodel] parent sweep equivalence | " @@ -1090,9 +1106,7 @@ def write_manifest() -> None: teacher_result_path=original_result_path, parent_result_path=parent_result_path, tolerances=tolerances, - hidden_basis_permuted=bool( - parent.get("hidden_basis_permuted", False) - ), + hidden_basis_permuted=bool(parent.get("hidden_basis_permuted", False)), ) mprint( "[solution/automodel] parent equivalence | " diff --git a/modelopt/torch/puzzletron/post_mip/reporting.py b/modelopt/torch/puzzletron/post_mip/reporting.py index 04d0edd29bb..aa32dca1dbb 100644 --- a/modelopt/torch/puzzletron/post_mip/reporting.py +++ b/modelopt/torch/puzzletron/post_mip/reporting.py @@ -119,7 +119,7 @@ def _candidate_label( width = width if width is not None else artifact.get("hidden_width") kind = kind or artifact.get("kind") current_id = str(revision.get("parent_revision_id") or "") - prefix = f"h{width} · " if width is not None else "" + prefix = f"hidden width {width} · " if width is not None else "" suffix = architecture_id.removeprefix("architecture_")[:8] return f"{prefix}{kind} · {suffix}" if kind else f"{prefix}{suffix}" @@ -185,6 +185,7 @@ def build_post_mip_report_payloads( registry = _json(root / "artifacts" / "post_mip" / "candidate_registry.json") revisions = dict(registry.get("revisions") or {}) + architectures = dict(registry.get("architectures") or {}) raw: dict[str, dict[str, Any]] = {} selected_by: dict[str, list[str]] = {} for node in nodes: @@ -222,12 +223,24 @@ def build_post_mip_report_payloads( for row in payload["observations"]: revision_id = str(row.get("input_revision_id") or row.get("source_revision_id") or "") architecture_id = _architecture_id(row, revisions) + architecture = architectures.get(architecture_id) + if not isinstance(architecture, Mapping): + architecture = {} + origin_kinds = sorted( + { + str(origin["kind"]) + for origin in architecture.get("origins", ()) + if isinstance(origin, Mapping) + and origin.get("kind") in {"heterogeneous", "homogeneous"} + } + ) observations.append( { **row, "architecture_id": architecture_id, "label": _candidate_label(architecture_id, revision_id, revisions), "color": _candidate_color(architecture_id), + "origin_kinds": origin_kinds, "selected_by": sorted(set(selected_by.get(revision_id, ()))), } ) @@ -277,34 +290,130 @@ def _status_summary(payload: Mapping[str, Any]) -> str: return f"

status={status}{f' · {outcomes}' if outcomes else ''}

" -def render_evaluation_report(section_id: str, payload: Mapping[str, Any]) -> str: - """Render one candidate-evaluation node.""" +_EVALUATION_METRIC_ORDER = ( + "lm_loss", + "token_accuracy_top_1", + "token_accuracy_top_10", + "token_accuracy_top_5", +) - observations = list(payload.get("observations") or ()) - metric_names = sorted( - { - str(metric) +_EVALUATION_METRIC_LABELS = { + "lm_loss": "LM loss", + "token_accuracy_top_1": "Top-1 token accuracy", + "token_accuracy_top_10": "Top-10 token accuracy", + "token_accuracy_top_5": "Top-5 token accuracy", +} + + +def _ordered_evaluation_metrics(names: set[str]) -> list[str]: + ordered = [name for name in _EVALUATION_METRIC_ORDER if name in names] + return ordered + sorted(names - set(ordered)) + + +def _evaluation_comparison_rows( + observations: list[Mapping[str, Any]], +) -> tuple[list[dict[str, Any]], list[str], bool]: + """Project structured observations into comparable alternative rows.""" + + candidate_metrics = { + str(name).removeprefix("candidate.") + for row in observations + for name in dict(row.get("metrics") or {}) + if str(name).startswith("candidate.") + } + reference_metrics = { + str(name).removeprefix("reference.") + for row in observations + for name in dict(row.get("metrics") or {}) + if str(name).startswith("reference.") + } + has_teacher = bool(reference_metrics) + has_namespaced_metrics = bool(candidate_metrics or reference_metrics) + metric_names = ( + candidate_metrics | reference_metrics + if has_namespaced_metrics + else { + str(name) for row in observations - for metric, value in dict(row.get("metrics") or {}).items() + for name, value in dict(row.get("metrics") or {}).items() if isinstance(value, (int, float)) and not isinstance(value, bool) } ) - rows = [] + rows: list[dict[str, Any]] = [] + teacher_signatures: set[str] = set() for row in observations: + metrics = dict(row.get("metrics") or {}) + if has_teacher and any(str(name).startswith("reference.") for name in metrics): + reference = {name: metrics.get(f"reference.{name}") for name in metric_names} + if any(value is not None for value in reference.values()): + signature = json.dumps(reference, sort_keys=True, default=str) + if signature not in teacher_signatures: + rows.append( + { + "alternative": "Teacher", + "candidate": "Reference checkpoint", + "status": row.get("status") or "pending", + "metrics": reference, + "evidence": "Matched reference measurement", + "error": row.get("error") or "", + } + ) + teacher_signatures.add(signature) + candidate_values = ( + {name: metrics.get(f"candidate.{name}") for name in metric_names} + if has_namespaced_metrics + else metrics + ) + origins = list(row.get("origin_kinds") or ()) or ["candidate"] + shared = len(origins) > 1 + candidate_label = ( + row.get("architecture_id") if shared else row.get("label") or row.get("architecture_id") + ) + for origin in origins: + rows.append( + { + "alternative": str(origin).replace("_", " ").title(), + "candidate": candidate_label or "unknown", + "status": row.get("status") or "pending", + "metrics": candidate_values, + "evidence": ( + "Shared physical measurement (same architecture)" + if shared + else "Measured candidate" + ), + "error": row.get("error") or "", + } + ) + return rows, _ordered_evaluation_metrics(metric_names), has_teacher + + +def render_evaluation_report(section_id: str, payload: Mapping[str, Any]) -> str: + """Render one candidate-evaluation node.""" + + observations = list(payload.get("observations") or ()) + comparison_rows, metric_names, has_teacher = _evaluation_comparison_rows(observations) + rows = [] + for row in comparison_rows: metrics = dict(row.get("metrics") or {}) cells = "".join(f"{_number(metrics.get(metric))}" for metric in metric_names) rows.append( "" - f"{_text(row.get('label') or row.get('architecture_id') or 'unknown')}" + f"{_text(row.get('alternative') or 'Candidate')}" + f"{_text(row.get('candidate') or 'unknown')}" f"{_text(row.get('status') or 'pending')}" f"{cells}" + f"{_text(row.get('evidence') or '')}" f"{_text(row.get('error') or '')}" "" ) - headings = "".join(f"{_text(metric)}" for metric in metric_names) + headings = "".join( + f"{_text(_EVALUATION_METRIC_LABELS.get(metric, metric))}" + for metric in metric_names + ) table = ( "
" - f"{headings}" + f"{headings}" + "" f"{''.join(rows)}
CandidateStatusErrorAlternativeCandidateStatusEvidenceError
" if rows else "

No evaluation observations are available yet.

" @@ -315,7 +424,15 @@ def render_evaluation_report(section_id: str, payload: Mapping[str, Any]) -> str if metric_names else "" ) - return f"

Candidate evaluation

{_status_summary(payload)}{plot}{table}" + comparison_note = ( + "

Teacher and candidate values come from the matched structured " + "evaluation observation. When heterogeneous and homogeneous solver results resolve " + "to the same architecture, both rows intentionally cite one shared physical " + "measurement.

" + if has_teacher or any(len(row.get("origin_kinds") or ()) > 1 for row in observations) + else "" + ) + return f"

Candidate evaluation

{_status_summary(payload)}{comparison_note}{plot}{table}" def render_aiperf_report(section_id: str, payload: Mapping[str, Any]) -> str: diff --git a/modelopt/torch/puzzletron/pruning/compact_runtime.py b/modelopt/torch/puzzletron/pruning/compact_runtime.py index 2ae4349d77b..1c66a1607cc 100644 --- a/modelopt/torch/puzzletron/pruning/compact_runtime.py +++ b/modelopt/torch/puzzletron/pruning/compact_runtime.py @@ -37,9 +37,24 @@ _SUPPORTED_GROUPED_ATTENTION_TYPES = { ("transformers.models.qwen3_5.modeling_qwen3_5", "Qwen3_5Attention"), + ("nemo_automodel.components.models.qwen3_next.layers", "Qwen3NextAttention"), } _SUPPORTED_GDN_TYPES = { ("transformers.models.qwen3_5.modeling_qwen3_5", "Qwen3_5GatedDeltaNet"), + ( + "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn", + "CPAwareGatedDeltaNet", + ), +} + +_NATIVE_GROUPED_ATTENTION_TYPES = { + ("nemo_automodel.components.models.qwen3_next.layers", "Qwen3NextAttention"), +} +_NATIVE_GDN_TYPES = { + ( + "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn", + "CPAwareGatedDeltaNet", + ), } @@ -58,6 +73,30 @@ def _index_select(tensor: torch.Tensor | None, dim: int, indices: torch.Tensor): return tensor.index_select(dim, _indices_on(indices, tensor)).contiguous() +def _is_dtensor(tensor: torch.Tensor) -> bool: + return ( + type(tensor).__module__.startswith("torch.distributed") + and type(tensor).__name__ == "DTensor" + ) + + +def _native_sdpa_attention(attention_module) -> bool: + if _type_key(attention_module) not in _NATIVE_GROUPED_ATTENTION_TYPES: + return True + projection_modules = tuple( + getattr(attention_module, name, None) for name in ("q_proj", "k_proj", "v_proj", "o_proj") + ) + if any(module is None or not hasattr(module, "weight") for module in projection_modules): + return False + backend = getattr(attention_module, "backend", None) + projections = tuple(getattr(module, "weight") for module in projection_modules) + return ( + getattr(backend, "attn", None) == "sdpa" + and getattr(attention_module, "attn_module", None) is None + and not any(_is_dtensor(weight) for weight in projections) + ) + + @contextmanager def _indexed_linear_forward( module, @@ -145,6 +184,7 @@ def supports_compact_grouped_attention( and orig_num_kv > 0 and head_dim > 0 and orig_num_q % orig_num_kv == 0 + and _native_sdpa_attention(attention_module) and _has_compact_grouped_attention_layout( attention_module, orig_num_q=orig_num_q, @@ -203,18 +243,21 @@ def resolve_compact_grouped_attention_target(layer, teacher_attention, child_att orig_num_kv=orig_num_kv, head_dim=head_dim, ) - if has_compact_layout and _type_key(attention_module) not in _SUPPORTED_GROUPED_ATTENTION_TYPES: + supported = supports_compact_grouped_attention( + attention_module, + orig_num_q=orig_num_q, + orig_num_kv=orig_num_kv, + head_dim=head_dim, + ) + if ( + has_compact_layout or _type_key(attention_module) in _SUPPORTED_GROUPED_ATTENTION_TYPES + ) and not supported: module_name, type_name = _type_key(attention_module) raise RuntimeError( "Compact grouped-attention scoring is unsupported for " f"{module_name}.{type_name}; refusing to score reduced geometry" ) - if not supports_compact_grouped_attention( - attention_module, - orig_num_q=orig_num_q, - orig_num_kv=orig_num_kv, - head_dim=head_dim, - ): + if not supported: return None return { "module": attention_module, @@ -300,14 +343,73 @@ def _compact_gdn_norm(norm, hidden_states, gate, value_dim_indices: torch.Tensor ) +def _gdn_gate_module(gdn_module): + holder = getattr(gdn_module, "_modules", {}).get("_fp32_params") + if holder is not None: + parameters = getattr(holder, "_parameters", {}) + if {"A_log", "dt_bias"}.issubset(parameters): + return holder + parameters = getattr(gdn_module, "_parameters", {}) + if {"A_log", "dt_bias"}.issubset(parameters): + return gdn_module + return None + + +def _gdn_gate_parameters(gdn_module) -> tuple[torch.Tensor, torch.Tensor] | None: + gate_module = _gdn_gate_module(gdn_module) + if gate_module is None: + return None + return gate_module._parameters["A_log"], gate_module._parameters["dt_bias"] + + +def _compact_gdn_gate(gdn_module, a: torch.Tensor, head_indices: torch.Tensor) -> torch.Tensor: + gate_module = _gdn_gate_module(gdn_module) + if gate_module is None: + raise RuntimeError(f"{type(gdn_module).__name__} has no supported GDN gate parameters") + a_log = _index_select(gate_module._parameters["A_log"], 0, head_indices) + dt_bias = _index_select(gate_module._parameters["dt_bias"], 0, head_indices) + if gate_module is gdn_module: + return -a_log.float().exp() * F.softplus(a.float() + dt_bias) + return torch.func.functional_call( + gate_module, + {"A_log": a_log, "dt_bias": dt_bias}, + (a,), + strict=False, + ) + + +def _native_gdn_single_device(gdn_module) -> bool: + if _type_key(gdn_module) not in _NATIVE_GDN_TYPES: + return True + cp_mesh = getattr(gdn_module, "_cp_mesh", None) + if cp_mesh is not None: + try: + if int(cp_mesh.size()) > 1: + return False + except (AttributeError, RuntimeError, TypeError, ValueError): + return False + tensors = ( + gdn_module.in_proj_qkv.weight, + gdn_module.in_proj_z.weight, + gdn_module.in_proj_a.weight, + gdn_module.in_proj_b.weight, + gdn_module.conv1d.weight, + gdn_module.norm.weight, + gdn_module.out_proj.weight, + ) + gate_parameters = _gdn_gate_parameters(gdn_module) + return gate_parameters is not None and not any( + _is_dtensor(tensor) for tensor in (*tensors, *gate_parameters) + ) + + def supports_compact_gated_delta_net(gdn_module, *, teacher_shape: GDNShape) -> bool: """Return whether this exact tested Qwen GDN layout supports compact execution.""" if _type_key(gdn_module) not in _SUPPORTED_GDN_TYPES: return False - if "_fp32_params" in getattr(gdn_module, "_modules", {}): - return False - if not {"A_log", "dt_bias"}.issubset(getattr(gdn_module, "_parameters", {})): + gate_parameters = _gdn_gate_parameters(gdn_module) + if gate_parameters is None: return False required = ( "in_proj_qkv", @@ -334,6 +436,7 @@ def supports_compact_gated_delta_net(gdn_module, *, teacher_shape: GDNShape) -> value_width = teacher_shape.num_value_heads * teacher_shape.value_head_dim return ( GDNShape.from_module(gdn_module) == teacher_shape + and _native_gdn_single_device(gdn_module) and int(gdn_module.in_proj_qkv.weight.shape[0]) == projection_width and int(gdn_module.in_proj_z.weight.shape[0]) == value_width and int(gdn_module.in_proj_a.weight.shape[0]) == teacher_shape.num_value_heads @@ -341,8 +444,8 @@ def supports_compact_gated_delta_net(gdn_module, *, teacher_shape: GDNShape) -> and int(gdn_module.conv1d.weight.shape[0]) == projection_width and int(gdn_module.norm.weight.shape[0]) == teacher_shape.value_head_dim and int(gdn_module.out_proj.weight.shape[1]) == value_width - and int(gdn_module.A_log.shape[0]) == teacher_shape.num_value_heads - and int(gdn_module.dt_bias.shape[0]) == teacher_shape.num_value_heads + and int(gate_parameters[0].shape[0]) == teacher_shape.num_value_heads + and int(gate_parameters[1].shape[0]) == teacher_shape.num_value_heads ) @@ -380,10 +483,29 @@ def compact_forward( cache_position=None, attention_mask: torch.Tensor | None = None, seq_idx=None, + position_ids=None, + qkv_format=None, + cu_seqlens=None, + cu_seqlens_cpu=None, + indices=None, + seq_index=None, **kwargs, ): + del cache_position, position_ids if kwargs: raise TypeError(f"Unsupported compact GDN forward arguments: {sorted(kwargs)}") + if ( + qkv_format not in (None, "bshd") + or seq_idx is not None + or seq_index is not None + or cu_seqlens is not None + or cu_seqlens_cpu is not None + or indices is not None + ): + raise RuntimeError( + "Compact native GDN packed execution is not supported; refusing to score " + "reduced geometry" + ) if ( attention_mask is not None and attention_mask.shape[1] > 1 @@ -459,7 +581,7 @@ def compact_forward( weight=conv_weight.squeeze(1), bias=conv_bias, activation=self.activation, - seq_idx=seq_idx, + seq_idx=None, ) else: mixed_qkv = F.silu( @@ -498,9 +620,7 @@ def compact_forward( ) beta = b.sigmoid() - g = -_index_select(self.A_log, 0, hidx).float().exp() * F.softplus( - a.float() + _index_select(self.dt_bias, 0, hidx) - ) + g = _compact_gdn_gate(self, a, hidx) repeats = target_shape.num_value_heads // target_shape.num_key_heads if repeats > 1: query = query.repeat_interleave(repeats, dim=2) diff --git a/modelopt/torch/puzzletron/scoring.py b/modelopt/torch/puzzletron/scoring.py index 549bd9accb9..7ee21b0f896 100644 --- a/modelopt/torch/puzzletron/scoring.py +++ b/modelopt/torch/puzzletron/scoring.py @@ -31,6 +31,8 @@ from .granularity import resolve_granularity from .tools.hydra_utils import register_hydra_resolvers from .tools.logger import mprint +from .tools.validate_puzzle_with_multi_replacements import load_puzzle_solutions +from .tools.validation_utils import scoring_result_matches __all__ = ["launch_scoring", "resolve_scoring_output_dir", "resolve_scoring_paths"] @@ -46,16 +48,27 @@ def extract_solution_id(filename): mprint(f"Couldn't extract solutions_id from file {filename}") -def find_missing_solutions(solutions_df, validation_dir): - all_solutions = np.arange(solutions_df.shape[0]) +def find_missing_solutions(solutions, validation_dir, scoring_args=None): + candidate_solutions = ( + solutions.to_dict(orient="records") if isinstance(solutions, pd.DataFrame) else solutions + ) + all_solutions = np.arange(len(candidate_solutions)) benchmarked_solutions = list(glob(f"{validation_dir}/solution*.json")) - benchmarked_solutions = [ - extract_solution_id(os.path.basename(s)) for s in benchmarked_solutions - ] - benchmarked_solutions = [s for s in benchmarked_solutions if s is not None] - - unbenchmarked_solutions = np.setdiff1d(all_solutions, benchmarked_solutions) + matching_solutions = [] + for result_path in benchmarked_solutions: + solution_id = extract_solution_id(os.path.basename(result_path)) + if solution_id is None or solution_id >= len(candidate_solutions): + continue + if scoring_args is not None and not scoring_result_matches( + result_path, + scoring_args, + expected_payload={"puzzle_solution": candidate_solutions[solution_id]}, + ): + continue + matching_solutions.append(solution_id) + + unbenchmarked_solutions = np.setdiff1d(all_solutions, matching_solutions) return unbenchmarked_solutions.tolist() @@ -103,13 +116,19 @@ def get_solutions_to_validate(cfg: DictConfig, num_nodes: int = 1, node_index: i _solutions_to_validate = cfg.scoring.solutions_to_validate if _solutions_to_validate is None: solutions_path, _ = resolve_scoring_paths(cfg) - single_block_replacement_solutions = pd.read_json(solutions_path) + single_block_replacement_solutions = load_puzzle_solutions( + solutions_path, + cfg.scoring.get("sort_solutions_by", None), + cfg.scoring.get("bigger_is_better", False), + ) if cfg.scoring.skip_existing_solutions: _solutions_to_validate = find_missing_solutions( - single_block_replacement_solutions, resolve_scoring_output_dir(cfg) + single_block_replacement_solutions, + resolve_scoring_output_dir(cfg), + cfg.scoring, ) else: - _solutions_to_validate = np.arange(single_block_replacement_solutions.shape[0]).tolist() + _solutions_to_validate = np.arange(len(single_block_replacement_solutions)).tolist() return partition_for_node(_solutions_to_validate, num_nodes, node_index) diff --git a/modelopt/torch/puzzletron/stages/diagnostics.py b/modelopt/torch/puzzletron/stages/diagnostics.py index b9773249c03..d49306428d8 100644 --- a/modelopt/torch/puzzletron/stages/diagnostics.py +++ b/modelopt/torch/puzzletron/stages/diagnostics.py @@ -1660,6 +1660,24 @@ def _write_reused_sort_equivalence( return merged +def _resolve_parent_sweep_sort_equivalence( + *, + parent_equivalence: dict[str, Any], + sort_summary_path: Path, + reuse_summary_path: Path, + reuse_sort_equivalence: bool, +) -> dict[str, Any]: + """Resolve sort evidence from the already-qualified stage when reuse is requested.""" + + if not reuse_sort_equivalence: + return parent_equivalence + return _write_reused_sort_equivalence( + sort_summary_path, + reuse_summary_path, + {"reused_source_summary": str(sort_summary_path)}, + ) + + def _parent_sweep_sanity_verdict(width_summary: dict[str, Any], sort_summary: dict[str, Any]): """Combine advisory width quality with blocking reused-sort correctness.""" @@ -2723,18 +2741,24 @@ def _activation_diagnostic_parent_sweep( diag_cfg=diag_cfg, ) - activation_equivalence = ( + parent_equivalence = ( (sweep_manifest.get("parents") or {}).get("activation") or {} ).get("equivalence") or {} + sort_equivalence_dir = puzzle_dir / "artifacts" / "sort_sanity" + sort_equivalence_dir.mkdir(parents=True, exist_ok=True) + sort_summary_path = sort_equivalence_dir / "summary.json" + reused_sort_summary_path = artifacts_dir / "reused_sort_equivalence.json" + activation_equivalence = _resolve_parent_sweep_sort_equivalence( + parent_equivalence=parent_equivalence, + sort_summary_path=sort_summary_path, + reuse_summary_path=reused_sort_summary_path, + reuse_sort_equivalence=bool(diag_cfg.get("reuse_sort_equivalence", False)), + ) equivalence_findings = [ {**finding, "stage": "sort_sanity", "severity": "error"} for finding in activation_equivalence.get("findings") or () ] sort_passed = activation_equivalence.get("passed") is True - sort_equivalence_dir = puzzle_dir / "artifacts" / "sort_sanity" - sort_equivalence_dir.mkdir(parents=True, exist_ok=True) - sort_summary_path = sort_equivalence_dir / "summary.json" - reused_sort_summary_path = artifacts_dir / "reused_sort_equivalence.json" reuse_sort_summary = { "passed": sort_passed, "reused_parent_sweep": True, diff --git a/modelopt/torch/puzzletron/tools/validation_utils.py b/modelopt/torch/puzzletron/tools/validation_utils.py index f7924a3d818..f0e6e95b318 100644 --- a/modelopt/torch/puzzletron/tools/validation_utils.py +++ b/modelopt/torch/puzzletron/tools/validation_utils.py @@ -1,10 +1,24 @@ # SPDX-FileCopyrightText: Copyright (c) 2024-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. """Stable validation-artifact writer shared by AutoModel scoring paths.""" from __future__ import annotations +import json +from collections.abc import Mapping from pathlib import Path from typing import Any @@ -12,7 +26,45 @@ from modelopt.torch.utils import json_dump -__all__ = ["write_results"] +__all__ = ["scoring_result_matches", "write_results"] + +_OPERATIONAL_SCORING_KEYS = frozenset( + {"force_rescore", "skip_existing_solutions", "solutions_to_validate"} +) + + +def _resolved_scoring_args(args: Any) -> dict[str, Any]: + if isinstance(args, DictConfig): + value = OmegaConf.to_container(args, resolve=True) + elif isinstance(args, Mapping): + value = dict(args) + else: + value = vars(args) + if not isinstance(value, dict): + raise TypeError(f"scoring arguments must resolve to a mapping, got {type(value).__name__}") + return {key: item for key, item in value.items() if key not in _OPERATIONAL_SCORING_KEYS} + + +def scoring_result_matches( + path: str | Path, + args: Any, + *, + expected_payload: Mapping[str, Any] | None = None, +) -> bool: + """Return whether an existing score matches its config and expected identity fields.""" + + try: + result = json.loads(Path(path).read_text()) + except (OSError, json.JSONDecodeError): + return False + recorded = result.get("args") if isinstance(result, dict) else None + if not isinstance(recorded, dict): + return False + if _resolved_scoring_args(recorded) != _resolved_scoring_args(args): + return False + return expected_payload is None or all( + result.get(key) == value for key, value in expected_payload.items() + ) def write_results( diff --git a/puzzletron_setup/inspection.py b/puzzletron_setup/inspection.py index bbb92766884..da5ab31aefb 100644 --- a/puzzletron_setup/inspection.py +++ b/puzzletron_setup/inspection.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. """Configuration-only model and dataset inspection for Puzzletron setup.""" @@ -13,7 +25,7 @@ from urllib.parse import unquote, urlparse import yaml -from huggingface_hub import HfApi +from huggingface_hub import HfApi, scan_cache_dir, try_to_load_from_cache from transformers import AutoConfig, PretrainedConfig from . import SetupError @@ -39,7 +51,6 @@ def _normalize_source(source: str, *, kind: str, url_prefix: str | None = None) -> str: """Normalize an existing local path or Hugging Face web URL.""" - source = source.strip() if not source: raise SetupError(f"Enter a {kind} path or Hugging Face URL.") @@ -80,13 +91,11 @@ def _normalize_source(source: str, *, kind: str, url_prefix: str | None = None) def normalize_model_source(source: str) -> str: """Normalize an existing model path or Hugging Face model web URL.""" - return _normalize_source(source, kind="model") def normalize_dataset_source(source: str) -> str: """Normalize an existing dataset path or Hugging Face dataset web URL.""" - return _normalize_source(source, kind="dataset", url_prefix="datasets") @@ -144,9 +153,30 @@ def _load_config_dict(source: str, *, revision: str | None, local: bool) -> dict return dict(config) +def _cached_remote_config(source: str, revision: str | None) -> tuple[Path, str] | None: + cached = try_to_load_from_cache(source, "config.json", revision=revision or "main") + if isinstance(cached, str) and Path(cached).is_file(): + # Preserve the snapshots//config.json path. Resolving its blob + # symlink would discard the commit directory needed for provenance. + path = Path(cached).expanduser().absolute() + return path, path.parent.name + if revision is not None: + return None + try: + matching = [ + (Path(item.snapshot_path).expanduser().absolute() / "config.json", item.commit_hash) + for repo in scan_cache_dir().repos + if repo.repo_type == "model" and repo.repo_id == source + for item in repo.revisions + ] + except Exception: + return None + available = [(path, commit) for path, commit in matching if path.is_file()] + return available[0] if len(available) == 1 else None + + def inspect_model(source: str, revision: str | None = None) -> InspectedModel: """Inspect a local path or Hugging Face URL without loading model weights.""" - source = normalize_model_source(source) expanded = Path(source).expanduser() is_local = expanded.exists() @@ -161,12 +191,21 @@ def inspect_model(source: str, revision: str | None = None) -> InspectedModel: try: resolved_revision = HfApi().model_info(config_source, revision=revision).sha except Exception as error: - raise SetupError(f"Cannot resolve Hugging Face model {source!r}: {error}") from error - config = _load_config_dict( - config_source, - revision=resolved_revision or revision, - local=is_local, - ) + cached = _cached_remote_config(config_source, revision) + if cached is None: + raise SetupError( + f"Cannot resolve Hugging Face model {source!r}: {error}" + ) from error + cached_path, resolved_revision = cached + config = _load_config_dict(str(cached_path), revision=None, local=True) + else: + config = _load_config_dict( + config_source, + revision=resolved_revision, + local=False, + ) + else: + config = _load_config_dict(config_source, revision=revision, local=True) profile = resolve_profile(config) return InspectedModel( source=effective_source, @@ -216,7 +255,6 @@ def _local_dataset_metadata(path: Path) -> Any: def infer_dataset_modality(source: str) -> ModalityFinding: """Infer text versus multimodal data and return explicit evidence.""" - source = normalize_dataset_source(source) path = Path(source).expanduser() if path.exists(): diff --git a/puzzletron_setup/v2/bundle.py b/puzzletron_setup/v2/bundle.py index cc6fdac527c..99ff2e262f0 100644 --- a/puzzletron_setup/v2/bundle.py +++ b/puzzletron_setup/v2/bundle.py @@ -232,15 +232,55 @@ def _render_experiment_v2( if config.mip_configured: rendered["mip"] = _deep_merge(rendered.get("mip") or {}, mip) rendered["mip"]["enabled"] = True + if budget == "smoke": + for run in _mapping(rendered["mip"].get("runs")).values(): + solver = _mapping(run.get("solver")) + solver["num_solutions"] = min(int(solver.get("num_solutions", 2) or 2), 2) + run["solver"] = solver + homogeneous = _mapping(run.get("homogeneous")) + if homogeneous: + homogeneous["keep"] = min(int(homogeneous.get("keep", 2) or 2), 2) + run["homogeneous"] = homogeneous flows = _plain(config.post_mip_flows) if config.post_mip_flows_configured: rendered["post_mip"] = {"flows": deepcopy(flows)} if budget == "smoke": for flow in rendered["post_mip"]["flows"].values(): for node in _mapping(flow.get("nodes")).values(): - if node.get("type") == "downstream_evaluation": - node_config = _mapping(node.get("config")) + node_type = node.get("type") + node_config = _mapping(node.get("config")) + if node_type == "downstream_evaluation" and "profile" not in node_config: node_config["limit"] = min(int(node_config.get("limit", 8) or 8), 8) + elif node_type == "evaluation": + node_config["eval_samples"] = min( + int(node_config.get("eval_samples", 2) or 2), 2 + ) + node_config["block_size"] = min( + int(node_config.get("block_size", 512) or 512), 512 + ) + elif node_type == "aiperf": + node_config["input_tokens"] = min( + int(node_config.get("input_tokens", 100) or 100), 100 + ) + node_config["output_tokens"] = min( + int(node_config.get("output_tokens", 80) or 80), 80 + ) + node_config["request_count"] = min( + int(node_config.get("request_count", 1) or 1), 1 + ) + extra_inputs = _mapping(node_config.get("extra_inputs")) + if extra_inputs: + extra_inputs["min_tokens"] = min( + int(extra_inputs.get("min_tokens", 80) or 80), 80 + ) + node_config["extra_inputs"] = extra_inputs + elif node_type == "global_kd": + node_config["max_steps"] = min(int(node_config.get("max_steps", 2) or 2), 2) + node_config["global_batch_size"] = 1 + node_config["checkpoint_every_steps"] = 2 + elif node_type == "filter": + node["top_k"] = min(int(node.get("top_k", 1) or 1), 1) + if node_config: node["config"] = node_config batch_mirrors = { @@ -306,7 +346,7 @@ def _render_execution_v2( for stage_id, resource in config.stage_resources.items(): entry = { "strategy": resource.strategy, - "instances": resource.instances, + "instances": min(resource.instances, 1) if budget == "smoke" else resource.instances, "resource": resource.resource, "gpus_per_node": ( resource.gpus_per_node if resource.gpus_per_node is not None else default_gpus @@ -366,6 +406,10 @@ def _bundle_readme( "# Puzzletron campaign", "", "Generated by `puzzletron_setup_v2.py`. The wizard did not launch any jobs.", + ( + "Run setup and orchestrator commands from the Python 3.10+ controller venv " + f"described in `{Path(repository) / 'examples/puzzletron/docs/environment_setup.md'}`." + ), ] acquisition = _mapping(acquisition) if acquisition: @@ -415,6 +459,11 @@ def _bundle_readme( "", "## Prepare dataset", "", + ( + "This is a worker command. Run it inside the worker venv or container " + "configured in `smoke/runner.yaml` and `production/runner.yaml`, not in " + "the lightweight controller venv." + ), "The command is idempotent only when the existing manifest matches these answers.", ( "If a nonempty destination contains a different manifest, choose a new output " @@ -463,6 +512,16 @@ def _bundle_readme( ), ) ) + if acquisition.get("adapter") == "nemotron_vlm_v2": + evaluation_guide = Path(repository) / ( + "examples/puzzletron/docs/vlm_checkpoint_evaluation.md" + ) + lines.extend( + ( + "", + f"For VLM evaluator caches, follow `{evaluation_guide}`.", + ) + ) for budget, heading, introduction, launch_introduction in sections: bundle = campaign_dir / budget orchestrator_args = [ diff --git a/puzzletron_setup/v2/post_mip.py b/puzzletron_setup/v2/post_mip.py index 072b2b0de0d..32b5d0716a8 100644 --- a/puzzletron_setup/v2/post_mip.py +++ b/puzzletron_setup/v2/post_mip.py @@ -334,8 +334,7 @@ def recommended_flow( **({"topology": deepcopy(serving["topology"])} if serving.get("topology") else {}), } serving_metric = ( - f"{serving_id}.images_{max(image_batch_sizes)}.concurrency_{concurrency[0]}." - "image_throughput" + f"{serving_id}.images_{max(image_batch_sizes)}.image_throughput" if multimodal else f"{serving_id}.output_token_throughput" ) diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index f1c93e6eada..9813db8db51 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -3582,7 +3582,11 @@ def post_mip_section(session: WizardSession, resolver: DefaultsResolver, context sequence = int(session.state.get_field("data.sequence_length", 4096)) workloads = _mapping_copy(session.state.collection("serving_workloads")) first_workload = next(iter(workloads.values()), {}) + serving_defaults = resolver.resolve("post_mip.serving", {}).value + if not isinstance(serving_defaults, Mapping): + raise SetupError("post_mip.serving defaults must be a mapping") serving = { + **deepcopy(dict(serving_defaults)), "input_tokens": int(first_workload.get("prefill_seq_len", sequence)), "output_tokens": int(first_workload.get("generation_seq_len", 1024)), "concurrency": [int(first_workload.get("max_num_seqs", 1))], diff --git a/tests/unit/torch/puzzletron/evaluation/vlm/preparation/test_benchmark_data.py b/tests/unit/torch/puzzletron/evaluation/vlm/preparation/test_benchmark_data.py index 6a83bab18b9..28d5003bae0 100644 --- a/tests/unit/torch/puzzletron/evaluation/vlm/preparation/test_benchmark_data.py +++ b/tests/unit/torch/puzzletron/evaluation/vlm/preparation/test_benchmark_data.py @@ -18,6 +18,7 @@ import hashlib import io import json +import os import stat import tarfile import threading @@ -84,6 +85,15 @@ def _emulate_atomic_exchange(first: Path, second: Path) -> bool: return True +def _rewrite_with_distinct_mtime(path: Path, payload: bytes) -> None: + """Make metadata-based invalidation deterministic on coarse-clock filesystems.""" + before = path.stat() + path.write_bytes(payload) + after = path.stat() + if after.st_mtime_ns == before.st_mtime_ns: + os.utime(path, ns=(after.st_atime_ns, before.st_mtime_ns + 1_000_000_000)) + + def test_atomic_exchange_directories_when_supported(tmp_path): first = tmp_path / "first" second = tmp_path / "second" @@ -226,6 +236,34 @@ def record_hash(path): assert hashed == [hf_home / "mmvu/videos/sample.mp4"] +def test_prepared_media_reuse_tolerates_distributed_filesystem_mtime_skew(tmp_path): + root = tmp_path / "prepared" + root.mkdir() + (root / "sample.mp4").write_bytes(b"video") + inventory = preparation._inventory(root) + inventory[0]["mtime_ns"] += 1_000_000_000 + + assert preparation._inventory_is_current(root, inventory) + + +def test_prepared_media_reuse_reports_paths_from_the_current_mount(tmp_path): + hf_home = tmp_path / "hf-home" + snapshot = preparation._hub_snapshot(hf_home, "mmvu_val") + snapshot.mkdir(parents=True) + _write_zip(snapshot / "videos.zip", {"videos/sample.mp4": b"video"}) + preparation._prepare(hf_home, "mmvu_val", snapshot) + marker = hf_home / "mmvu" / preparation._MARKER_NAME + payload = json.loads(marker.read_text()) + payload["snapshot"] = "/stale-mount/snapshot" + payload["media_root"] = "/stale-mount/media" + marker.write_text(json.dumps(payload)) + + reused = preparation._prepare(hf_home, "mmvu_val", snapshot) + + assert reused["snapshot"] == str(snapshot) + assert reused["media_root"] == str(hf_home / "mmvu") + + @pytest.mark.parametrize("damage", ["missing", "corrupt", "unexpected"]) def test_complete_media_marker_repairs_owned_root_from_pinned_snapshot( tmp_path, monkeypatch, damage @@ -242,7 +280,7 @@ def test_complete_media_marker_repairs_owned_root_from_pinned_snapshot( if damage == "missing": media.unlink() elif damage == "corrupt": - media.write_bytes(b"wrong") + _rewrite_with_distinct_mtime(media, b"wrong") elif damage == "unexpected": (target / "unexpected.bin").write_bytes(b"stale") report = preparation._prepare(hf_home, "mmvu_val", snapshot) @@ -265,7 +303,7 @@ def test_repair_without_atomic_exchange_preserves_live_root(monkeypatch, tmp_pat preparation._prepare(hf_home, "mmvu_val", snapshot) target = hf_home / "mmvu" media = target / "videos/sample.mp4" - media.write_bytes(b"wrong") + _rewrite_with_distinct_mtime(media, b"wrong") monkeypatch.setattr(preparation, "_atomic_exchange_directories", lambda *_args: False) with pytest.raises(RuntimeError, match="atomic media-directory exchange is unavailable"): @@ -343,7 +381,7 @@ def test_snapshot_inventory_rejects_partial_and_same_size_corruption(tmp_path): second.write_bytes(b"two") report = preparation._snapshot_inventory_report(hf_home, "realworldqa", snapshot) assert preparation._snapshot_inventory_is_current(report) - first.write_bytes(b"bad") + _rewrite_with_distinct_mtime(first, b"bad") assert not preparation._snapshot_inventory_is_current(report) @@ -376,7 +414,7 @@ def record_hash(path): assert hashed == [sample] hashed.clear() - sample.write_bytes(b"two") + _rewrite_with_distinct_mtime(sample, b"two") refreshed = preparation._snapshot_inventory_report(hf_home, "realworldqa", snapshot) assert hashed == [sample] assert refreshed["files"][0]["sha256"] == hashlib.sha256(b"two").hexdigest() diff --git a/tests/unit/torch/puzzletron/evaluation/vlm/test_post_mip.py b/tests/unit/torch/puzzletron/evaluation/vlm/test_post_mip.py index 6eb4d2ef635..b79ab4b3c25 100644 --- a/tests/unit/torch/puzzletron/evaluation/vlm/test_post_mip.py +++ b/tests/unit/torch/puzzletron/evaluation/vlm/test_post_mip.py @@ -66,6 +66,7 @@ def fake_runner(checkpoint_path, *, output_root, settings): "batch_size": 1, "timeout_seconds": 900, "dtype": "bfloat16", + "gdn_prefill_backend": "triton", "topology": {"tensor_parallel_size": 1}, }, ) @@ -76,6 +77,7 @@ def fake_runner(checkpoint_path, *, output_root, settings): assert captured["settings"]["limit"] == 2 assert captured["settings"]["timeout_seconds"] == 900 assert captured["settings"]["dtype"] == "bfloat16" + assert captured["settings"]["gdn_prefill_backend"] == "triton" assert captured["settings"]["topology"] == {"tensor_parallel_size": 1} assert result["metrics"] == {"modelopt_vlm_benchmark_realworldqa.accuracy": 0.5} assert result["profile_path"] == str(output / "profile.json") @@ -83,6 +85,7 @@ def fake_runner(checkpoint_path, *, output_root, settings): assert result["contract"]["model_backend"] == "vllm" assert result["contract"]["post_mip_runner_overrides"] == { "dtype": "bfloat16", + "gdn_prefill_backend": "triton", "topology": {"tensor_parallel_size": 1}, } diff --git a/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py b/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py index 0e12684f870..271d63bdbab 100644 --- a/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py +++ b/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py @@ -21,8 +21,10 @@ from contextlib import contextmanager from types import SimpleNamespace +import pandas as pd import pytest import torch +from omegaconf import OmegaConf from safetensors.torch import save_file from torch import nn @@ -33,7 +35,7 @@ MLAConfig, MoEConfig, ) -from modelopt.torch.puzzletron.plugins.automodel import solution_recipe +from modelopt.torch.puzzletron.plugins.automodel import solution_launch, solution_recipe from modelopt.torch.puzzletron.plugins.automodel.solution_launch import ( _candidate_execution_context, _quarantine_failed_realization, @@ -51,11 +53,10 @@ ) 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.scoring import find_missing_solutions def test_baseline_only_scoring_does_not_require_candidate_solutions(tmp_path): - from modelopt.torch.puzzletron.plugins.automodel import solution_launch - solutions, pending_ids = solution_launch._load_solution_work( {"baseline_only": True}, tmp_path, @@ -65,7 +66,47 @@ def test_baseline_only_scoring_does_not_require_candidate_solutions(tmp_path): assert pending_ids == [] -def test_native_automodel_gdn_rejects_compact_runtime_candidate(): +def test_solution_work_reuses_only_matching_config_and_identity(monkeypatch, tmp_path): + scoring_args = { + "solutions_path": str(tmp_path / "solutions.json"), + "eval_samples": 16, + "skip_existing_solutions": True, + "solutions_to_validate": None, + } + scoring = OmegaConf.create(scoring_args) + solutions = [{"width": 1}, {"width": 2}, {"width": 3}] + monkeypatch.setattr(solution_launch, "load_puzzle_solutions", lambda *args: solutions) + recorded = {**scoring_args, "solutions_to_validate": [0, 1]} + (tmp_path / "solution_0.json").write_text( + json.dumps({"args": recorded, "puzzle_solution": solutions[0]}) + ) + stale = {**scoring_args, "eval_samples": 64} + (tmp_path / "solution_1.json").write_text( + json.dumps({"args": stale, "puzzle_solution": solutions[1]}) + ) + (tmp_path / "solution_2.json").write_text( + json.dumps({"args": recorded, "puzzle_solution": {"width": 4}}) + ) + + solutions, pending_ids = solution_launch._load_solution_work(scoring, tmp_path) + + assert pending_ids == [1, 2] + + assert find_missing_solutions(pd.DataFrame(solutions), tmp_path, scoring) == [1, 2] + + parent_identity = {"parent_role": "sorted", "checkpoint_dir": "/checkpoints/sorted"} + (tmp_path / "parent.json").write_text(json.dumps({"args": recorded, **parent_identity})) + assert solution_launch.scoring_result_matches( + tmp_path / "parent.json", scoring, expected_payload=parent_identity + ) + assert not solution_launch.scoring_result_matches( + tmp_path / "parent.json", + scoring, + expected_payload={**parent_identity, "checkpoint_dir": "/checkpoints/changed"}, + ) + + +def test_native_automodel_gdn_supports_single_device_compact_runtime_candidate(): # Optional dependency: native AutoModel Qwen modules are not installed in every test env. pytest.importorskip("nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn") from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import CPAwareGatedDeltaNet @@ -107,7 +148,7 @@ def test_native_automodel_gdn_rejects_compact_runtime_candidate(): child = BlockConfig( subblock_configs=( MambaConfig( - num_heads=shape.num_value_heads, + num_heads=shape.num_value_heads // 2, head_dim=shape.value_head_dim, num_groups=shape.num_key_heads // 2, state_dim=shape.key_head_dim, @@ -117,9 +158,14 @@ def test_native_automodel_gdn_rejects_compact_runtime_candidate(): original_forward = gdn.forward.__func__ original_state = set(vars(gdn)) original_forward_hooks = dict(gdn._forward_hooks) + hidden_states = torch.randn(1, 4, config.hidden_size) - with pytest.raises(RuntimeError, match="refusing to score reduced geometry"): - apply_runtime_candidate(layer, teacher, child) + handle = apply_runtime_candidate(layer, teacher, child) + assert "forward" in vars(gdn) + with torch.no_grad(): + output = gdn(hidden_states, qkv_format="bshd") + assert output.shape == hidden_states.shape + handle.remove() assert gdn.forward.__func__ is original_forward assert set(vars(gdn)) == original_state diff --git a/tests/unit/torch/puzzletron/test_compact_runtime.py b/tests/unit/torch/puzzletron/test_compact_runtime.py index 6352df3bd0c..ce7662703cb 100644 --- a/tests/unit/torch/puzzletron/test_compact_runtime.py +++ b/tests/unit/torch/puzzletron/test_compact_runtime.py @@ -100,7 +100,7 @@ def test_compact_grouped_attention_target_requires_reduced_supported_geometry(): assert resolve_compact_grouped_attention_target(layer, teacher, teacher) is None -def test_compact_grouped_attention_rejects_native_automodel_backend(): +def test_compact_grouped_attention_dispatches_native_automodel_sdpa_backend(): # Optional dependency: native AutoModel Qwen modules are not installed in every test env. pytest.importorskip("nemo_automodel.components.models.qwen3_next.layers") from nemo_automodel.components.models.common import BackendConfig @@ -143,14 +143,21 @@ def test_compact_grouped_attention_rejects_native_automodel_backend(): original_forward = attention.forward.__func__ original_state = set(vars(attention)) - assert not supports_compact_grouped_attention( + assert supports_compact_grouped_attention( attention, orig_num_q=4, orig_num_kv=2, head_dim=8, ) - with pytest.raises(RuntimeError, match="refusing to score reduced geometry"): - resolve_compact_grouped_attention_target(layer, teacher, child) + target = resolve_compact_grouped_attention_target(layer, teacher, child) + assert target == { + "module": attention, + "orig_num_q": 4, + "orig_num_kv": 2, + "target_num_q": 2, + "target_num_kv": 1, + "head_dim": 8, + } assert attention.forward.__func__ is original_forward assert set(vars(attention)) == original_state diff --git a/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py b/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py index 7ff8832e3df..6a7b98c0ee7 100644 --- a/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py +++ b/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py @@ -25,6 +25,7 @@ _merge_reused_sort_equivalence, _parent_sweep_sanity_verdict, _ratio_aligned_hidden_widths, + _resolve_parent_sweep_sort_equivalence, _select_layers, _write_hidden_only_diagnostic_artifacts, _write_reused_sort_equivalence, @@ -237,6 +238,24 @@ def test_reused_parent_sweep_does_not_mutate_completed_sort_artifact(tmp_path): assert merged["reused_parent_sweep"] is True +def test_parent_sweep_reuses_canonical_sort_verdict_when_equivalence_was_skipped(tmp_path): + sort_summary_path = tmp_path / "sort_sanity" / "summary.json" + reuse_summary_path = tmp_path / "width_sanity" / "reused_sort_equivalence.json" + sort_summary_path.parent.mkdir() + sort_summary_path.write_text(json.dumps({"passed": True, "findings": []})) + + resolved = _resolve_parent_sweep_sort_equivalence( + parent_equivalence={}, + sort_summary_path=sort_summary_path, + reuse_summary_path=reuse_summary_path, + reuse_sort_equivalence=True, + ) + + assert resolved["passed"] is True + assert resolved["reused_source_summary"] == str(sort_summary_path) + assert json.loads(reuse_summary_path.read_text()) == resolved + + @pytest.mark.parametrize("invalid_summary", [None, [], "not-an-object"]) def test_reused_parent_sweep_rejects_non_object_sort_artifacts(tmp_path, invalid_summary): sort_summary_path = tmp_path / "sort_sanity" / "summary.json" diff --git a/tests/unit/torch/puzzletron/test_lmms_evaluation.py b/tests/unit/torch/puzzletron/test_lmms_evaluation.py index bd86a120693..b336ce311fb 100644 --- a/tests/unit/torch/puzzletron/test_lmms_evaluation.py +++ b/tests/unit/torch/puzzletron/test_lmms_evaluation.py @@ -20,6 +20,7 @@ import shutil import subprocess import sys +import time from pathlib import Path from types import SimpleNamespace @@ -168,6 +169,7 @@ def test_command_maps_checkpoint_and_vllm_topology(tmp_path): }, "model_args": {"dtype": "bfloat16"}, "chat_template": "/templates/qwen35_no_thinking.jinja", + "gdn_prefill_backend": "triton", }, checkpoint="/ckpts/candidate", output_path=tmp_path / "results", @@ -182,6 +184,7 @@ def test_command_maps_checkpoint_and_vllm_topology(tmp_path): assert "tensor_parallel_size=4" in model_args assert "pipeline_parallel_size=2" in model_args assert "chat_template=/templates/qwen35_no_thinking.jinja" in model_args + assert "gdn_prefill_backend=triton" in model_args assert "gpu_group_size" not in model_args assert env["LMMS_EVAL_HOME"] == str(tmp_path / "cache") assert timeout == 123 @@ -357,26 +360,84 @@ def test_command_rejects_unsupported_backend_contract(tmp_path, settings, expect @pytest.mark.skipif(os.name != "posix", reason="process groups are POSIX-specific") def test_timeout_kills_ignored_process_group_members(monkeypatch, tmp_path): script = ( - "import signal,time; " + "import signal,subprocess,sys,time; " "signal.signal(signal.SIGTERM, signal.SIG_IGN); " "print('partial stdout', flush=True); " + "subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(4)'], " + "start_new_session=True); " "time.sleep(60)" ) monkeypatch.setattr(lmms, "_PROCESS_CLEANUP_TIMEOUT_SECONDS", 0.1) + started = time.monotonic() with pytest.raises(lmms.LmmsEvalTimeoutError) as exc_info: lmms._run_process( [sys.executable, "-c", script], cwd=str(tmp_path), env=os.environ.copy(), - timeout=1.0, + timeout=0.5, ) - assert exc_info.value.timeout == 1.0 + assert time.monotonic() - started < 2.0 + assert exc_info.value.timeout == 0.5 assert exc_info.value.output == "partial stdout\n" assert exc_info.value.stderr == "" +def test_run_process_streams_progress_and_bounds_inherited_pipes(monkeypatch, tmp_path, capsys): + progress_path = tmp_path / "progress.json" + env = { + **os.environ, + lmms._PROGRESS_PATH_ENV: str(progress_path), + lmms._PROGRESS_TASKS_ENV: json.dumps( + [{"name": "realworldqa", "total": 64}, {"name": "mmmu_val", "total": 120}] + ), + } + script = ( + "import subprocess,sys; " + "print('evaluator started', flush=True); " + "print('Model Responding: 21%|##| 38/184 [00:10<00:40, 3.80it/s]', " + "file=sys.stderr, flush=True); " + "subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(4)'], " + "start_new_session=sys.platform != 'win32')" + ) + monkeypatch.setattr(lmms, "_PROCESS_CLEANUP_TIMEOUT_SECONDS", 0.1) + + started = time.monotonic() + result = lmms._run_process( + [sys.executable, "-c", script], + cwd=str(tmp_path), + env=env, + timeout=5.0, + ) + + assert time.monotonic() - started < 2.0 + captured = capsys.readouterr() + progress = json.loads(progress_path.read_text()) + assert captured.out == "evaluator started\n" + assert "38/184" in captured.err + assert result.stdout == captured.out + assert result.stderr == captured.err + assert progress["unit"] == "samples" + assert progress["current"] == 38 + assert progress["total"] == 184 + assert progress["rate_per_second"] == 3.8 + assert progress["task"] == {"name": "realworldqa", "current": 38, "total": 64} + + +def test_evaluation_progress_without_denominator_does_not_invent_one(): + progress = lmms._evaluation_progress_payload( + "Model Responding: 38it [00:45, 1.20s/it]", + [{"name": "realworldqa", "total": 64}], + ) + + assert progress is not None + assert progress["current"] == 38 + assert progress["total"] is None + assert progress["rate_per_second"] == pytest.approx(1 / 1.2) + assert "task" not in progress + + def test_run_checkpoint_flattens_metrics_and_preserves_artifacts(monkeypatch, tmp_path): checkpoint = tmp_path / "checkpoint" checkpoint.mkdir() @@ -401,7 +462,14 @@ def fake_run(argv, *, cwd, env, timeout): result = lmms.run_lmms_eval_checkpoint( checkpoint, output_root=tmp_path / "results", - settings={**_settings("ifeval", "gsm8k"), "limit": 4}, + settings={ + **_settings("ifeval", "gsm8k"), + "limit": 4, + "progress_tasks": [ + {"name": "ifeval", "total": 4}, + {"name": "gsm8k", "total": 4}, + ], + }, ) assert result["metrics"] == { @@ -416,6 +484,10 @@ def fake_run(argv, *, cwd, env, timeout): assert summary["raw_result_path"] == result["raw_result_path"] assert "result_path" not in summary assert summary["sample_counts"] == {"gsm8k": 4.0, "ifeval": 4.0} + progress = json.loads((Path(result["result_path"]).parent / "progress.json").read_text()) + assert progress["status"] == "completed" + assert progress["current"] == progress["total"] == 8 + assert progress["task"] == {"name": "gsm8k", "current": 4, "total": 4} def test_run_checkpoint_executes_real_process(tmp_path): @@ -472,6 +544,8 @@ def test_run_checkpoint_preserves_failure_artifacts(monkeypatch, tmp_path): assert Path(error.command_path).is_file() assert Path(error.stdout_path).read_text() == "partial evaluator output\n" assert Path(error.stderr_path).read_text() == "backend failed\n" + progress = json.loads((Path(error.stderr_path).parent / "progress.json").read_text()) + assert progress["status"] == "failed" def test_run_checkpoint_preserves_timeout_artifacts(monkeypatch, tmp_path): @@ -499,6 +573,8 @@ def time_out(argv, **_kwargs): assert Path(error.command_path).is_file() assert Path(error.stdout_path).read_text() == "partial evaluator output\n" assert Path(error.stderr_path).read_text() == "evaluation timed out\n" + progress = json.loads((Path(error.stderr_path).parent / "progress.json").read_text()) + assert progress["status"] == "timed_out" def test_run_checkpoint_preserves_timeout_artifacts_if_attempt_directory_disappears( diff --git a/tests/unit/torch/puzzletron/test_orchestration_reporting.py b/tests/unit/torch/puzzletron/test_orchestration_reporting.py index 38a4c7c010f..de11fe7b0a3 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_reporting.py +++ b/tests/unit/torch/puzzletron/test_orchestration_reporting.py @@ -200,6 +200,20 @@ def test_clean_completion_regenerates_tampered_final_report(tmp_path: Path): assert result["report_status"] == "completed" +def test_clean_completion_regenerates_after_stage_state_changes(tmp_path: Path): + plan = _plan(tmp_path) + executor = _ReportExecutor(JobState.COMPLETED) + CampaignController(plan, executor=executor, poll_interval_seconds=0).run() + stage_root = plan.puzzle_dir / "orchestration/stages" + stage_root.mkdir(parents=True) + (stage_root / "later-stage.json").write_text('{"status": "completed"}\n') + + result = CampaignController(plan, executor=executor, poll_interval_seconds=0).run() + + assert [attempt.stage_id for attempt in executor.submitted] == ["final_report", "final_report"] + assert result["report_status"] == "completed" + + def test_clean_completion_regenerates_oversized_completion_record(tmp_path: Path): plan = _plan(tmp_path) executor = _ReportExecutor(JobState.COMPLETED) diff --git a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py index b86ffb11b84..c76b142d514 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py +++ b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py @@ -36,7 +36,9 @@ load_runner_config, ) from puzzletron_orchestrator.controller import CampaignController +from puzzletron_orchestrator.dashboard import format_eta from puzzletron_orchestrator.executors.base import Executor +from puzzletron_orchestrator.progress import summarize_stage_artifacts from puzzletron_orchestrator.schema import ( AttemptSpec, JobHandle, @@ -50,6 +52,69 @@ from collections.abc import Sequence +def _seed_evaluation_progress( + run_dir: Path, + *, + current: int, + total: int | None, + task: dict[str, object] | None, +) -> dict[str, object]: + input_root = run_dir / "artifacts/post_mip/nodes/selected" + (input_root / "executions/input-1").mkdir(parents=True) + (input_root / "current.json").write_text(json.dumps({"execution_identity": "input-1"})) + (input_root / "executions/input-1/candidate_set.json").write_text( + json.dumps({"revision_ids": ["candidate-1"]}) + ) + progress_root = ( + run_dir + / "artifacts/post_mip/nodes/eval/executions/post_mip_execution_1" + / "raw/candidate-1/lmms_eval/attempt_1" + ) + progress_root.mkdir(parents=True) + payload: dict[str, object] = { + "schema": "modelopt.puzzletron.evaluation-progress/v1", + "status": "running", + "unit": "samples", + "current": current, + "total": total, + "rate_per_second": 3.8, + } + if task is not None: + payload["task"] = task + (progress_root / "progress.json").write_text(json.dumps(payload)) + return { + "post_mip": { + "flows": { + "quality": { + "nodes": {"eval": {"type": "downstream_evaluation", "input": "selected"}} + } + } + } + } + + +def test_evaluation_progress_uses_structured_task_samples(tmp_path): + config = _seed_evaluation_progress( + tmp_path, + current=38, + total=184, + task={"name": "realworldqa", "current": 38, "total": 64}, + ) + + detail = summarize_stage_artifacts(tmp_path, "post.quality.eval", config=config) + + assert detail == "evaluation realworldqa 38/64 samples at 3.80 samples/s" + + +def test_evaluation_progress_preserves_unknown_denominator(tmp_path): + config = _seed_evaluation_progress(tmp_path, current=38, total=None, task=None) + + detail = summarize_stage_artifacts(tmp_path, "post.quality.eval", config=config) + + assert detail == "evaluation 38 samples (total unavailable) at 3.80 samples/s" + assert format_eta(None) == "unavailable" + + class _FakeExecutor(Executor): backend = "fake" diff --git a/tests/unit/torch/puzzletron/test_post_mip_reporting.py b/tests/unit/torch/puzzletron/test_post_mip_reporting.py new file mode 100644 index 00000000000..6fc4a926ead --- /dev/null +++ b/tests/unit/torch/puzzletron/test_post_mip_reporting.py @@ -0,0 +1,101 @@ +# 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. + +import json +from pathlib import Path +from types import SimpleNamespace + +from modelopt.torch.puzzletron.post_mip.reporting import ( + build_post_mip_report_payloads, + render_evaluation_report, +) + + +def _write_json(path: Path, payload: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload), encoding="utf-8") + + +def test_evaluation_report_compares_teacher_and_deduplicated_mip_origins(tmp_path: Path): + architecture_id = "architecture_shared" + revision_id = "revision_candidate" + execution_id = "post_mip_execution_test" + node_root = tmp_path / "artifacts/post_mip/nodes/quality" + _write_json( + tmp_path / "artifacts/post_mip/candidate_registry.json", + { + "architectures": { + architecture_id: { + "origins": [ + {"kind": "heterogeneous"}, + {"kind": "homogeneous"}, + ] + }, + "architecture_invalid": "invalid", + }, + "revisions": { + revision_id: { + "architecture_id": architecture_id, + "artifact": {"hidden_width": 1024, "kind": "heterogeneous"}, + }, + "failed_revision": {"architecture_id": "architecture_invalid"}, + }, + }, + ) + _write_json(node_root / "summary.json", {"execution_identity": execution_id}) + _write_json( + node_root / f"executions/{execution_id}/observations.json", + [ + { + "input_revision_id": revision_id, + "status": "success", + "metrics": { + "candidate.lm_loss": 0.7, + "candidate.token_accuracy_top_1": 0.8, + "candidate.token_accuracy_top_10": 0.95, + "reference.lm_loss": 0.5, + "reference.token_accuracy_top_1": 0.9, + "delta.lm_loss": 0.2, + }, + }, + { + "input_revision_id": "failed_revision", + "status": "failed", + "metrics": {"reference.lm_loss": None}, + "error": "evaluation failed", + }, + ], + ) + node = SimpleNamespace( + stage_id="post.quality", + node_id="quality", + flow_id="post", + node_type="evaluation", + ) + + payload = build_post_mip_report_payloads(tmp_path, (node,))[node.stage_id] + rendered = render_evaluation_report(str(payload["section_id"]), payload) + + assert payload["observations"][0]["origin_kinds"] == ["heterogeneous", "homogeneous"] + assert payload["observations"][1]["origin_kinds"] == [] + assert rendered.count("Reference checkpoint") == 1 + assert "Teacher" in rendered + assert "Heterogeneous" in rendered + assert "Homogeneous" in rendered + assert "Top-1 token accuracy" in rendered + assert "Top-10 token accuracy" in rendered + assert "0.95" in rendered + assert "evaluation failed" in rendered + assert rendered.count("Shared physical measurement (same architecture)") == 2 diff --git a/tests/unit/torch/puzzletron/test_qwen3p5_0p8b_full_vlm_smoke_plan.py b/tests/unit/torch/puzzletron/test_qwen3p5_0p8b_full_vlm_smoke_plan.py index 9e80ab073f9..3b7f3513c83 100644 --- a/tests/unit/torch/puzzletron/test_qwen3p5_0p8b_full_vlm_smoke_plan.py +++ b/tests/unit/torch/puzzletron/test_qwen3p5_0p8b_full_vlm_smoke_plan.py @@ -42,7 +42,6 @@ def _compile(monkeypatch, tmp_path: Path, experiment: Path, execution: Path): monkeypatch.setenv("PUZZLETRON_RUN_ROOT", str(tmp_path / experiment.stem)) - monkeypatch.setenv("PUZZLETRON_DATASET_PATH", str(tmp_path / "dataset")) monkeypatch.setenv("PUZZLETRON_DATASET_REVISION", "fixture-revision") monkeypatch.setenv("HF_HOME", str(tmp_path / "hf-home")) return compile_campaign_plan( @@ -60,8 +59,12 @@ def test_full_vlm_smoke_compiles_one_complete_bounded_lifecycle( stages = {stage.stage_id: stage for stage in plan.stages} config = plan.experiment_config nodes = config["post_mip"]["flows"]["params-90"]["nodes"] + serving_args = nodes["vlm_serving"]["config"]["topology"]["extra_vllm_args"] assert "tokenize_data" not in stages + assert config["dataset_path"] == str(tmp_path / "full_vlm_smoke/datasets/nemotron_vlm_v2") + assert config["prepare_dataset"]["output"] == config["dataset_path"] + assert serving_args[serving_args.index("--gdn-prefill-backend") + 1] == "triton" assert tuple(node for node in stages if node.startswith("post.")) == ( "post.params-90.image_eval", "post.params-90.best_vlm_loss", @@ -75,10 +78,16 @@ def test_full_vlm_smoke_compiles_one_complete_bounded_lifecycle( "post.params-90.best", ) assert config["prepare_dataset"]["num_samples"] == 8 + assert config["sort_sanity"]["max_abs_lm_loss_delta"] == 0.003 + assert config["sort_sanity"]["max_abs_reverse_lm_loss_delta"] == 0.003 assert nodes["image_eval"]["config"]["eval_samples"] == 2 assert nodes["checkpoint_eval"]["config"]["profile"] == "qwen35_vlm_core3_24row_smoke_v2" + assert nodes["checkpoint_eval"]["config"]["gdn_prefill_backend"] == "triton" assert nodes["post_kd_checkpoint_eval"]["config"] == nodes["checkpoint_eval"]["config"] assert nodes["vlm_serving"]["config"]["request_count"] == 1 + assert nodes["fastest_vlm"]["metric"] == ( + "vlm_serving.images_12.concurrency_1.image_throughput" + ) assert nodes["short_vlm_kd"]["config"]["max_steps"] == 2 cpu_stages = [stage for stage in stages.values() if stage.resource == "cpu"] assert cpu_stages @@ -90,11 +99,16 @@ def test_vlm_campaign_compiles_the_multi_axis_flow(monkeypatch, tmp_path: Path) plan = _compile(monkeypatch, tmp_path, CAMPAIGN_PATH, CAMPAIGN_EXECUTION_PATH) stages = {stage.stage_id: stage for stage in plan.stages} config = plan.experiment_config - assert set(config["post_mip"]["flows"]) == {"candidates"} + assert config["dataset_path"] == str(tmp_path / "vlm_campaign/datasets/nemotron_vlm_v2") + assert config["prepare_dataset"]["output"] == config["dataset_path"] assert config["prepare_dataset"]["evaluation_hf_home"] == str(tmp_path / "hf-home") + assert set(config["post_mip"]["flows"]) == {"candidates"} candidates = config["post_mip"]["flows"]["candidates"]["nodes"] + serving_args = candidates["serving"]["config"]["topology"]["extra_vllm_args"] + assert serving_args[serving_args.index("--gdn-prefill-backend") + 1] == "triton" assert set(config["mip"]["runs"]) == {"params-90"} + assert config["replacement_scoring"]["eval_samples"] == 16 assert candidates["best_image_loss"]["top_k"] == 5 assert candidates["kd"]["config"]["max_steps"] == 128 assert candidates["pre_kd_eval"]["config"] == config["vlm_quality_evaluation"] @@ -105,7 +119,8 @@ def test_vlm_campaign_compiles_the_multi_axis_flow(monkeypatch, tmp_path: Path) assert candidates["selected"]["input"] == "post_kd_eval" assert candidates["selected"]["top_k"] == 1 assert stages["post.candidates.serving"].parents == ("post.candidates.result",) - assert stages["post.candidates.kd"].total_gpus == 2 + assert stages["post.candidates.kd"].total_gpus == 1 + assert max(stage.total_gpus for stage in stages.values()) == 1 profile_rows = contracts.load_profile("core-3_344-examples_r1-vllm").exact_rows assert profile_rows is not None diff --git a/tests/unit/torch/puzzletron/test_setup_inspection.py b/tests/unit/torch/puzzletron/test_setup_inspection.py index b169ad51040..2fda8835702 100644 --- a/tests/unit/torch/puzzletron/test_setup_inspection.py +++ b/tests/unit/torch/puzzletron/test_setup_inspection.py @@ -1,11 +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. """Focused source-normalization tests for the Puzzletron setup wizard.""" +from types import SimpleNamespace + import pytest -from puzzletron_setup import SetupError +from puzzletron_setup import SetupError, inspection from puzzletron_setup.inspection import normalize_dataset_source, normalize_model_source @@ -13,12 +27,11 @@ def test_normalizes_hugging_face_web_urls(): assert normalize_model_source("https://huggingface.co/Qwen/Qwen3.5-0.8B") == ( "Qwen/Qwen3.5-0.8B" ) - assert normalize_model_source("huggingface.com/Qwen/Qwen3.5-0.8B/") == ( - "Qwen/Qwen3.5-0.8B" + assert normalize_model_source("huggingface.com/Qwen/Qwen3.5-0.8B/") == ("Qwen/Qwen3.5-0.8B") + assert ( + normalize_dataset_source("https://huggingface.com/datasets/nvidia/Some-Dataset") + == "nvidia/Some-Dataset" ) - assert normalize_dataset_source( - "https://huggingface.com/datasets/nvidia/Some-Dataset" - ) == "nvidia/Some-Dataset" def test_normalizes_existing_local_paths_and_rejects_other_uris(tmp_path): @@ -33,3 +46,71 @@ def test_normalizes_existing_local_paths_and_rejects_other_uris(tmp_path): normalize_model_source("s3://bucket/model") with pytest.raises(SetupError, match="does not exist"): normalize_dataset_source("../missing-dataset") + + +def test_inspect_model_uses_cached_config_when_hub_resolution_is_unavailable(tmp_path, monkeypatch): + revision = "a" * 40 + config_path = tmp_path / "models--Qwen--cached" / "snapshots" / revision / "config.json" + config_path.parent.mkdir(parents=True) + config_path.write_text('{"model_type": "cached"}\n') + + class OfflineApi: + def model_info(self, source, revision=None): + del source, revision + raise RuntimeError("offline") + + monkeypatch.setattr(inspection, "HfApi", OfflineApi) + monkeypatch.setattr( + inspection, + "try_to_load_from_cache", + lambda source, filename, revision: None, + ) + monkeypatch.setattr( + inspection, + "scan_cache_dir", + lambda: SimpleNamespace( + repos=[ + SimpleNamespace( + repo_type="model", + repo_id="Qwen/cached", + revisions=[ + SimpleNamespace( + snapshot_path=config_path.parent, + commit_hash=revision, + ) + ], + ) + ] + ), + ) + monkeypatch.setattr( + inspection, + "resolve_profile", + lambda config: SimpleNamespace(inventory=lambda value: "cached-inventory"), + ) + + model = inspection.inspect_model("Qwen/cached") + + assert model.source == "Qwen/cached" + assert model.resolved_revision == revision + assert model.config == {"model_type": "cached"} + assert model.inventory == "cached-inventory" + + +def test_cached_model_ref_preserves_snapshot_commit_across_blob_symlink(tmp_path, monkeypatch): + revision = "b" * 40 + blob = tmp_path / "blobs" / "config" + blob.parent.mkdir() + blob.write_text("{}\n") + config_path = tmp_path / "snapshots" / revision / "config.json" + config_path.parent.mkdir(parents=True) + config_path.symlink_to(blob) + monkeypatch.setattr( + inspection, + "try_to_load_from_cache", + lambda source, filename, revision: str(config_path), + ) + + cached = inspection._cached_remote_config("Qwen/cached", None) + + assert cached == (config_path.absolute(), revision) diff --git a/tests/unit/torch/puzzletron/test_setup_v2_data.py b/tests/unit/torch/puzzletron/test_setup_v2_data.py index 899581bdb6a..19bba98373e 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_data.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_data.py @@ -409,6 +409,9 @@ def test_bundle_readme_emits_bounded_vlm_materialization_command(tmp_path): assert "--revision sha" in document assert "--num-samples 64" in document assert "--max-shards-per-subset 2" in document + assert "Python 3.10+ controller venv" in document + assert "This is a worker command" in document + assert "vlm_checkpoint_evaluation.md" in document def test_checkbox_rejects_a_disabled_scripted_selection(tmp_path): diff --git a/tests/unit/torch/puzzletron/test_setup_v2_post_mip.py b/tests/unit/torch/puzzletron/test_setup_v2_post_mip.py index 64c7c479a30..ac9020b106f 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_post_mip.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_post_mip.py @@ -79,6 +79,7 @@ def test_recommended_multimodal_flow_uses_image_serving_and_vlm_selection(): "concurrency": [1, 4], "request_count": 32, "image_batch_sizes": [1, 6, 12], + "topology": {"extra_vllm_args": ["--gdn-prefill-backend", "triton"]}, }, modality="multimodal", quality_comparison=comparison, @@ -97,8 +98,12 @@ def test_recommended_multimodal_flow_uses_image_serving_and_vlm_selection(): ) assert flow.nodes["vlm_serving"].config["endpoint_type"] == "chat" assert flow.nodes["vlm_serving"].config["image_batch_sizes"] == [1, 6, 12] + assert flow.nodes["vlm_serving"].config["topology"]["extra_vllm_args"] == [ + "--gdn-prefill-backend", + "triton", + ] assert flow.nodes["fastest_vlm"].selector["metric"] == ( - "vlm_serving.images_12.concurrency_1.image_throughput" + "vlm_serving.images_12.image_throughput" ) assert flow.nodes["quality_benchmarks"].input_id == "best" assert flow.nodes["quality_benchmarks"].config == comparison diff --git a/tests/unit/torch/puzzletron/test_setup_v2_quick.py b/tests/unit/torch/puzzletron/test_setup_v2_quick.py index af913b9c0d9..cd6e6894c53 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_quick.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_quick.py @@ -17,11 +17,13 @@ from __future__ import annotations +from pathlib import Path from types import SimpleNamespace import pytest import yaml +import puzzletron_setup.v2.bundle as bundle_module import puzzletron_setup.v2.cli as cli_module import puzzletron_setup.v2.wizard as wizard_module from puzzletron_orchestrator.compiler import ( @@ -469,7 +471,36 @@ def test_guided_wizard_runs_real_sections_and_generates_valid_bundles( ) smoke = yaml.safe_load((campaign / "smoke" / "experiment.yaml").read_text()) smoke_flow = next(iter(smoke["post_mip"]["flows"].values())) + assert ( + smoke_flow["nodes"]["serving"]["config"]["topology"]["server_context_overhead_tokens"] + == 16384 + ) + assert smoke_flow["nodes"]["serving"]["config"]["topology"]["extra_vllm_args"] == [ + "-cc.cudagraph_mode=NONE", + "--no-enable-flashinfer-autotune", + "--gdn-prefill-backend", + "triton", + "--gpu-memory-utilization", + "0.5", + "--reasoning-parser", + "qwen3", + "--default-chat-template-kwargs", + '{"enable_thinking": false}', + ] smoke_comparison = smoke_flow["nodes"]["quality_benchmarks"] + smoke_nodes = smoke_flow["nodes"] + smoke_mip_run = next(iter(smoke["mip"]["runs"].values())) + assert smoke_mip_run["solver"]["num_solutions"] == 2 + assert smoke_mip_run["homogeneous"]["keep"] == 2 + assert smoke_nodes["online_eval"]["config"]["eval_samples"] == 2 + assert smoke_nodes["online_eval"]["config"]["block_size"] == 32 + assert smoke_nodes["serving"]["config"]["input_tokens"] == 100 + assert smoke_nodes["serving"]["config"]["output_tokens"] == 80 + assert smoke_nodes["serving"]["config"]["request_count"] == 1 + assert smoke_nodes["short_kd"]["config"]["max_steps"] == 2 + assert smoke_nodes["short_kd"]["config"]["global_batch_size"] == 1 + assert smoke_nodes["short_kd"]["config"]["checkpoint_every_steps"] == 2 + assert smoke_comparison["config"]["limit"] == 8 assert "recorded_observation" not in smoke_comparison["config"] smoke_runner = yaml.safe_load((campaign / "smoke" / "runner.yaml").read_text()) assert smoke_runner["runner"]["slurm"]["job_name_prefix"] == "acct-puzzletron" @@ -570,7 +601,7 @@ def test_guided_wizard_generates_the_complete_qwen_vlm_flow(tmp_path, monkeypatc "sequence_length": 512, }, "infrastructure": { - "gpus_per_node": 1, + "gpus_per_node": 8, "execution_contract": { "repository": "/worker/modelopt", "venv": "/worker/venv", @@ -586,19 +617,23 @@ def test_guided_wizard_generates_the_complete_qwen_vlm_flow(tmp_path, monkeypatc defaults_path=defaults, backend=NonInteractiveBackend(), campaign_dir=campaign, - setup_profile="smoke", + setup_profile="balanced", ) assert result == campaign.resolve() smoke = yaml.safe_load((campaign / "smoke" / "experiment.yaml").read_text()) + smoke_execution = yaml.safe_load((campaign / "smoke" / "execution.yaml").read_text()) smoke_flow = next(iter(smoke["post_mip"]["flows"].values())) smoke_quality = smoke_flow["nodes"]["quality_benchmarks"]["config"] assert smoke_quality["profile"] == "qwen35_vlm_realworldqa64_mmmu120_mvbench160_frozen_rows_v3" - assert smoke_quality["limit"] == 8 + assert "limit" not in smoke_quality assert smoke_quality["limit_mm_per_prompt"] == {"image": 32} assert smoke_quality["max_model_len"] == 32768 assert "recorded_observation" not in smoke_quality + assert smoke_execution["execution"]["stages"]["replacement_scoring"]["instances"] == 1 + assert smoke_execution["execution"]["stages"]["post.params-90.short_kd"]["instances"] == 1 production = yaml.safe_load((campaign / "production" / "experiment.yaml").read_text()) + production_execution = yaml.safe_load((campaign / "production" / "execution.yaml").read_text()) flow_id, flow = next(iter(production["post_mip"]["flows"].items())) assert flow_id == "params-90" nodes = flow["nodes"] @@ -611,6 +646,8 @@ def test_guided_wizard_generates_the_complete_qwen_vlm_flow(tmp_path, monkeypatc assert "model" not in quality["config"] assert "log_samples" not in quality["config"] assert "recorded_observation" not in quality["config"] + assert production_execution["execution"]["stages"]["replacement_scoring"]["instances"] == 8 + assert production_execution["execution"]["stages"]["post.params-90.short_kd"]["instances"] == 8 assert production["global_distillation"]["domain"] == "vlm" assert production["global_distillation"]["freeze_policy"] == "train_all" assert production["tokenize_data"]["enabled"] is False @@ -628,7 +665,70 @@ def test_guided_wizard_generates_the_complete_qwen_vlm_flow(tmp_path, monkeypatc assert f"post.{flow_id}.vlm_serving" in stage_ids assert f"post.{flow_id}.short_kd" in stage_ids assert stage_ids[-1] == f"post.{flow_id}.quality_benchmarks" - assert all(stage.total_gpus <= 1 for stage in plan.stages) + mip = next(stage for stage in plan.stages if stage.stage_id == "mip") + assert mip.resource == "cpu" + assert mip.total_gpus == 0 + dry_run = (campaign / "production" / "dry-run-plan.txt").read_text() + assert "mip: 1 submission(s), strategy=single, resource=cpu" in dry_run + assert '"resource": "cpu"' in dry_run + assert "scheduler_script" in dry_run + assert str(campaign / "production" / "experiment.yaml") in dry_run + assert ".puzzletron-v2-" not in dry_run + readme = (campaign / "README.md").read_text() + assert "generation-time snapshot" in readme + snapshots = { + budget: (campaign / budget / "dry-run-plan.txt").read_text() + for budget in ("smoke", "production") + } + wizard_module.build_bundles_v2(campaign, WizardState.resume(campaign)) + assert snapshots == { + budget: (campaign / budget / "dry-run-plan.txt").read_text() + for budget in ("smoke", "production") + } + + changed_state = WizardState.resume(campaign) + changed_state.set_field( + "infrastructure.runner.slurm.partition", + "replacement-gpu-partition", + source="user", + ) + readme_path = campaign / "README.md" + readme_path.write_text(f"{readme_path.read_text()}\nrollback sentinel\n") + + def campaign_files(): + return { + path.relative_to(campaign): path.read_bytes() + for path in campaign.rglob("*") + if path.is_file() + } + + published_files = campaign_files() + assert b"replacement-gpu-partition" not in published_files[Path("smoke/runner.yaml")] + + render_plan = bundle_module.dry_run_bundle + + def fail_production_plan(bundle): + if bundle.parent == campaign and bundle.name == "production": + raise RuntimeError("dry-run rendering failed") + return render_plan(bundle) + + monkeypatch.setattr(bundle_module, "dry_run_bundle", fail_production_plan) + with pytest.raises(RuntimeError, match="dry-run rendering failed"): + wizard_module.build_bundles_v2(campaign, changed_state) + assert campaign_files() == published_files + + monkeypatch.setattr(bundle_module, "dry_run_bundle", render_plan) + replace = bundle_module.os.replace + + def fail_readme_publish(source, target): + if target == campaign / "README.md" and source.name == "README.md": + raise RuntimeError("README publication failed") + return replace(source, target) + + monkeypatch.setattr(bundle_module.os, "replace", fail_readme_publish) + with pytest.raises(RuntimeError, match="README publication failed"): + wizard_module.build_bundles_v2(campaign, changed_state) + assert campaign_files() == published_files # Interactive prompt navigation From 4b871ee75a32ad55cf1e6c989d749e68400ddb99 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Tue, 8 Sep 2026 00:40:08 +0200 Subject: [PATCH 2/4] Clarify Puzzletron example expectations Signed-off-by: Johannes Rausch --- examples/puzzletron/README.md | 6 ++++-- examples/puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md | 10 ++++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 2f14ff0bac9..c3cdbee5712 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -20,7 +20,9 @@ Start with the maintained Qwen 3.5 0.8B VLM smoke. It exercises the complete image-text lifecycle with bounded workloads and at most one GPU per stage. Once it succeeds, the optional longer example campaign uses the same environment, runner, orchestrator, progress display, report, and resume command. That larger -campaign is a scheduled, multi-hour example, not a routine smoke or presubmit. +campaign takes longer than the smoke; its duration depends on worker hardware, +scheduler availability, cache state, and the execution profile. It is intended +for scheduled integration validation, not routine smoke or presubmit use. ### 1. Create the controller environment @@ -138,7 +140,7 @@ report. For a failed or interrupted run, follow the actionable checks in The [Qwen VLM example guide](docs/qwen3p5_0p8b_vlm_smoke.md) lists the smoke's expected lifecycle checks and explains how to interpret its bounded results. -### 6. Optional: run the longer multi-hour example +### 6. Optional: run the longer example After the smoke succeeds, keep the controller venv and runner and select a new run root plus the longer example files: diff --git a/examples/puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md b/examples/puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md index a1db391e493..6265e2a11eb 100644 --- a/examples/puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md +++ b/examples/puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md @@ -5,7 +5,13 @@ The Qwen 3.5 0.8B VLM example has two experiment files with distinct jobs: | Experiment | Purpose | Execution profile | | --- | --- | --- | | `full_vlm_smoke.yaml` | Check the complete lifecycle on one GPU | `execution.single_gpu.yaml` | -| `vlm_campaign.yaml` | Run a scheduled, multi-hour multi-axis example | `qwen3p5_0p8b/execution.vlm_campaign.yaml` | +| `vlm_campaign.yaml` | Run the longer scheduled multi-axis example | `qwen3p5_0p8b/execution.vlm_campaign.yaml` | + +Both recipes select vLLM's Triton GDN prefill backend. On a fresh worker, the +FlashInfer GDN kernels can still be compiling when the server readiness check +expires. Selecting Triton avoids that cold-start failure and makes startup +predictable in the reviewed worker image; it is not a general performance +recommendation. Start with `full_vlm_smoke.yaml`. It uses small workloads to check dataset preparation, pruning, MIP, materialization, checkpoint evaluation, serving, @@ -99,7 +105,7 @@ After completion, inspect the campaign report and verify that: The smoke deliberately uses tiny workloads. Run a separate benchmark with representative requests before drawing performance conclusions. -## Run the longer multi-hour illustrative campaign +## Run the longer illustrative campaign The campaign enables hidden width (the model's shared transformer hidden size), heterogeneous FFN width (per-block feed-forward intermediate sizes), depth, From 474d8da0a068c4f7feccbda50f851142b09b3d23 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Tue, 8 Sep 2026 01:08:19 +0200 Subject: [PATCH 3/4] Fix smoke rendering and tighten tests Signed-off-by: Johannes Rausch --- puzzletron_setup/v2/bundle.py | 3 +- .../vlm/preparation/test_benchmark_data.py | 19 ++-- .../test_automodel_solution_scoring.py | 91 ++++++++++++------- .../torch/puzzletron/test_compact_runtime.py | 70 ++++++++------ .../test_hidden_width_diagnostic.py | 47 ++++------ .../test_orchestration_shutdown_progress.py | 27 ------ .../puzzletron/test_post_mip_reporting.py | 6 +- .../torch/puzzletron/test_setup_inspection.py | 49 ++++++++-- .../torch/puzzletron/test_setup_v2_data.py | 41 ++++----- .../torch/puzzletron/test_setup_v2_quick.py | 22 +---- .../test_setup_v2_resolved_config.py | 64 +++++++++++++ 11 files changed, 258 insertions(+), 181 deletions(-) diff --git a/puzzletron_setup/v2/bundle.py b/puzzletron_setup/v2/bundle.py index 99ff2e262f0..ef93b1e1760 100644 --- a/puzzletron_setup/v2/bundle.py +++ b/puzzletron_setup/v2/bundle.py @@ -239,7 +239,8 @@ def _render_experiment_v2( run["solver"] = solver homogeneous = _mapping(run.get("homogeneous")) if homogeneous: - homogeneous["keep"] = min(int(homogeneous.get("keep", 2) or 2), 2) + keep = homogeneous.get("keep", 2) or 2 + homogeneous["keep"] = 2 if keep == "all" else min(int(keep), 2) run["homogeneous"] = homogeneous flows = _plain(config.post_mip_flows) if config.post_mip_flows_configured: diff --git a/tests/unit/torch/puzzletron/evaluation/vlm/preparation/test_benchmark_data.py b/tests/unit/torch/puzzletron/evaluation/vlm/preparation/test_benchmark_data.py index 28d5003bae0..16eea88f38e 100644 --- a/tests/unit/torch/puzzletron/evaluation/vlm/preparation/test_benchmark_data.py +++ b/tests/unit/torch/puzzletron/evaluation/vlm/preparation/test_benchmark_data.py @@ -168,9 +168,11 @@ def prepare(root, task, snapshot, *, verify_content=False): assert prepared == ["mvbench"] -def test_prepare_benchmark_datasets_rejects_symlinked_hf_home(tmp_path): +@pytest.mark.parametrize("dangling", [False, True]) +def test_prepare_benchmark_datasets_rejects_symlinked_hf_home(tmp_path, dangling): target = tmp_path / "target" - target.mkdir() + if not dangling: + target.mkdir() alias = tmp_path / "hf-home" alias.symlink_to(target, target_is_directory=True) @@ -178,14 +180,6 @@ def test_prepare_benchmark_datasets_rejects_symlinked_hf_home(tmp_path): preparation.prepare_benchmark_datasets(alias, ("realworldqa",)) -def test_prepare_benchmark_datasets_rejects_dangling_symlinked_hf_home(tmp_path): - alias = tmp_path / "hf-home" - alias.symlink_to(tmp_path / "missing", target_is_directory=True) - - with pytest.raises(ValueError, match="must not be a symlink"): - preparation.prepare_benchmark_datasets(alias, ("realworldqa",)) - - def test_zip_preparation_is_revision_bound_idempotent_and_byte_verified(tmp_path): hf_home = tmp_path / "hf-home" snapshot = preparation._hub_snapshot(hf_home, "mmvu_val") @@ -244,6 +238,8 @@ def test_prepared_media_reuse_tolerates_distributed_filesystem_mtime_skew(tmp_pa inventory[0]["mtime_ns"] += 1_000_000_000 assert preparation._inventory_is_current(root, inventory) + inventory[0]["sha256"] = "0" * 64 + assert not preparation._inventory_is_current(root, inventory) def test_prepared_media_reuse_reports_paths_from_the_current_mount(tmp_path): @@ -590,7 +586,8 @@ def run(*, second=False): second.join(timeout=5) assert not errors - assert not first.is_alive() and not second.is_alive() + assert not first.is_alive() + assert not second.is_alive() assert extraction_count == 1 assert len(results) == 2 assert all(result["status"] == "complete" for result in results) diff --git a/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py b/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py index 271d63bdbab..320836707ff 100644 --- a/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py +++ b/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py @@ -21,7 +21,6 @@ from contextlib import contextmanager from types import SimpleNamespace -import pandas as pd import pytest import torch from omegaconf import OmegaConf @@ -92,7 +91,7 @@ def test_solution_work_reuses_only_matching_config_and_identity(monkeypatch, tmp assert pending_ids == [1, 2] - assert find_missing_solutions(pd.DataFrame(solutions), tmp_path, scoring) == [1, 2] + assert find_missing_solutions(solutions, tmp_path, scoring) == [1, 2] parent_identity = {"parent_role": "sorted", "checkpoint_dir": "/checkpoints/sorted"} (tmp_path / "parent.json").write_text(json.dumps({"args": recorded, **parent_identity})) @@ -106,7 +105,9 @@ def test_solution_work_reuses_only_matching_config_and_identity(monkeypatch, tmp ) -def test_native_automodel_gdn_supports_single_device_compact_runtime_candidate(): +def test_native_automodel_gdn_executes_compact_geometry_and_rejects_context_parallelism( + monkeypatch, +): # Optional dependency: native AutoModel Qwen modules are not installed in every test env. pytest.importorskip("nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn") from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import CPAwareGatedDeltaNet @@ -156,21 +157,67 @@ def test_native_automodel_gdn_supports_single_device_compact_runtime_candidate() ) ) original_forward = gdn.forward.__func__ - original_state = set(vars(gdn)) original_forward_hooks = dict(gdn._forward_hooks) hidden_states = torch.randn(1, 4, config.hidden_size) + conv_shapes = [] + kernel_shapes = [] + + def compact_conv(*, x, weight, bias, **kwargs): + conv_shapes.append((x.shape, weight.shape, bias.shape)) + return x + + def compact_kernel(query, key, value, *, g, beta, **kwargs): + kernel_shapes.append( + { + "query": query.shape, + "key": key.shape, + "value": value.shape, + "gate": g.shape, + "beta": beta.shape, + } + ) + return value, None + + monkeypatch.setattr(gdn, "causal_conv1d_fn", compact_conv) + monkeypatch.setattr(gdn, "chunk_gated_delta_rule", compact_kernel) + original_state = set(vars(gdn)) handle = apply_runtime_candidate(layer, teacher, child) - assert "forward" in vars(gdn) - with torch.no_grad(): - output = gdn(hidden_states, qkv_format="bshd") - assert output.shape == hidden_states.shape - handle.remove() + try: + with torch.no_grad(): + output = gdn(hidden_states, qkv_format="bshd") + finally: + handle.remove() + target_heads = shape.num_value_heads // 2 + target_projection_width = ( + 2 * (shape.num_key_heads // 2) * shape.key_head_dim + target_heads * shape.value_head_dim + ) + assert conv_shapes == [ + ( + torch.Size((1, target_projection_width, 4)), + torch.Size((target_projection_width, config.linear_conv_kernel_dim)), + torch.Size((target_projection_width,)), + ) + ] + assert kernel_shapes == [ + { + "query": torch.Size((1, 4, target_heads, shape.key_head_dim)), + "key": torch.Size((1, 4, target_heads, shape.key_head_dim)), + "value": torch.Size((1, 4, target_heads, shape.value_head_dim)), + "gate": torch.Size((1, 4, target_heads)), + "beta": torch.Size((1, 4, target_heads)), + } + ] + assert output.shape == hidden_states.shape assert gdn.forward.__func__ is original_forward assert set(vars(gdn)) == original_state assert dict(gdn._forward_hooks) == original_forward_hooks + monkeypatch.setattr(gdn, "_cp_mesh", SimpleNamespace(size=lambda: 2)) + with pytest.raises(RuntimeError, match="refusing to score reduced geometry"): + apply_runtime_candidate(layer, teacher, child) + def test_block_checkpoint_overlay_merges_split_hf_experts_and_restores(tmp_path): class GroupedExperts(nn.Module): @@ -579,26 +626,6 @@ def test_runtime_mla_prefix_slice_matches_physical_projection_and_restores() -> torch.testing.assert_close(got, expected) -class _MLAHeadsProjection(nn.Module): - def __init__(self, num_heads: int, v_head_dim: int, hidden: int): - super().__init__() - self.o_proj = nn.Linear( - num_heads * v_head_dim, - hidden, - bias=False, - dtype=torch.float64, - ) - - def forward(self, head_outputs): - return self.o_proj(head_outputs) - - -class _MLAHeadsLayer(nn.Module): - def __init__(self): - super().__init__() - self.self_attn = _MLAHeadsProjection(num_heads=4, v_head_dim=3, hidden=5) - - class _GroupedRMSNormMamba(nn.Module): """Small fused-Mamba analogue whose norm retains the static teacher shape.""" @@ -648,12 +675,6 @@ def forward(self, x): ) -class _MambaLayer(nn.Module): - def __init__(self): - super().__init__() - self.mixer = _GroupedRMSNormMamba() - - class _NativeFusedMamba(_GroupedRMSNormMamba): """Native-shape harness that imports the fused function inside forward.""" diff --git a/tests/unit/torch/puzzletron/test_compact_runtime.py b/tests/unit/torch/puzzletron/test_compact_runtime.py index ce7662703cb..b4e480813ca 100644 --- a/tests/unit/torch/puzzletron/test_compact_runtime.py +++ b/tests/unit/torch/puzzletron/test_compact_runtime.py @@ -32,7 +32,6 @@ compact_grouped_attention_forward, resolve_compact_grouped_attention_target, supports_compact_gated_delta_net, - supports_compact_grouped_attention, ) from modelopt.torch.puzzletron.pruning.gated_delta_net import ( GDNShape, @@ -100,7 +99,7 @@ def test_compact_grouped_attention_target_requires_reduced_supported_geometry(): assert resolve_compact_grouped_attention_target(layer, teacher, teacher) is None -def test_compact_grouped_attention_dispatches_native_automodel_sdpa_backend(): +def test_compact_grouped_attention_executes_native_sdpa_and_rejects_other_backends(): # Optional dependency: native AutoModel Qwen modules are not installed in every test env. pytest.importorskip("nemo_automodel.components.models.qwen3_next.layers") from nemo_automodel.components.models.common import BackendConfig @@ -140,27 +139,46 @@ def test_compact_grouped_attention_dispatches_native_automodel_sdpa_backend(): num_query_heads=2, num_kv_heads=1, ) - original_forward = attention.forward.__func__ - original_state = set(vars(attention)) - - assert supports_compact_grouped_attention( - attention, - orig_num_q=4, - orig_num_kv=2, - head_dim=8, - ) target = resolve_compact_grouped_attention_target(layer, teacher, child) - assert target == { - "module": attention, - "orig_num_q": 4, - "orig_num_kv": 2, - "target_num_q": 2, - "target_num_kv": 1, - "head_dim": 8, - } - - assert attention.forward.__func__ is original_forward - assert set(vars(attention)) == original_state + assert target is not None + attention = target.pop("module") + kernel_shapes = [] + + def compact_attention(query, key, value, **kwargs): + kernel_shapes.append((query.shape, key.shape, value.shape)) + return torch.zeros_like(query) + + attention.attn_func = compact_attention + dtype = attention.q_proj.weight.dtype + hidden_states = torch.randn(1, 4, config.hidden_size, dtype=dtype) + freqs_cis = torch.cat( + ( + torch.ones(1, 4, config.head_dim // 2, dtype=dtype), + torch.zeros(1, 4, config.head_dim // 2, dtype=dtype), + ), + dim=-1, + ) + with compact_grouped_attention_forward(attention, **target): + output = attention(hidden_states, freqs_cis=freqs_cis) + restored_output = attention(hidden_states, freqs_cis=freqs_cis) + + assert kernel_shapes == [ + ( + torch.Size((1, 2, 4, 8)), + torch.Size((1, 1, 4, 8)), + torch.Size((1, 1, 4, 8)), + ), + ( + torch.Size((1, 4, 4, 8)), + torch.Size((1, 2, 4, 8)), + torch.Size((1, 2, 4, 8)), + ), + ] + assert output.shape == restored_output.shape == hidden_states.shape + + attention.backend.attn = "te" + with pytest.raises(RuntimeError, match="refusing to score reduced geometry"): + resolve_compact_grouped_attention_target(layer, teacher, child) @pytest.mark.parametrize(("target_num_q", "target_num_kv"), [(2, 1), (2, 2)]) @@ -213,8 +231,6 @@ def test_compact_grouped_attention_matches_physical_projection_geometry( ) attention_mask = torch.zeros(1, 1, 4, 4, dtype=dtype) original_shapes = {name: tuple(tensor.shape) for name, tensor in teacher.state_dict().items()} - projection_modules = (teacher.q_proj, teacher.k_proj, teacher.v_proj, teacher.o_proj) - assert all("forward" not in vars(module) for module in projection_modules) with torch.no_grad(): teacher_output = teacher( @@ -235,7 +251,6 @@ def test_compact_grouped_attention_matches_physical_projection_geometry( target_num_kv=target_num_kv, head_dim=8, ): - assert all("forward" in vars(module) for module in projection_modules) runtime_output = teacher( hidden_states, position_embeddings, @@ -297,7 +312,6 @@ def test_compact_grouped_attention_matches_physical_projection_geometry( original_shapes ) assert teacher.num_key_value_groups == 2 - assert all("forward" not in vars(module) for module in projection_modules) @pytest.mark.parametrize( @@ -341,8 +355,6 @@ def test_compact_gdn_matches_physical_projection_and_kernel_geometry( hidden_states = torch.randn(2, 4, teacher_config.hidden_size, dtype=dtype) attention_mask = torch.tensor([[1, 1, 1, 1], [1, 1, 0, 0]]) original_shapes = {name: tuple(tensor.shape) for name, tensor in teacher.state_dict().items()} - assert supports_compact_gated_delta_net(teacher, teacher_shape=teacher_shape) - assert "forward" not in vars(teacher) with torch.no_grad(): teacher_output = teacher(hidden_states, attention_mask=attention_mask) @@ -352,7 +364,6 @@ def test_compact_gdn_matches_physical_projection_and_kernel_geometry( teacher_shape=teacher_shape, target_shape=target_shape, ): - assert "forward" in vars(teacher) runtime_output = teacher(hidden_states, attention_mask=attention_mask) physical_cache = DynamicCache(config=target_config) @@ -385,7 +396,6 @@ def test_compact_gdn_matches_physical_projection_and_kernel_geometry( original_shapes ) assert GDNShape.from_module(teacher) == teacher_shape - assert "forward" not in vars(teacher) @pytest.mark.parametrize( diff --git a/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py b/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py index 6a7b98c0ee7..fefddbaf4fa 100644 --- a/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py +++ b/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py @@ -22,7 +22,6 @@ _hidden_only_diagnostic_ready, _hidden_width_ranking_verdict, _hidden_width_result_metrics, - _merge_reused_sort_equivalence, _parent_sweep_sanity_verdict, _ratio_aligned_hidden_widths, _resolve_parent_sweep_sort_equivalence, @@ -193,27 +192,6 @@ def test_hidden_width_diagnostic_preserves_all_available_solution_metrics(): assert all(metrics[name] == raw[name]["avg"] for name in metric_names) -def test_reused_parent_sweep_preserves_existing_sort_diagnosis_metrics(): - existing = { - "passed": True, - "teacher": {"lm_loss": 1.2}, - "sorted_teacher": {"lm_loss": 1.2001}, - "reverse_sorted": {"lm_loss": 1.5}, - } - reuse = { - "passed": True, - "reused_parent_sweep": True, - "equivalence": {"passed": True}, - } - - merged = _merge_reused_sort_equivalence(existing, reuse) - - assert merged["teacher"] == existing["teacher"] - assert merged["sorted_teacher"] == existing["sorted_teacher"] - assert merged["reverse_sorted"] == existing["reverse_sorted"] - assert merged["reused_parent_sweep"] is True - - def test_reused_parent_sweep_does_not_mutate_completed_sort_artifact(tmp_path): sort_summary_path = tmp_path / "sort_sanity" / "summary.json" reuse_summary_path = tmp_path / "width_sanity" / "reused_sort_equivalence.json" @@ -222,6 +200,7 @@ def test_reused_parent_sweep_does_not_mutate_completed_sort_artifact(tmp_path): "passed": True, "teacher": {"lm_loss": 1.2}, "sorted_teacher": {"lm_loss": 1.2001}, + "reverse_sorted": {"lm_loss": 1.5}, } sort_summary_path.write_text(json.dumps(original, indent=2, sort_keys=True) + "\n") original_bytes = sort_summary_path.read_bytes() @@ -234,8 +213,7 @@ def test_reused_parent_sweep_does_not_mutate_completed_sort_artifact(tmp_path): assert sort_summary_path.read_bytes() == original_bytes assert json.loads(reuse_summary_path.read_text()) == merged - assert merged["teacher"] == original["teacher"] - assert merged["reused_parent_sweep"] is True + assert merged == {**original, "reused_parent_sweep": True} def test_parent_sweep_reuses_canonical_sort_verdict_when_equivalence_was_skipped(tmp_path): @@ -253,15 +231,28 @@ def test_parent_sweep_reuses_canonical_sort_verdict_when_equivalence_was_skipped assert resolved["passed"] is True assert resolved["reused_source_summary"] == str(sort_summary_path) - assert json.loads(reuse_summary_path.read_text()) == resolved -@pytest.mark.parametrize("invalid_summary", [None, [], "not-an-object"]) -def test_reused_parent_sweep_rejects_non_object_sort_artifacts(tmp_path, invalid_summary): +def test_parent_sweep_keeps_manifest_equivalence_when_sort_reuse_is_disabled(tmp_path): + parent_equivalence = {"passed": False, "findings": [{"message": "parent mismatch"}]} + reuse_summary_path = tmp_path / "reused_sort_equivalence.json" + + resolved = _resolve_parent_sweep_sort_equivalence( + parent_equivalence=parent_equivalence, + sort_summary_path=tmp_path / "missing_sort_summary.json", + reuse_summary_path=reuse_summary_path, + reuse_sort_equivalence=False, + ) + + assert resolved == parent_equivalence + assert not reuse_summary_path.exists() + + +def test_reused_parent_sweep_rejects_non_object_sort_artifacts(tmp_path): sort_summary_path = tmp_path / "sort_sanity" / "summary.json" reuse_summary_path = tmp_path / "width_sanity" / "reused_sort_equivalence.json" sort_summary_path.parent.mkdir() - sort_summary_path.write_text(json.dumps(invalid_summary)) + sort_summary_path.write_text("[]") with pytest.raises(ValueError, match="expected a JSON object"): _write_reused_sort_equivalence( diff --git a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py index c76b142d514..ae01bc14c95 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py +++ b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py @@ -330,18 +330,6 @@ def _stage_id(handle: JobHandle) -> str: return str(handle.metadata.get("work_id", "")).split(":", 1)[0] -def _blocked_descendants(plan, failed_stages: set[str]) -> set[str]: - blocked = set(failed_stages) - changed = True - while changed: - changed = False - for node in plan.stages: - if node.stage_id not in blocked and any(parent in blocked for parent in node.parents): - blocked.add(node.stage_id) - changed = True - return blocked - - def _compile_test_plan( tmp_path: Path, *, @@ -952,18 +940,6 @@ def poll(self, handles: Sequence[JobHandle]) -> list[JobStatus]: and attempt["work_id"].split(":", 1)[0] != failed_stage for attempt in controller.store.list_attempts() ) - blocked = {failed_stage} - changed = True - while changed: - changed = False - for node in plan.stages: - if node.stage_id not in blocked and any(parent in blocked for parent in node.parents): - blocked.add(node.stage_id) - changed = True - attempted_stages = { - attempt["work_id"].split(":", 1)[0] for attempt in controller.store.list_attempts() - } - assert not ((blocked - {failed_stage}) & attempted_stages) def test_controller_fatal_failure_cancels_other_jobs_in_fail_fast_mode( @@ -1066,9 +1042,6 @@ def poll(self, handles: Sequence[JobHandle]) -> list[JobStatus]: assert result["failed_stages"] == ["sort_sanity"] assert "bypass_sanity" in executor.submitted_stage_ids assert "width_sanity" not in executor.submitted_stage_ids - blocked = _blocked_descendants(plan, {"sort_sanity"}) - assert "width_sanity" in blocked - assert "bypass_sanity" not in blocked def test_controller_sigint_during_recovery_cancels_jobs(tmp_path: Path): diff --git a/tests/unit/torch/puzzletron/test_post_mip_reporting.py b/tests/unit/torch/puzzletron/test_post_mip_reporting.py index 6fc4a926ead..836cd94de04 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_reporting.py +++ b/tests/unit/torch/puzzletron/test_post_mip_reporting.py @@ -91,11 +91,13 @@ def test_evaluation_report_compares_teacher_and_deduplicated_mip_origins(tmp_pat assert payload["observations"][0]["origin_kinds"] == ["heterogeneous", "homogeneous"] assert payload["observations"][1]["origin_kinds"] == [] assert rendered.count("Reference checkpoint") == 1 - assert "Teacher" in rendered + assert rendered.count("Teacher") == 1 assert "Heterogeneous" in rendered assert "Homogeneous" in rendered assert "Top-1 token accuracy" in rendered assert "Top-10 token accuracy" in rendered - assert "0.95" in rendered + assert rendered.count("0.8") == 2 + assert rendered.count("0.9") == 1 + assert rendered.count("0.95") == 2 assert "evaluation failed" in rendered assert rendered.count("Shared physical measurement (same architecture)") == 2 diff --git a/tests/unit/torch/puzzletron/test_setup_inspection.py b/tests/unit/torch/puzzletron/test_setup_inspection.py index 2fda8835702..197cf1f867c 100644 --- a/tests/unit/torch/puzzletron/test_setup_inspection.py +++ b/tests/unit/torch/puzzletron/test_setup_inspection.py @@ -23,6 +23,12 @@ from puzzletron_setup.inspection import normalize_dataset_source, normalize_model_source +class _OfflineApi: + def model_info(self, source, revision=None): + del source, revision + raise RuntimeError("offline") + + def test_normalizes_hugging_face_web_urls(): assert normalize_model_source("https://huggingface.co/Qwen/Qwen3.5-0.8B") == ( "Qwen/Qwen3.5-0.8B" @@ -54,12 +60,7 @@ def test_inspect_model_uses_cached_config_when_hub_resolution_is_unavailable(tmp config_path.parent.mkdir(parents=True) config_path.write_text('{"model_type": "cached"}\n') - class OfflineApi: - def model_info(self, source, revision=None): - del source, revision - raise RuntimeError("offline") - - monkeypatch.setattr(inspection, "HfApi", OfflineApi) + monkeypatch.setattr(inspection, "HfApi", _OfflineApi) monkeypatch.setattr( inspection, "try_to_load_from_cache", @@ -97,6 +98,42 @@ def model_info(self, source, revision=None): assert model.inventory == "cached-inventory" +def test_inspect_model_rejects_ambiguous_cached_revisions(tmp_path, monkeypatch): + revisions = ("a" * 40, "b" * 40) + snapshots = [] + for revision in revisions: + snapshot = tmp_path / "snapshots" / revision + snapshot.mkdir(parents=True) + (snapshot / "config.json").write_text("{}\n") + snapshots.append(SimpleNamespace(snapshot_path=snapshot, commit_hash=revision)) + + monkeypatch.setattr(inspection, "HfApi", _OfflineApi) + monkeypatch.setattr(inspection, "try_to_load_from_cache", lambda *args, **kwargs: None) + monkeypatch.setattr( + inspection, + "scan_cache_dir", + lambda: SimpleNamespace( + repos=[SimpleNamespace(repo_type="model", repo_id="Qwen/cached", revisions=snapshots)] + ), + ) + + with pytest.raises(SetupError, match="Cannot resolve Hugging Face model"): + inspection.inspect_model("Qwen/cached") + + +def test_inspect_model_does_not_substitute_another_cached_revision(monkeypatch): + monkeypatch.setattr(inspection, "HfApi", _OfflineApi) + monkeypatch.setattr(inspection, "try_to_load_from_cache", lambda *args, **kwargs: None) + monkeypatch.setattr( + inspection, + "scan_cache_dir", + lambda: pytest.fail("an explicit revision miss must not scan other cached revisions"), + ) + + with pytest.raises(SetupError, match="Cannot resolve Hugging Face model"): + inspection.inspect_model("Qwen/cached", revision="requested") + + def test_cached_model_ref_preserves_snapshot_commit_across_blob_symlink(tmp_path, monkeypatch): revision = "b" * 40 blob = tmp_path / "blobs" / "config" diff --git a/tests/unit/torch/puzzletron/test_setup_v2_data.py b/tests/unit/torch/puzzletron/test_setup_v2_data.py index 19bba98373e..a6a28cf16c1 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_data.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_data.py @@ -30,6 +30,7 @@ from puzzletron_setup.v2.state import WizardState from puzzletron_setup.v2.wizard import ( _CUSTOM_DATA_SOURCE, + _DEFAULT_DATA_SOURCE, _NEMOTRON_VLM_DATA_SOURCE, _PUZZLE_KD_DATA_SOURCE, _data_source_choices, @@ -43,7 +44,7 @@ def __init__(self, answers): self.checkbox_calls = [] def checkbox(self, message, choices, defaults): - self.checkbox_calls.append((message, tuple(choices), tuple(defaults))) + self.checkbox_calls.append((tuple(choices), tuple(defaults))) return super().checkbox(message, choices, defaults) @@ -99,10 +100,10 @@ def test_data_choices_include_first_class_sources_and_deduplicate_default(): choices = _data_source_choices(resolver) - assert [choice.title for choice in choices] == [ - f"Default — {_PUZZLE_KD_DATA_SOURCE}", - "NVIDIA Nemotron-VLM v2: recommended image-text dataset", - "Custom dataset: choose a local path or Hugging Face dataset", + assert [choice.value for choice in choices] == [ + _DEFAULT_DATA_SOURCE, + _NEMOTRON_VLM_DATA_SOURCE, + _CUSTOM_DATA_SOURCE, ] @@ -259,11 +260,13 @@ def test_nemotron_vlm_subset_prompt_uses_catalog_choices_and_defaults(tmp_path): _, backend = _run_nemotron_vlm_data_section(tmp_path) assert len(backend.checkbox_calls) == 1 - message, choices, defaults = backend.checkbox_calls[0] - assert message == "Dataset subsets:" - assert len(choices) == 46 - assert choices[0].title == "sparsetables — 100 rows — 1.00 KiB" - assert choices[3].disabled == "external media required" + choices, defaults = backend.checkbox_calls[0] + assert tuple(choice.value for choice in choices) == tuple( + subset.name for subset in _nemotron_catalog().subsets + ) + assert next(choice for choice in choices if choice.value == "external").disabled == ( + "external media required" + ) assert defaults == ("sparsetables", "plotqa_cot", "wiki_en") @@ -319,7 +322,7 @@ def test_generic_hugging_face_dataset_uses_dynamic_subset_checkbox( "weight": 0.9, }, ] - assert backend.checkbox_calls[0][2] == ("small",) + assert backend.checkbox_calls[0][1] == ("small",) def test_guided_explicit_invalid_subset_fails_instead_of_falling_back( @@ -436,7 +439,7 @@ def test_checkbox_rejects_a_disabled_scripted_selection(tmp_path): def test_interactive_checkbox_passes_disabled_reason_to_questionary(monkeypatch): - rendered = [] + shown_choices = [] class _KeyBindings: @staticmethod @@ -459,7 +462,6 @@ def ask(): class _Questionary: @staticmethod def Choice(**kwargs): # noqa: N802 - mirrors questionary's public constructor - rendered.append(kwargs) return kwargs @staticmethod @@ -472,11 +474,8 @@ def Style(rules): # noqa: N802 - mirrors questionary's public constructor @staticmethod def checkbox(message, choices, *, instruction, style): - assert message == "Subsets:" - assert choices[:2] == rendered - assert choices[2] == {"separator": " ← Back (press Esc)"} - assert " to go back" in instruction - assert style + del message, instruction, style + shown_choices.extend(choices) return _Question() monkeypatch.setattr( @@ -502,6 +501,6 @@ def checkbox(message, choices, *, instruction, style): ) assert selected == ["hosted"] - assert rendered[0]["checked"] - assert rendered[0]["disabled"] is None - assert rendered[1]["disabled"] == "external media required" + assert shown_choices[0]["checked"] + assert shown_choices[0]["disabled"] is None + assert shown_choices[1]["disabled"] == "external media required" diff --git a/tests/unit/torch/puzzletron/test_setup_v2_quick.py b/tests/unit/torch/puzzletron/test_setup_v2_quick.py index cd6e6894c53..830d0171ab4 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_quick.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_quick.py @@ -65,7 +65,7 @@ def test_only_cpu_slurm_integer_defaults_accept_null() -> None: "cpu_cpus_per_task": None, "cpu_memory_mb": None, } - with pytest.raises(SetupError, match="data.sequence_length must be an integer"): + with pytest.raises(SetupError, match=r"data\.sequence_length must be an integer"): validate_defaults({"schema_version": 1, "data": {"sequence_length": None}}) @@ -346,7 +346,7 @@ def test_guided_data_rejects_an_explicit_modality_incompatible_with_the_model( lambda source: SimpleNamespace(modality="text", evidence="local fixture"), ) - with pytest.raises(SetupError, match="multimodal.*incompatible"): + with pytest.raises(SetupError, match=r"multimodal.*incompatible"): data_section( WizardSession( state, @@ -450,9 +450,6 @@ def test_guided_wizard_runs_real_sections_and_generates_valid_bundles( ) assert result == campaign.resolve() - assert (campaign / "smoke" / "experiment.yaml").is_file() - assert (campaign / "production" / "experiment.yaml").is_file() - assert (campaign / "resolved_defaults.yaml").is_file() generated = WizardState.resume(campaign) assert generated.collection("pruning")["depth_remove"] == 0 assert generated.collection("pruning")["width_importance_samples"] == 8 @@ -488,19 +485,6 @@ def test_guided_wizard_runs_real_sections_and_generates_valid_bundles( '{"enable_thinking": false}', ] smoke_comparison = smoke_flow["nodes"]["quality_benchmarks"] - smoke_nodes = smoke_flow["nodes"] - smoke_mip_run = next(iter(smoke["mip"]["runs"].values())) - assert smoke_mip_run["solver"]["num_solutions"] == 2 - assert smoke_mip_run["homogeneous"]["keep"] == 2 - assert smoke_nodes["online_eval"]["config"]["eval_samples"] == 2 - assert smoke_nodes["online_eval"]["config"]["block_size"] == 32 - assert smoke_nodes["serving"]["config"]["input_tokens"] == 100 - assert smoke_nodes["serving"]["config"]["output_tokens"] == 80 - assert smoke_nodes["serving"]["config"]["request_count"] == 1 - assert smoke_nodes["short_kd"]["config"]["max_steps"] == 2 - assert smoke_nodes["short_kd"]["config"]["global_batch_size"] == 1 - assert smoke_nodes["short_kd"]["config"]["checkpoint_every_steps"] == 2 - assert smoke_comparison["config"]["limit"] == 8 assert "recorded_observation" not in smoke_comparison["config"] smoke_runner = yaml.safe_load((campaign / "smoke" / "runner.yaml").read_text()) assert smoke_runner["runner"]["slurm"]["job_name_prefix"] == "acct-puzzletron" @@ -546,7 +530,6 @@ def test_guided_wizard_runs_real_sections_and_generates_valid_bundles( "mmlu_pro_history", ] assert comparison["config"]["limit"] == 256 - assert smoke_comparison["config"]["limit"] < comparison["config"]["limit"] assert "recorded_observation" not in comparison["config"] resolved_defaults = yaml.safe_load((campaign / "resolved_defaults.yaml").read_text()) assert resolved_defaults["pruning.depth_remove"] == { @@ -630,7 +613,6 @@ def test_guided_wizard_generates_the_complete_qwen_vlm_flow(tmp_path, monkeypatc assert smoke_quality["limit_mm_per_prompt"] == {"image": 32} assert smoke_quality["max_model_len"] == 32768 assert "recorded_observation" not in smoke_quality - assert smoke_execution["execution"]["stages"]["replacement_scoring"]["instances"] == 1 assert smoke_execution["execution"]["stages"]["post.params-90.short_kd"]["instances"] == 1 production = yaml.safe_load((campaign / "production" / "experiment.yaml").read_text()) production_execution = yaml.safe_load((campaign / "production" / "execution.yaml").read_text()) diff --git a/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py b/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py index 0327eae9364..8cf142fc5c9 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py @@ -360,6 +360,70 @@ def test_resolved_sections_take_precedence_over_compatibility_overrides(tmp_path assert experiment["pruning"]["automodel"]["parallel"]["tp"] == 2 +def test_smoke_renderer_caps_bounded_work_without_changing_production(tmp_path: Path) -> None: + state = _campaign_state(tmp_path) + mip = deepcopy(state.collection("mip_config")) + run = mip["runs"]["params-90"] + run["solver"] = {"num_solutions": 8} + run["homogeneous"] = {"enabled": True, "keep": "all", "rank_by": "objective"} + state.set_collection("mip_config", mip) + nodes = { + "evaluation": { + "type": "evaluation", + "config": {"eval_samples": 128, "block_size": 2048}, + }, + "serving": { + "type": "aiperf", + "config": { + "input_tokens": 256, + "output_tokens": 160, + "request_count": 32, + "extra_inputs": {"min_tokens": 160}, + }, + }, + "filter": {"type": "filter", "top_k": 5}, + "kd": { + "type": "global_kd", + "config": { + "max_steps": 128, + "global_batch_size": 128, + "checkpoint_every_steps": 128, + }, + }, + "comparison": {"type": "downstream_evaluation", "config": {"limit": 64}}, + } + state.set_collection( + "post_mip_flows", + {"selection": {"source": {"run": "params-90"}, "nodes": nodes}}, + ) + + smoke = render_experiment_v2(state, "smoke") + production = render_experiment_v2(state, "production") + smoke_run = smoke["mip"]["runs"]["params-90"] + smoke_nodes = smoke["post_mip"]["flows"]["selection"]["nodes"] + production_run = production["mip"]["runs"]["params-90"] + + assert smoke_run["solver"]["num_solutions"] == 2 + assert smoke_run["homogeneous"]["keep"] == 2 + assert smoke_nodes["evaluation"]["config"] == {"eval_samples": 2, "block_size": 512} + assert smoke_nodes["serving"]["config"] == { + "input_tokens": 100, + "output_tokens": 80, + "request_count": 1, + "extra_inputs": {"min_tokens": 80}, + } + assert smoke_nodes["filter"]["top_k"] == 1 + assert smoke_nodes["kd"]["config"] == { + "max_steps": 2, + "global_batch_size": 1, + "checkpoint_every_steps": 2, + } + assert smoke_nodes["comparison"]["config"]["limit"] == 8 + assert production_run["solver"]["num_solutions"] == 8 + assert production_run["homogeneous"]["keep"] == "all" + assert production["post_mip"]["flows"]["selection"]["nodes"] == nodes + + def test_stage_batches_update_shared_data_sections(tmp_path: Path) -> None: state = _campaign_state(tmp_path) state.set_collection( From 784abe770ae3757f7027aafeb9441e825a5575e9 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Tue, 8 Sep 2026 01:32:04 +0200 Subject: [PATCH 4/4] Fix no-bias GDN test path Signed-off-by: Johannes Rausch --- .../unit/torch/puzzletron/test_automodel_solution_scoring.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py b/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py index 320836707ff..55b11fbc923 100644 --- a/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py +++ b/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py @@ -163,7 +163,7 @@ def test_native_automodel_gdn_executes_compact_geometry_and_rejects_context_para kernel_shapes = [] def compact_conv(*, x, weight, bias, **kwargs): - conv_shapes.append((x.shape, weight.shape, bias.shape)) + conv_shapes.append((x.shape, weight.shape, None if bias is None else bias.shape)) return x def compact_kernel(query, key, value, *, g, beta, **kwargs): @@ -197,7 +197,7 @@ def compact_kernel(query, key, value, *, g, beta, **kwargs): ( torch.Size((1, target_projection_width, 4)), torch.Size((target_projection_width, config.linear_conv_kernel_dim)), - torch.Size((target_projection_width,)), + None, ) ] assert kernel_shapes == [