diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 6e589c526dd..c3cdbee5712 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -7,164 +7,125 @@ 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 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. -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 +137,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 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 +172,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 +218,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 +242,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..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 an illustrative multi-axis campaign | `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, @@ -13,8 +19,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 +39,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 +105,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 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 +125,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 +144,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 +161,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..ef93b1e1760 100644 --- a/puzzletron_setup/v2/bundle.py +++ b/puzzletron_setup/v2/bundle.py @@ -232,15 +232,56 @@ 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: + 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: 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 +347,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 +407,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 +460,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 +513,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..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 @@ -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" @@ -158,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) @@ -168,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") @@ -226,6 +230,36 @@ 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) + 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): + 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 +276,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 +299,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 +377,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 +410,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() @@ -552,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/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..55b11fbc923 100644 --- a/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py +++ b/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py @@ -23,6 +23,7 @@ import pytest import torch +from omegaconf import OmegaConf from safetensors.torch import save_file from torch import nn @@ -33,7 +34,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 +52,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 +65,49 @@ 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(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_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 @@ -107,7 +149,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, @@ -115,16 +157,67 @@ 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) + conv_shapes = [] + kernel_shapes = [] - with pytest.raises(RuntimeError, match="refusing to score reduced geometry"): - apply_runtime_candidate(layer, teacher, child) + def compact_conv(*, x, weight, bias, **kwargs): + 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): + 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) + 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)), + None, + ) + ] + 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): @@ -533,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.""" @@ -602,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 6352df3bd0c..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_rejects_native_automodel_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,21 +139,47 @@ def test_compact_grouped_attention_rejects_native_automodel_backend(): num_query_heads=2, num_kv_heads=1, ) - original_forward = attention.forward.__func__ - original_state = set(vars(attention)) - - assert not 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 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) - assert attention.forward.__func__ is original_forward - assert set(vars(attention)) == original_state - @pytest.mark.parametrize(("target_num_q", "target_num_kv"), [(2, 1), (2, 2)]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) @@ -206,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( @@ -228,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, @@ -290,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( @@ -334,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) @@ -345,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) @@ -378,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 7ff8832e3df..fefddbaf4fa 100644 --- a/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py +++ b/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py @@ -22,9 +22,9 @@ _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, _select_layers, _write_hidden_only_diagnostic_artifacts, _write_reused_sort_equivalence, @@ -192,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" @@ -221,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() @@ -233,16 +213,46 @@ 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): + 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) + + +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() -@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_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_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..ae01bc14c95 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" @@ -265,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, *, @@ -887,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( @@ -1001,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 new file mode 100644 index 00000000000..836cd94de04 --- /dev/null +++ b/tests/unit/torch/puzzletron/test_post_mip_reporting.py @@ -0,0 +1,103 @@ +# 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 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 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_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..197cf1f867c 100644 --- a/tests/unit/torch/puzzletron/test_setup_inspection.py +++ b/tests/unit/torch/puzzletron/test_setup_inspection.py @@ -1,24 +1,43 @@ # 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 +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" ) - 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 +52,102 @@ 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') + + 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_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" + 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..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( @@ -409,6 +412,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): @@ -433,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 @@ -456,7 +462,6 @@ def ask(): class _Questionary: @staticmethod def Choice(**kwargs): # noqa: N802 - mirrors questionary's public constructor - rendered.append(kwargs) return kwargs @staticmethod @@ -469,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( @@ -499,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_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..830d0171ab4 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 ( @@ -63,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}}) @@ -344,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, @@ -448,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 @@ -469,6 +468,22 @@ 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"] assert "recorded_observation" not in smoke_comparison["config"] smoke_runner = yaml.safe_load((campaign / "smoke" / "runner.yaml").read_text()) @@ -515,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"] == { @@ -570,7 +584,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 +600,22 @@ 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"]["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 +628,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 +647,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 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(