diff --git a/.agents/skills/running-puzzletron/SKILL.md b/.agents/skills/running-puzzletron/SKILL.md deleted file mode 100644 index aa430454cbb..00000000000 --- a/.agents/skills/running-puzzletron/SKILL.md +++ /dev/null @@ -1,245 +0,0 @@ ---- -name: running-puzzletron -description: Use when onboarding a new LLM or multimodal model to Puzzletron, running or resuming a Puzzletron pruning campaign, adding model descriptors or pruning axes, validating activation hooks and physical slicing, collecting vLLM statistics, running MIP/evaluation/AIPerf/distillation, or debugging Puzzletron on Slurm or SSH-managed bare metal. ---- - -# Run Puzzletron - -## Overview - -Turn an exact model checkpoint and user requirements into a complete, -reproducible Puzzletron pruning campaign. Inspect actual model/backend source, -admit only physically valid axes, emit hashed smoke and production bundles, -execute the stage DAG, resume durable work, and keep the cumulative report -verified. - -If the user asks only for planning or documentation, produce those artifacts -and stop. Do not launch campaign work without explicit authorization. - -## Required References - -Read each applicable reference completely before acting: - -- For model onboarding, descriptors, pruning axes, hooks, sorting, width - ranking, physical slicing, multimodal models, or MTP, read - [model-and-axis-validation.md](references/model-and-axis-validation.md). -- For execution environments, AutoModel topology, model-instance semantics, - the stage DAG, RPC depth evaluation, resumability, Slurm, or bare metal, read - [distributed-execution.md](references/distributed-execution.md). -- For vLLM candidate construction, exact runtime measurement, cost/memory, - worker sharding, cache identity, or negative marginals, read - [vllm-runtime-statistics.md](references/vllm-runtime-statistics.md) and the - distributed-execution reference. -- For bypass, named MIP profiles, typed depth selection, artifact invalidation, - reports, trackers, tutorials, or handoff, read - [bypass-mip-and-reporting.md](references/bypass-mip-and-reporting.md) and the - distributed-execution reference. -- For a complete campaign, read all four references. - -## Core Contract - -1. Ask the two-phase intake below; do not invent consequential choices. -2. Inspect exact model and backend source before proposing axes. -3. Admit only axes valid through training, physical export, accounting, and all - required inference backends. -4. Treat physical materialization as ground truth for dynamic slicing. -5. Run a tiny full-coverage smoke and deliberately test durable resume. -6. Present one explicit readiness gate before expensive production work. -7. Emit runnable smoke and production bundles with immutable identities. -8. Execute the authoritative stage DAG and parallelize independent writers. -9. Regenerate and semantically verify the cumulative report after each stage. -10. Preserve or quarantine incompatible artifacts; never silently overwrite or - merge them. -11. Use the v2 campaign orchestrator for Slurm and bare-metal execution; do not - reconstruct scheduler submission with ad hoc shell loops. - -## Orchestrator Contract - -Use `examples/puzzletron/orchestrate.py` as the public campaign entry point and -`modelopt/torch/puzzletron/orchestration/` as the scheduler-neutral control -plane, exposed through the dependency-light `puzzletron_orchestrator` import -path. The controller requires only Python and PyYAML; do not install PyTorch or -the full ModelOpt runtime on a login node. The orchestrator binds three separate -inputs: - -- the canonical Puzzletron experiment YAML, which owns algorithm semantics and - every stage-local `automodel.parallel` mesh; -- a runner YAML, which owns exactly one Slurm or bare-metal environment plus - the repository, venv, container, mounts, setup contract, and optional - `prerun_commands` / `postrun_commands` shell hooks; and -- an execution YAML, which owns each stage's `single`, `sharded`, or - `persistent_pool` strategy, independent `instances`, and failure policy. - -Treat `instances` as independent model/process-group copies. Never multiply it -into AutoModel DP. Derive one instance's allocation as -`PP * DP_REPLICATE * DP_SHARD * CP * TP`; EP overlays `DP_SHARD` and requires -`DP_SHARD % EP == 0`. Let the orchestrator pack independent instances onto -nodes—for example, sixteen one-GPU vLLM-statistics instances on eight-GPU nodes -must resolve to two fully utilized nodes. - -Before any real submission: - -1. parse the experiment with the pipeline loader; -2. run the orchestrator with `--dry-run`; -3. verify every stage's strategy, instance count, mesh, node/GPU allocation, - command, artifact writer, and DAG dependency; -4. verify Slurm account/partition/time limit/container/mounts or every - bare-metal host/GPU lease and rendezvous input; and -5. run the smoke bundle before the production bundle. - -Canonical commands are: - -```bash -python examples/puzzletron/orchestrate.py \ - --experiment "$EXPERIMENT" \ - --runner "$RUNNER" \ - --execution "$EXECUTION" \ - --stage full --dry-run - -python examples/puzzletron/orchestrate.py \ - --experiment "$EXPERIMENT" \ - --runner "$RUNNER" \ - --execution "$EXECUTION" \ - --stage full -``` - -Use `${puzzle_dir}/orchestration/` for controller state, attempt provenance, -recovery, and logs. Resume with the same three hashed inputs and command; do not -fall back to `main.py --stage full` or manual `sbatch`/SSH commands when the -orchestrator supports the stage. Stop if dry-run emits an invalid command, -incorrect packing, unresolved path, unsupported adapter, or conflicting -canonical writer. - -## Phase 1: Infrastructure and Goals - -Ask one grouped infrastructure message covering: - -- exact checkpoint and tokenizer/processor revisions; -- remote-code trust policy and dataset identity/access; -- experiment ID, shared artifact root, and existing-artifact policy; -- Slurm account/partitions/node types/wall limits or bare-metal hosts/slots; -- container image or host Python, mounts, virtual environment, site setup, and - sibling AutoModel/vLLM/AIPerf checkouts; -- allowed code changes, available GPU/storage budget, deadline, and cost limit; -- required outcome: search, evaluation, AIPerf, bypass, KD, post-KD evaluation, - or the complete DAG. - -When using Slurm, read `nv-internal/CLUSTER_GUIDE_NV.md` if present. For bare -metal, require passwordless SSH, identical shared paths, and an explicit -rendezvous plan. - -## Inspect Before Phase 2 - -Resolve the exact checkpoint revision and inspect every distinct computational -layer: embeddings, attention, dense/MoE experts and routers, recurrent/SSM -layers, norms, multimodal towers/projectors, LM heads, and MTP heads. Read, in -order, the exact Hugging Face implementation, existing vLLM implementation, -AutoModel implementation, and Puzzletron descriptors/hooks/materializers/cost -models. - -Build an inventory of semantic axes, legal domains, coupled dimensions, -alignment/grouping constraints, distributed implications, backend support, -dynamic hooks, physical export, and accounting. Record accepted and rejected -axes with evidence. - -Use an AutoModel descriptor only when AutoModel already supports the unpruned -model; otherwise use the HF-native path. Do not add a new native AutoModel or -vLLM base model as incidental campaign work. If the base model runs in vLLM, -every admitted pruned configuration must work through the repository's existing -AnyModel vLLM path. - -## Phase 2: Model-Aware Questions - -Present only options valid for the inspected model. Ask, in one grouped search -message: - -- real datasets, splits, modalities, sampling weights, fixed/padded/packed - layouts, and native chat template; -- legal candidate domains per admitted axis and maximum depth removals; -- bypass, vLLM statistics, and replace-one granularity (`block` or `subblock`); -- parameter, runtime, memory, and active-parameter constraints plus workloads; -- number of solutions per profile for evaluation, AIPerf, and distillation; -- smoke and production TP/CP/PP/EP/DP/sequence-parallel settings; -- evaluation metrics, acceptance gaps, seeds, deadline, and compute envelope. - -Explain trade-offs briefly. Ask only follow-ups that materially change the -campaign. - -## Emit Runnable Bundles - -Create separate `smoke` and `production` namespaces. Each contains: - -- one canonical experiment YAML accepted by `examples/puzzletron/main.py`; -- one runner YAML and one execution YAML accepted by - `examples/puzzletron/orchestrate.py`; -- stage-local `automodel.parallel` meshes (never a shared AutoModel recipe YAML), data - manifests, exact revisions, and execution contract; -- resolved immutable snapshot plus hash and deterministic artifact identities; -- exact orchestrator dry-run, launch, resume, monitor, report-only, and safe - cleanup commands; -- a stage/resource table with dependencies, topology, outputs, and status. - -Parse both experiment configs with the pipeline loader and compile both -orchestrator plans before presenting them. Keep a live task tracker and a -reproduction tutorial from the start; the reporting reference defines when -each may be updated. - -## Execute and Validate - -Use `modelopt/torch/puzzletron/stages/graph.py` as the dependency authority. -Launch the DAG through `examples/puzzletron/orchestrate.py`. The controller -schedules independent branches concurrently when artifact writers do not -conflict. One writer publishes each canonical identity; immutable worker shards -are aggregated transactionally. - -Smoke data and iteration counts may be tiny, but coverage must include every -distinct layer/axis, requested modality/layout, requested distributed path, -conversion, tokenization, importance, sorting, dynamic/physical slicing, vLLM -and scoring when enabled, MIP/evaluation/AIPerf, bypass/KD sanity and short real -runs, interruption/resume, and report regeneration. - -Fail readiness for wrong equivalence, broken resume, invalid artifacts, backend -incompatibility, or missing report sections. Warnings about ranking quality may -remain visible without blocking when correctness gates pass and the user has -accepted that policy. - -## Readiness Gate - -Before production, present: - -- accepted/rejected axes and source evidence; -- smoke outcomes and unresolved warnings; -- production DAG with optional/disabled nodes; -- exact config paths and hashes; -- resource/topology estimates and total cost envelope; -- resumability/checkpoint plan; -- orchestrator dry-run, launch, monitor, and resume commands. - -Ask for one explicit approval. After approval, operate autonomously inside the -authorized cost, monitor long jobs, diagnose from first-rank evidence, resume -durable work, cancel redundant continuations, and never rerun expensive complete -shards. - -## Stop Conditions - -Stop and request direction rather than silently reducing coverage when: - -- checkpoint/data access or licensing is unresolved; -- the unpruned model cannot run in a required existing backend; -- a required layer type has no valid pruning operation; -- physical and dynamic slicing disagree materially; -- multimodal or MTP target boundaries are incorrect; -- smoke resume is not durable; -- production config has unresolved placeholders; -- requested resources exceed the approved envelope. - -## Repository Navigation - -Start with `examples/puzzletron/README.md`, `examples/puzzletron/main.py`, and -`examples/puzzletron/orchestrate.py`, then read -`modelopt/torch/puzzletron/orchestration/` and -`modelopt/torch/puzzletron/stages/graph.py`. Read the cluster guide when present -and `nv-internal/PUZZLETRON_V2_ENGINEERING_GUIDE.md` for historical findings, -validating each claim against current source. Use symbol search because paths -evolve. Follow repository coding/test instructions, preserve dirty work, keep -changes generic, and validate proportionally to risk. diff --git a/.agents/skills/running-puzzletron/agents/openai.yaml b/.agents/skills/running-puzzletron/agents/openai.yaml deleted file mode 100644 index 2cb331fe8fb..00000000000 --- a/.agents/skills/running-puzzletron/agents/openai.yaml +++ /dev/null @@ -1,4 +0,0 @@ -interface: - display_name: "Run Puzzletron" - short_description: "Onboard and run resumable Puzzletron campaigns" - default_prompt: "Use $running-puzzletron to onboard this model and produce a runnable Puzzletron campaign configuration." diff --git a/.agents/skills/running-puzzletron/references/bypass-mip-and-reporting.md b/.agents/skills/running-puzzletron/references/bypass-mip-and-reporting.md deleted file mode 100644 index 6f9093efb78..00000000000 --- a/.agents/skills/running-puzzletron/references/bypass-mip-and-reporting.md +++ /dev/null @@ -1,101 +0,0 @@ -# Bypass, MIP, and Reporting - -## Configure Bypass - -Bypass replaces a selected block/subblock input with the matching teacher input -and computes a local loss. Ensure PP batch/microbatch compatibility, identical -CP partitions and masks, required TP sequence parallelism without duplicate -reductions, and valid per-DP architecture sampling. - -Every observation records step, DP rank, layer/subblock, canonical architecture -hash, human-readable config, normalized parameter ratio, component/total loss, -seed, and sample identity. - -Run two sanity modes before production: - -1. a fixed smallest config that clearly overfits; -2. a diverse resampled schedule whose trend decreases despite scatter. - -Define production scale with total tokens, sequence length, tokens/update, -microbatch, gradient accumulation, and derived target steps. On a fresh run, -initialize the current runner's state cursors at `step_num=1`, `iter_num=0`, and -`token_count=0`; these are resume counters, not target totals. - -Overfit success validates boundaries and optimization mechanics, not production -LR. Before spending the full budget, require an early finite production trend -with improvement represented across every sampled width/config family. Record -configured maximum LR, observed LR, scheduler phase, loss windows, and gradient -norms. A successful exit with flat/worsening loss is rejected evidence. - -Measure tokens/second before adding nodes. Communication, restore, and routing -overhead may make a larger DP topology slower. Keep the smallest topology with -the best verified throughput that fits memory. When monitoring a resume, join -the durable checkpoint/history with newer live-log records across the boundary. - -## Compile Named MIP Profiles - -Read `examples/puzzletron/docs/mip_profiles.md`. Define named workloads in the -root MIP namespace and reference them from runtime/memory constraints. Multiple -constraints may target different workloads; resolve each percentage against -that workload's teacher total. Missing workload data is invalid, not a reason to -relax the constraint. - -Support exact values, percentages, ranges, and lists. Scalar shorthand is -equality; use explicit `min`/`max` for inequalities. Use -`num_homogeneous_solutions` for homogeneous Cartesian baselines (`-1` means all) -and restricted-axis domains to compare mix-and-match against single-axis -searches. Omitted domains inherit all legal values; explicitly fixed domains -stay at the teacher value. - -Depth collection remains one global iterative trajectory. Profiles may request -a total prefix or, only with subblock depth importance, typed attention, -Mamba/SSM, MoE, or dense-FFN counts. Expand lists by Cartesian product, select -each kind's prefix from the global trajectory, and union while preserving global -order. Omitted kinds mean zero. Never mix `total` and typed counts. Validate -availability before solving. - -Identity scenarios with the complete typed depth map, profile, width, domains, -constraints, and workload—not only total removals. Equal-total typed selections -remain distinct through resume, materialization, reports, evaluation, AIPerf, -and KD. Keep profile identity even when two profiles produce the same structure. - -## Artifact Invalidation - -Changing LR, token budget, estimator schema, physical candidate construction, -workload, or execution contract creates a new semantic identity. Quarantine -incompatible canonical outputs and selected downstream descendants. Never merge -old/new shards because filenames or shapes match. - -Before rerunning an invalidated stage, regenerate the report and verify old -results disappeared and the stage is pending. Rebuild canonical aggregates only -from compatible completed shards. Preserve rejected evidence under a clear -non-canonical namespace. - -## Report Verification - -The cumulative HTML is generated from canonical artifact contracts after each -completed stage and may show partial long-running observations without marking -them complete. Include experiment identity, DAG, resolved config, model/axis -inventory, provenance, cell-level warnings, sanity/results, and pending, -disabled, optional, and completed nodes. Do not duplicate post-KD metrics. - -Use model identity precedence: configured display name, declared Hugging Face -repository plus revision, model metadata, then local path. Never title the -report with a resolved HF cache snapshot when repository identity exists. - -Verification is semantic: parse the report and assert model identity, stage -state, candidate/width/depth counts, profile names, warnings, and canonical -paths. Partial caches and worker shards are progress only. Report-only fixes do -not invalidate expensive compatible compute. - -## Tracker, Tutorial, and Handoff - -Maintain a live task tracker containing jobs, progress, evidence, dependencies, -and next actions. It may record partial observations. Maintain a reproduction -tutorial containing only commands and conclusions whose artifacts and report -sections have been verified. - -At handoff provide bundle paths/hashes, run/resume/monitor/report/cleanup -commands, artifact/report paths, DAG states, measured resources/runtime, -warnings, rejected axes, and the recommended next action. - diff --git a/.agents/skills/running-puzzletron/references/distributed-execution.md b/.agents/skills/running-puzzletron/references/distributed-execution.md deleted file mode 100644 index 42eedcdff1e..00000000000 --- a/.agents/skills/running-puzzletron/references/distributed-execution.md +++ /dev/null @@ -1,116 +0,0 @@ -# Distributed Execution - -## Persist the Execution Contract - -Before any task, record the container image/digest, mounts, repository root, -venv, required site setup, cache roots, sibling revisions, and experiment- -scoped environment overrides. Apply the same ordered prologue to every Python -command, including CPU conversion, tokenization, reports, and tests. Enforce any -cluster policy that permits only shell utilities or Python standard library on -login nodes. - -Hash this contract into resumable identities. Source site setup before venv -activation when it sets caches or compiler state. Reuse a persistent allocation -for short iterations; use interactive capacity for unstable debugging and batch -for validated work. Bind disjoint `CUDA_VISIBLE_DEVICES` for partial-node tasks -and avoid exclusive nodes when GPUs are intentionally unused. - -## Parallelism and Model Instances - -Ask for TP, CP, PP, EP, DP, and sequence parallelism separately for smoke and -production. Validate backend constraints rather than assuming all named degrees -multiply independently. - -For AutoModel FSDP/HSDP, record the realized mesh as -`(PP, DP_REPLICATE, DP_SHARD, CP, TP)`: - -- `DP_SHARD` is the FSDP shard group; it collectively owns one model replica. -- `DP_REPLICATE` is the HSDP replica axis and creates sample-parallel lanes. -- EP overlays `DP_SHARD`; it is not another multiplicative allocation axis. -- Logical sample DP is `DP_REPLICATE × (DP_SHARD / EP)` when EP fully overlays - the shard axis. -- Allocation is `PP × DP_REPLICATE × DP_SHARD × CP × TP`. - -For example, `(2, 8, 4, 1, 1)` with EP4 consumes 64 GPUs and has logical DP8. -Verify the realized mesh in logs. Derive PP microbatch divisibility and update -counts from logical sample DP, not raw shards or GPU count. Increase batch only -after measuring memory headroom and preserving sample/identity semantics. - -Classify each stage by model instance: - -- Width importance is one coordinated instance with one global hook/reduction - set. Scale through logical DP/HSDP; never merge independent width models. -- RPC depth uses multiple independent persistent model instances, each scoring - different cumulative replacement requests. -- Sharded vLLM statistics may use many independent one-GPU instances because - each immutable layout measurement is self-contained. - -## Execute the DAG - -Use `modelopt/torch/puzzletron/stages/graph.py` as dependency authority. -Conversion/tokenization may overlap vLLM stats; depth may overlap coordinated -width importance; vLLM may overlap bypass when allocations and writers are -disjoint. Only one writer publishes a canonical identity. - -## RPC Iterative Depth - -Use `modelopt.torch.puzzletron.distributed_eval`. One worker group is one -persistent distributed model instance that loads the sorted teacher and target -cache once, then evaluates many cumulative no-op requests. A group may use part -of a node, one node, or multiple nodes. Use the coordinator; do not add a second -filesystem scheduler. - -Depth iterations are sequential; candidates within one iteration are parallel: - -1. Load the durable selected prefix and cached results. -2. Submit the prefix baseline and each `prefix + candidate` request. Merge - multiple removed sublayers in one layer into one replacement. -3. Apply cumulative replacements in temporary prune contexts, score, and - restore the resident model before another request. -4. Rank the complete iteration atomically, publish trajectory, then continue. - -Identity every request by model, data, metric, precision, cumulative removals, -and evaluator revision. Retry only missing/transient requests. The coordinator -alone writes rankings and trajectory. - -When multiple groups share a node, assign disjoint CUDA devices, unique HTTP -ports, rendezvous IDs/ports, and worker IDs. All groups in one campaign must -match its declared world size/topology. Drain and terminate them explicitly. - -## Resumable Expensive Stages - -Bypass, KD, scoring, and vLLM stats use immutable shards and transactional -checkpoints. A training checkpoint stores model/optimizer shards, scheduler, -scaler, global step, Python/NumPy/CPU/CUDA/sampler/per-rank RNG, dataloader -cursor/sample order, topology/world size, config/code/model/data identities, -manifest, and atomic completion marker. - -Write to a temporary transaction, validate expected shards, publish atomically, -then update `latest`. Quarantine incomplete transactions. Smoke must interrupt -and resume at least bypass and global KD, verifying exact step/sample continuity -and report validity. - -Any semantic change—LR, token budget, estimator schema, physical construction, -workload, or execution contract—creates a new identity and invalidates selected -downstream dependencies. Never merge old/new shards based only on matching -filenames or shapes. - -## Slurm - -Derive launch commands from the cluster guide. Capture job IDs/logs with -`pipefail`. Checkpoint well before wall time and queue compatible resumable -continuations. Cancel redundant continuations after valid completion. - -Use persistent allocations when startup dominates short tests, but release them -when GPUs would idle. Run stable paths on batch and debug unresolved paths on -interactive. For sharded single-GPU work, request one GPU per task, avoid -exclusive nodes, and verify scheduler binding inside every worker. - -## Bare Metal - -Verify passwordless SSH, paths, clocks, ports, GPU visibility, and environment -compatibility on every host. Choose one rendezvous host/port, create -deterministic rank mappings, launch one process group, record PIDs/logs/hashes, -health-check all ranks, and terminate all peers after a local failure. Shared -storage does not prove identical Python/CUDA environments. - diff --git a/.agents/skills/running-puzzletron/references/model-and-axis-validation.md b/.agents/skills/running-puzzletron/references/model-and-axis-validation.md deleted file mode 100644 index 2a42924c90f..00000000000 --- a/.agents/skills/running-puzzletron/references/model-and-axis-validation.md +++ /dev/null @@ -1,102 +0,0 @@ -# Model and Axis Validation - -## Inventory and Admission - -For every candidate axis, verify: - -1. semantic capacity removed and tensors changed; -2. legal discrete values with alignment/grouping constraints; -3. dynamic slicing or masking; -4. physical materialization/export; -5. state-dict conversion and reload; -6. total/active parameter, memory, and runtime accounting; -7. required HF, AutoModel, and vLLM compatibility; -8. physical-versus-dynamic equivalence. - -Treat norms, rotary encodings, grouped heads, tied weights, residual -projections, recurrent state, and multimodal projectors as coupled systems. -Zeroing output channels is not equivalent to physically reducing a normalized -input. Reject an invalid axis with evidence and continue with other axes. - -## MoE Contracts - -Inspect the actual latent contract. If `moe_latent_size` is absent or `None`, -expose no latent-MoE axis and rank/slice intermediate channels independently per -expert. If a latent projection exists, validate it as a coupled axis. Never -infer latent MoE from another model in the same family. - -Preserve native router semantics. Correction bias may affect selection without -becoming a route weight. Grouped top-k may require fixed group membership. With -multiple groups, sort and compact experts within each original group unless -source inspection proves whole-group permutation is a symmetry. Run selection -on the physically compact vector and map compact IDs back only for resident- -teacher execution. Apply identical indices to router rows, correction bias, -auxiliary tensors, and expert weights. - -Sort routed-expert intermediate channels independently inside each expert -before sorting expert identities. Keep shared-expert channel sorting separate. -Never reuse stale inner-channel scores after an outer identity permutation. - -## Attention and Mamba/SSM Contracts - -For grouped attention, sort query heads inside each KV group before sorting KV -groups. Compose that group-major order consistently across Q/K/V rows, output -columns, biases, rotary state, and auxiliary tensors. - -For Mamba/SSM axes, preserve fixed group topology. Sort channels within the -head representation, then sort heads within fixed groups, and apply one -group-preserving composition to every coupled tensor. A physical reduction must -compact projected inputs, convolution tensors, state parameters, normalization -weight, and output-projection columns passed to the fused kernel. Teacher-sized -masking plus compensation is not sufficient BF16 evidence. - -## Activation Hooks - -Define the exact module boundary, measured tensor, reductions, statistic, -valid-token/media mask, packed-document boundaries, accumulation dtype, -numerical guards, shard ownership, and durable schema. - -- DP ranks may see different samples; combine only commutative statistics. -- CP ranks own sequence shards; exclude padding and reduce correctly. -- TP ranks may own or replicate features; avoid double counting. -- PP stages observe local modules; use globally unique layer identities. -- EP ranks observe local experts; retain expert and routing identity. - -The manifest includes model/data/config hashes, topology, expected/completed -shards, and a completion marker. Resume missing shards only. - -## Sorting, Width Quality, and Physical Slicing - -Keep three separate checks: - -- **Sort sanity:** original, sorted, and reverse-sorted full-width teachers must - agree within dtype-aware tolerance. Reverse sorting is an equivalence control. -- **Width sanity:** compare activation-sorted, unsorted/random, and reverse at - reduced width. Poor ranking is a visible quality warning. -- **Slicing sanity:** compare dynamic sorted, dynamic unsorted, dynamic reverse, - and physically materialized sorted against the same original. Dynamic versus - physical disagreement is a correctness failure. - -Verify every permutation is bijective and every coupled tensor receives the -same composition. By default, test three representative layers and two legal -targets per axis in production. For global embedding/residual width, use about -seven-eighths and three-quarters of teacher width when aligned. A smoke may be -smaller only while covering every layer implementation and axis. - -Require exact config changes, parameter reduction, strict state-dict reload, -and all required backend forward/export paths. For native fused Mamba/SSM, -compare compact runtime arguments with physical export and require bitwise or -tightly bounded kernel-output parity. - -## Multimodal and MTP - -Use processor-native real examples for each modality. Preserve modality order, -placeholder tokens, grids/timestamps, cross-attention masks, labels, and packed -boundaries. Classify each tower/projector as prunable, immutable, or coupled to -language width, and include its costs even when not searchable. - -For MTP, inspect depth, shifted targets, ties/sharing, projectors/norms, backend -support, and topology. Targets must not cross padding, document, or modality -boundaries. Report `main_ce`, `mtp_ce`, `main_kd`, and `mtp_kd` separately plus -their weighted total. Include nonzero MTP CE and KD in smoke when supported. - diff --git a/.agents/skills/running-puzzletron/references/vllm-runtime-statistics.md b/.agents/skills/running-puzzletron/references/vllm-runtime-statistics.md deleted file mode 100644 index 034b83db93b..00000000000 --- a/.agents/skills/running-puzzletron/references/vllm-runtime-statistics.md +++ /dev/null @@ -1,90 +0,0 @@ -# vLLM Runtime Statistics - -## Inspect Runtime and Cost Code - -Read the current implementations before changing formulas: - -- `modelopt/torch/puzzletron/export/vllm.py` -- `modelopt/torch/puzzletron/utils/vllm_adapter.py` -- `modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py` -- `modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py` - -Count total and active parameters separately. Include embeddings, projectors, -norms, routers, experts and top-k activation, tied weights, MTP, KV/SSM caches, -activations, and server overhead. State topology and workload assumptions. - -## Exact Production Construction - -Instantiate every production candidate at exact physical dimensions. Proxy caps -for expert count, routed/shared width, top-k, sequence length, or batch size are -legal only behind an explicit smoke-only flag. Put the flag and caps in cache -identity and never aggregate proxy measurements into production profiles. - -Build temporary configs through the model descriptor. Derive the exact hybrid -pattern and mutually exclusive attention, Mamba/SSM, MoE, or dense-FFN -placement. Do not emit dense-FFN candidates/report rows for models without dense -FFNs. Keep a minimal valid scaffold only when homogeneous candidates cannot be -instantiated alone. - -## Paired Marginal Estimation - -Use paired `N` and `2N` layouts to cancel fixed overhead, defaulting to `N=4` -for production after a smaller numerical gate. Measure the identical candidate -and workload at both endpoints. Prefer candidate-only homogeneous layouts; -otherwise keep the same descriptor-defined scaffold in both so it cancels. - -Changing `N`, warmups, iterations, scaffold, hybrid pattern, workload, or -estimator schema changes cache identity. Recover and validate the fixed overhead -and repeated controls; do not trust close rankings when control spread is large. - -## One-GPU Worker Isolation - -For each independent worker: - -- bind exactly one explicit CUDA device; -- assign each vLLM subprocess a unique rendezvous port; -- isolate Triton, TorchInductor, TileLang, vLLM, FlashInfer, CUDA, temporary, - PyTorch-kernel, and XDG caches; -- share only immutable model/data caches; -- stagger local starts when compilation/filesystem pressure is high; -- write immutable phase results and transactional shard completion markers. - -Keep backend overrides scoped to this experiment. Do not globally disable a -fused path to make one benchmark pass. Verify `pip check`, vLLM and MoE native -extension imports, and one real-model fused-kernel gate before a large sweep. - -## Cache and Aggregation - -Cache identity covers the finalized model config, exact candidate, workload, -topology, estimator, environment overrides, GPU type, and software revisions. -Temporary paths and ports are non-semantic. If construction/timing code changes -without a schema bump, quarantine the cache manually. - -Resume only compatible missing endpoints. Publish phase JSON transactionally, -then shard markers. One aggregator validates all shard identities, indices, -counts, finiteness, and candidate coverage before publishing canonical -`subblock_stats.json`. Exclude structural no-op controls from candidate counts -and MIP costs. A no-op or scaffold must not create a fake FFN row for a model -without FFNs. - -## Anomalous Marginals - -Never clamp negative or non-finite marginals. Identify and remeasure only the -exact original `N`/`2N` endpoints on quiet, pinned hardware while retaining -unaffected compatible caches. A higher-iteration diagnostic has a different -identity and cannot silently replace production data. - -Non-positive total cost, attention/control failures, or widespread anomalies -fail the gate. An explicit `ignore_negatives` policy may retain a characterized -phase-only noise-dominated value with warning/provenance, but that value cannot -support a constraint on the affected phase. Record endpoints, remeasurement, -and acceptance rationale in the report. - -## Completion Gate - -Require the exact active-candidate count at every width/workload, all expected -layout phases and shard markers, matching identities, finite metrics, valid -fixed-overhead controls, and parameter/cache/memory accounting. Then publish the -canonical summary, stage manifest, and regenerated cumulative report. Partial -caches are progress, not stage completion. - diff --git a/.claude/skills/running-puzzletron b/.claude/skills/running-puzzletron deleted file mode 120000 index bde32cd6376..00000000000 --- a/.claude/skills/running-puzzletron +++ /dev/null @@ -1 +0,0 @@ -../../.agents/skills/running-puzzletron \ No newline at end of file diff --git a/docs/source/guides/3_pruning.rst b/docs/source/guides/3_pruning.rst index f866d219214..e4a07b43fa6 100644 --- a/docs/source/guides/3_pruning.rst +++ b/docs/source/guides/3_pruning.rst @@ -29,7 +29,8 @@ attention heads of the model. More details on these pruning modes are as follows FastNAS finds the subnet which maximizes the score function while meeting the given constraints. The remainder of this guide covers the unified Minitron and FastNAS API. -Puzzletron v2 instead uses its setup wizard and campaign runner. +Puzzletron uses checked-in recipes for maintained routes and +``python examples/puzzletron/puzzletron.py setup`` for custom-model onboarding. Follow the steps described below to obtain the optimal model satisfying your requirements using :mod:`mtp`: diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 09509106c85..95a41c8b959 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -1,269 +1,110 @@ -# Puzzletron v2 +# Puzzletron -Puzzletron v2 creates resumable pruning campaigns that search model shapes, -validate the selected candidates, and optionally evaluate, benchmark, -materialize, or distill them. Start here to choose a supported route. Each -linked guide owns the complete commands and result checks for its task. The -maintained VLM smoke commands are repeated here to keep the first run linear; -the model guide remains their canonical reference. +Puzzletron runs resumable pruning campaigns. A maintained run has two authored +inputs: -## Table of contents +- a small recipe that selects the model, workflow, mode, run directory, and + resource profile; +- one reusable site file for worker paths, environment, scheduler, and + available capacity. -- [Start here: lifecycle smoke](#start-here-lifecycle-smoke) -- [Choose the next task](#choose-the-next-task) -- [Documentation map](#documentation-map) +Puzzletron resolves the full stage configuration, validates it, and stores an +immutable copy with the run. -## Start here: lifecycle smoke +## Quickstart -Start with the maintained Qwen 3.5 0.8B VLM smoke. It exercises the complete -image-text lifecycle with bounded workloads inside one reusable Slurm allocation -that reserves two GPUs on a single node. Logical stages still have separate -attempts, logs, artifacts, and resume checks. Once it succeeds, the optional -longer example campaign uses the same environment, runner, orchestrator, -progress display, report, and resume command. Logical stage progress is emitted -in the reusable allocation log. 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. - -The commands in this first-run path explicitly select maintained execution -files that enable `reusable_allocation`. This is not Puzzletron's global -default. Wizard-generated bundles use `per_attempt`, which submits each stage -attempt as a separate Slurm job, unless the generated execution file is changed -after checking that the site supports a reusable single-node allocation. - -### 1. Create the controller environment - -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: +Use Python 3.10 through 3.14 for the lightweight controller environment: ```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 controller environment does not need PyTorch or model weights. Model -conversion, evaluation, serving, and distillation run in the worker environment -selected by the runner. - -### 2. Configure the worker runner - -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: +Create one site file and fill in its worker-visible checkout, virtual +environment, shared Hugging Face cache, Slurm account, and partition: ```bash -cp examples/puzzletron/configs/orchestration/runner.slurm.example.yaml \ - runner.slurm.yaml +cp examples/puzzletron/configs/site.example.yaml puzzletron.site.yaml ``` -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. - -These accelerated profiles are qualified for a Slurm environment with -eight-GPU nodes: the smoke reserves two GPUs on one node, while the -representative campaign reserves all eight GPUs on one node. They are not -qualified for arbitrary Slurm sites. See [Slurm -configuration](docs/slurm_configuration.md#reusable-single-node-allocations) -for the assumptions and settings another site may need to adapt. - -The three campaign inputs have separate responsibilities: - -- 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 selects reusable allocation and maps each logical stage - onto a GPU lease or zero-GPU CPU execution inside the shared allocation. - -### 3. Select shared paths - -Set a Hugging Face cache visible to workers and a new run root on shared -storage: +Choose a checked-in recipe. Start with the Qwen 3.5 0.8B VLM smoke, which +checks the full image-text lifecycle on two colocated GPUs: ```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"]))')" +cp examples/puzzletron/configs/recipes/qwen3p5_0p8b_vlm_smoke.yaml puzzletron.recipe.yaml +# Edit puzzletron.recipe.yaml and set run_root to a new shared path. ``` -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). - -### 4. Inspect and run the smoke - -Activate `.venv-puzzletron`, select the maintained files, and inspect the exact -plan before it requests resources: +Validate the inputs, inspect their provenance and resources, and preview the +scheduler commands: ```bash -EXPERIMENT=examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/full_vlm_smoke.yaml -EXECUTION=examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_smoke.yaml -RUNNER=runner.slurm.yaml - -python examples/puzzletron/orchestrate.py \ - --experiment "$EXPERIMENT" \ - --runner "$RUNNER" \ - --execution "$EXECUTION" \ - --stage full --dry-run +python examples/puzzletron/puzzletron.py validate puzzletron.recipe.yaml \ + --site puzzletron.site.yaml +python examples/puzzletron/puzzletron.py explain puzzletron.recipe.yaml \ + --site puzzletron.site.yaml +python examples/puzzletron/puzzletron.py dry-run puzzletron.recipe.yaml \ + --site puzzletron.site.yaml ``` -Review the resolved paths, CPU and GPU requests, worker image, mounts, and log -locations. This profile's dry-run shows one outer allocation and the logical -attempts that will execute inside it. Then run the same plan without -`--dry-run`: +Launch the same pair: ```bash -python examples/puzzletron/orchestrate.py \ - --experiment "$EXPERIMENT" \ - --runner "$RUNNER" \ - --execution "$EXECUTION" \ - --stage full +python examples/puzzletron/puzzletron.py launch puzzletron.recipe.yaml \ + --site puzzletron.site.yaml ``` -### 5. Resume and inspect results - -Run the launch command above again to recover an interrupted smoke or verify a completed one. It reattaches to a compatible active outer allocation. After Slurm reports that allocation failed or was cancelled, it starts a replacement unless a compatible worker result records clean completion or an uncancelled terminal stage failure. The replacement reruns only unfinished compatible work. A recorded cancellation and a final-report failure remain retryable; completed logical stages are not rerun. Changing the configuration or run root creates a new plan identity; compatible completed stages in the same run root are still not submitted again. Puzzletron stores its structured runtime state under `$PUZZLETRON_RUN_ROOT/orchestration/`. - -The reusable-allocation watcher shows the outer Slurm state and allocation log -path. Follow that log to see the inner controller's completed/total stages, -queued and running stages, elapsed time, 30-second heartbeats, and native inner -units such as 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 -`/artifacts/campaign_report/campaign_report.html`. A report failure -does not invalidate completed logical stages, but it is recorded in the run -result and the command exits nonzero so the report can be retried. See -[run and recovery options](docs/orchestration_operations.md) for individual -stages, `--once`, logging controls, security options, and recovery details, or -[campaign reports](docs/campaign_reports.md) to regenerate and interpret a -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 integration example - -After the smoke succeeds, keep the controller venv and runner and select a new -run root plus the longer example files: +Resume and inspect by run directory. Resume uses the sealed inputs stored with +the run, not the current recipe or site file. ```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 +python examples/puzzletron/puzzletron.py resume /shared/puzzle_runs/my-run +python examples/puzzletron/puzzletron.py inspect /shared/puzzle_runs/my-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. Its execution profile uses -one eight-GPU node and configured within-node candidate concurrency when -dependencies are ready, but does not introduce another operational path. Use -focused tests and the lifecycle smoke for routine development; reserve this -campaign for scheduled integration validation. - -## Choose the next task - -These guides own their model-specific configuration, launch commands, expected -artifacts, and interpretation limits: +Use a new run directory whenever an authored input changes. -| Model and route | Use it for | -| --- | --- | -| [Qwen 3.5 0.8B text smoke](docs/qwen3p5_0p8b_smoke.md) | Bounded text pruning, checkpoint reload, IFEval, serving, short KD, and resume | -| [Qwen 3.5 0.8B VLM smoke and campaign](docs/qwen3p5_0p8b_vlm_smoke.md) | Bounded lifecycle smoke and a broader illustrative multi-axis VLM campaign | -| [Qwen 3.5 4B VLM example](docs/qwen3p5_4b_vlm_example.md) | FFN search with opt-in lifecycle and longer campaign routes | +## Choose a recipe -Run the smoke for a model in the same worker environment before starting its -larger campaign. A campaign guide may require more data, GPUs, walltime, or -manual approval than its smoke. +Run `python examples/puzzletron/puzzletron.py routes` to list the maintained +model, workflow, and mode combinations. Five checked-in recipes cover Qwen 3.5 +0.8B text and VLM smokes, a larger 0.8B VLM campaign, and 4B VLM smoke and +campaign routes. See [maintained recipes](docs/maintained_recipes.md) for their +requirements and interpretation limits. -### Evaluate a checkpoint +Recipes are the public run interface. Files under `configs/families/` are +internal composition templates and should not be edited or launched directly. -Evaluation can run inside a post-MIP flow so candidate lineage and metrics stay -with the campaign. See [evaluate saved checkpoints](docs/post_mip_pipeline.md#evaluate-saved-checkpoints) -for that route. +## Custom models -For a standalone local Hugging Face checkpoint, choose the matching evaluator: +If no maintained recipe matches the model, use the existing setup wizard +through the same main command: -| Checkpoint | Guide | Owner of exact commands and outputs | -| --- | --- | --- | -| Text | [Text checkpoint evaluation](docs/checkpoint_evaluation.md) | `python -m examples.puzzletron.evaluation.text` and NeMo Evaluator preparation | -| Qwen 3.5 VLM | [VLM checkpoint evaluation](docs/vlm_checkpoint_evaluation.md) | Benchmark profiles, cache preparation, preflight, execution, and result interpretation | - -Keep evaluator, task contract, dataset revision, sample selection, model -revision, generation settings, and judge identity with every result. Scores -from different evaluator paths or profiles are not interchangeable. - -## Documentation map +```bash +python examples/puzzletron/puzzletron.py setup +``` -The documentation is organized by user task. Follow links from this page -rather than searching configuration or implementation directories. +The wizard creates smoke and production bundles and a README with their launch +commands. It does not submit jobs. Existing `puzzletron_setup.py`, +`puzzletron_setup_v2.py`, and three-file configurations remain supported for +compatibility, but they are not additional maintained-recipe workflows. -### Set up and prepare +## Documentation -- [Environment setup](docs/environment_setup.md) explains the controller and +- [Maintained recipes](docs/maintained_recipes.md): choose and run the five + supported recipe routes. +- [Environment setup](docs/environment_setup.md): configure controller and worker environments. -- [Worker image](docs/worker_image.md) owns image build, export, GPU check, and - identity instructions. -- [Setup wizard](docs/setup_wizard.md) owns profiles, model and dataset inputs, - advanced automation, setup resume, and generated files. - -### Run and recover - -- The [maintained campaign routes](#choose-the-next-task) own tracked - model-specific commands and expected artifacts. -- [Run and recovery](docs/orchestration_operations.md) owns stage selection, - dry-run behavior, progress, interruption, recovery, and execution records. - -### Evaluate and read results - -- [Text checkpoint evaluation](docs/checkpoint_evaluation.md) and - [VLM checkpoint evaluation](docs/vlm_checkpoint_evaluation.md) own standalone - evaluation. -- [Post-MIP pipelines](docs/post_mip_pipeline.md) owns evaluation and other - processing attached to candidate lineage. -- [Campaign reports](docs/campaign_reports.md) owns report regeneration and the - navigation boundary for retained results. Campaign and run leaves own their - own measurements, provenance, reproduction evidence, and limitations. - -### Configure and diagnose - -- [Configuration and overrides](docs/configuration_overrides.md) covers config - layering, output roots, and temporary experiment changes. -- [Slurm configuration](docs/slurm_configuration.md) covers partitions, - CPU-routed stages, logs, and scheduler settings. -- [MIP runs](docs/mip_profiles.md) covers objectives, constraints, search - spaces, variants, and homogeneous search. -- [Sanity validation](docs/sanity_validation.md) explains sorting, width - ranking, physical slicing, warnings, and correctness failures. - -### Understand and extend - -- [Architecture](docs/v2_architecture.md) maps the campaign DAG and components - to their implementation locations. -- [Legacy Nano campaign](docs/legacy_nano_campaign.md) documents the separate - retained Nano workflow; it is not the default Puzzletron v2 route. -- To run a campaign with an agent, ask it to use - [`running-puzzletron`](../../.agents/skills/running-puzzletron/SKILL.md) and - provide the model, dataset, compute environment, search space, resource - constraints, and required downstream stages. +- [Configuration](docs/configuration.md): recipe and site fields, + advanced changes, provenance, and sealed bundles. +- [Run and recovery](docs/orchestration_operations.md): progress, retries, + interruption, inspection, and resume. +- [Evaluation](docs/post_mip_pipeline.md): checkpoint materialization, + evaluation, serving, and distillation stages. +- [Campaign reports](docs/campaign_reports.md): generate and interpret the + cumulative report. + +Additional evaluator and implementation references live under +`examples/puzzletron/docs/`. diff --git a/examples/puzzletron/configs/families/nemotron3/setup_v2_defaults.yaml b/examples/puzzletron/configs/families/nemotron3/setup_v2_defaults.yaml index 8c143a3d286..a48c15aa26e 100644 --- a/examples/puzzletron/configs/families/nemotron3/setup_v2_defaults.yaml +++ b/examples/puzzletron/configs/families/nemotron3/setup_v2_defaults.yaml @@ -1,3 +1,4 @@ +# Defaults for the custom-model setup wizard. Maintained recipes do not read this file. schema_version: 2 profiles: 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 6f127c6a095..f2fae6c15fe 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 @@ -33,9 +33,6 @@ pruning: sort: deferred_axes: [] -vllm_stats: - model_hidden_sizes: [1024, 960, 896] - sort_sanity: eval_samples: 16 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/vlm_smoke.yaml similarity index 100% rename from examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/full_vlm_smoke.yaml rename to examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/vlm_smoke.yaml diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/smoke_base.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/smoke_base.yaml index 359e562064e..59651481a29 100644 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/smoke_base.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/smoke_base.yaml @@ -25,6 +25,7 @@ capability_validation: data: modality: text layout: fixed + revision: ${oc.env:PUZZLETRON_DATASET_REVISION,unpublished} max_sample_length: 512 packing: diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml index 473ac3876c5..edc2e04a717 100644 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml @@ -1,7 +1,7 @@ # @package _global_ defaults: - - mip_vlm_smoke + - vlm_pruning_base - _self_ # Qwen 3.5 4B VLM FFN-width 10%-to-20% KD search. Every MIP diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/full_vlm_smoke.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/full_vlm_smoke.yaml deleted file mode 100644 index febb3a460ad..00000000000 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/full_vlm_smoke.yaml +++ /dev/null @@ -1,6 +0,0 @@ -# @package _global_ - -# Deprecated compatibility alias. Use vlm_lifecycle_smoke. -defaults: - - vlm_lifecycle_smoke - - _self_ diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_campaign.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_campaign.yaml deleted file mode 100644 index d780df55d09..00000000000 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_campaign.yaml +++ /dev/null @@ -1,6 +0,0 @@ -# @package _global_ - -# Deprecated compatibility alias. Use ffn_width_10to20pct_kd_search. -defaults: - - ffn_width_10to20pct_kd_search - - _self_ diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/mip_vlm_smoke.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_pruning_base.yaml similarity index 96% rename from examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/mip_vlm_smoke.yaml rename to examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_pruning_base.yaml index 77f243103f3..bf7a7a0a389 100644 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/mip_vlm_smoke.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_pruning_base.yaml @@ -6,8 +6,7 @@ defaults: - /families/qwen3_5/qwen3p5_4b/model@_global_ - _self_ -# Bounded one-GPU compilation and MIP route for the pinned public checkpoint. -# The default stops before materialization, evaluation, and distillation. +# Shared pruning and MIP settings for maintained lifecycle routes. puzzle_dir: ${oc.env:PUZZLETRON_RUN_ROOT,puzzle_runs/qwen3p5_4b_vlm_smoke} dataset_path: ${oc.env:PUZZLETRON_DATASET_PATH} display_name: Qwen3p5_4B_VLM diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_lifecycle_smoke.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_smoke.yaml similarity index 97% rename from examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_lifecycle_smoke.yaml rename to examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_smoke.yaml index 86337a9ef37..83da9c63edd 100644 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_lifecycle_smoke.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_smoke.yaml @@ -1,10 +1,10 @@ # @package _global_ defaults: - - mip_vlm_smoke + - vlm_pruning_base - _self_ -# Opt-in lifecycle check of FFN pruning through physical materialization, +# Bounded lifecycle check of FFN pruning through physical materialization, # checkpoint reload, bounded RealWorldQA evaluation, and two-step VLM KD. puzzle_dir: ${oc.env:PUZZLETRON_RUN_ROOT} data: 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 bd2c0f79df6..31f759839d0 100644 --- a/examples/puzzletron/configs/families/qwen3_5/setup_v2_defaults.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/setup_v2_defaults.yaml @@ -1,3 +1,4 @@ +# Defaults for the custom-model setup wizard. Maintained recipes do not read this file. schema_version: 2 profiles: diff --git a/examples/puzzletron/configs/orchestration/execution.example.yaml b/examples/puzzletron/configs/orchestration/execution.example.yaml deleted file mode 100644 index 413036096c9..00000000000 --- a/examples/puzzletron/configs/orchestration/execution.example.yaml +++ /dev/null @@ -1,48 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Example execution config for the v2 orchestrator. - -execution: - defaults: - failure_policy: strict - halt_policy: drain - artifact_settling_timeout_seconds: 300 - gpus_per_node: 8 - stages: - # CPU-only stages request no GPUs and can override the runner partition. - # A list lets Slurm select any eligible CPU partition. - convert: - strategy: single - instances: 1 - resource: cpu - partition: - - REPLACE_WITH_PRIMARY_CPU_SLURM_PARTITION - - REPLACE_WITH_ALTERNATE_CPU_SLURM_PARTITION - tokenize_data: - strategy: single - instances: 1 - resource: cpu - partition: REPLACE_WITH_PRIMARY_CPU_SLURM_PARTITION - # Final report generation is always CPU-only and accepts a partition - # override without resource or strategy fields. - final_report: - partition: REPLACE_WITH_PRIMARY_CPU_SLURM_PARTITION - vllm_stats: - strategy: sharded - instances: 16 - width_importance: - strategy: single - instances: 1 - depth_importance: - strategy: persistent_pool - instances: 8 - replacement_scoring: - strategy: persistent_pool - instances: 8 - zero_shot_evaluation: - strategy: sharded - instances: 8 - aiperf: - strategy: sharded - instances: 8 diff --git a/examples/puzzletron/configs/orchestration/execution.single_gpu.yaml b/examples/puzzletron/configs/orchestration/execution.single_gpu.yaml deleted file mode 100644 index e7b40631aef..00000000000 --- a/examples/puzzletron/configs/orchestration/execution.single_gpu.yaml +++ /dev/null @@ -1,16 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Shared one-GPU execution profile for bounded Puzzletron examples. Most stages -# already default to one instance with the strategy implied by their stage or -# post-MIP node type. These two stages otherwise use a persistent worker pool. -execution: - defaults: - failure_policy: strict - halt_policy: drain - gpus_per_node: 1 - stages: - prepare_dataset: {strategy: single, resource: cpu} - mip: {strategy: single, resource: cpu} - depth_importance: {strategy: single} - replacement_scoring: {strategy: single} diff --git a/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_campaign.yaml b/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_campaign.yaml deleted file mode 100644 index b552aecfb8f..00000000000 --- a/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_campaign.yaml +++ /dev/null @@ -1,30 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Parallel execution profile for the illustrative VLM campaign. -execution: - mode: reusable_allocation - defaults: - failure_policy: strict - halt_policy: drain - gpus_per_node: 8 - stages: - prepare_dataset: {strategy: single, resource: cpu} - convert: {strategy: single, resource: cpu} - mip: {strategy: single, resource: cpu} - depth_importance: {strategy: single} - # Keep the available GPUs busy while scoring the configured width scenarios. - # The adapter distributes this configurable worker count across those scenarios. - replacement_scoring: {strategy: persistent_pool, instances: 8} - # Five shards let the retained candidates use disjoint GPUs concurrently. - # Lower this value when a site provides fewer GPUs or memory is constrained. - post.candidates.image_eval: &candidates {instances: 5} - post.candidates.best_image_loss: {strategy: single, resource: cpu} - post.candidates.materialized: *candidates - post.candidates.pre_kd_eval: *candidates - post.candidates.kd: *candidates - post.candidates.post_kd_eval: *candidates - post.candidates.selected: {strategy: single, resource: cpu} - post.candidates.result: {strategy: single, resource: cpu} - # Serving consumes only the selected top-1 candidate. - post.candidates.serving: {instances: 1} diff --git a/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_smoke.yaml b/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_smoke.yaml deleted file mode 100644 index 78648c3fc4a..00000000000 --- a/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_smoke.yaml +++ /dev/null @@ -1,17 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Compact lifecycle smoke inside one reusable two-GPU Slurm allocation. Each -# logical stage keeps its own attempt and log; independent work may share the -# configured GPU capacity without additional scheduler submissions. -execution: - mode: reusable_allocation - defaults: - failure_policy: strict - halt_policy: drain - gpus_per_node: 2 - stages: - prepare_dataset: {strategy: single, resource: cpu} - mip: {strategy: single, resource: cpu} - depth_importance: {strategy: single} - replacement_scoring: {strategy: single} diff --git a/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/runner.slurm.yaml b/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/runner.slurm.yaml deleted file mode 100644 index 15cccbc5d61..00000000000 --- a/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/runner.slurm.yaml +++ /dev/null @@ -1,21 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Portable public half of the Qwen 3.5 0.8B smoke run packet. Replace every -# placeholder with the reviewed site contract before crossing the GPU gate. -runner: - kind: slurm - slurm: - account: REPLACE_WITH_SLURM_ACCOUNT - partition: REPLACE_WITH_SLURM_PARTITION - cpu_cpus_per_task: 4 - cpu_memory_mb: 32768 - max_nodes: 1 - time_limit: "4:00:00" - execution_contract: - repository: /opt/puzzletron/src/modelopt - venv: /venv - container: REPLACE_WITH_PUZZLETRON_IMAGE - container_mounts: REPLACE_WITH_REQUIRED_CONTAINER_MOUNTS - prerun_commands: [] - postrun_commands: [] diff --git a/examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.campaign.yaml b/examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.campaign.yaml deleted file mode 100644 index 5005c045829..00000000000 --- a/examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.campaign.yaml +++ /dev/null @@ -1,21 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Deprecated compatibility alias. Use execution.ffn_width_10to20pct_kd_search.yaml. -execution: - defaults: - failure_policy: strict - halt_policy: drain - # Node capacity used by the orchestrator when packing independent - # candidate instances. Partial-node stages request only their active GPUs. - gpus_per_node: 8 - stages: - width_importance: {strategy: single} - replacement_scoring: {strategy: single} - post.candidate-evaluation.online_eval: &candidate_instances {instances: 4} - post.candidate-evaluation.materialized: *candidate_instances - post.candidate-evaluation.serving: *candidate_instances - post.candidate-evaluation.screening_kd: *candidate_instances - post.candidate-evaluation.screening_eval: *candidate_instances - post.candidate-evaluation.quality_screen: *candidate_instances - final_report: {partition: cpu} diff --git a/examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.full_vlm_smoke.yaml b/examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.full_vlm_smoke.yaml deleted file mode 100644 index 4eed293c3c4..00000000000 --- a/examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.full_vlm_smoke.yaml +++ /dev/null @@ -1,12 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Deprecated compatibility alias. Use execution.vlm_lifecycle_smoke.yaml. -execution: - defaults: - failure_policy: strict - halt_policy: drain - gpus_per_node: 2 - stages: - depth_importance: {strategy: single} - replacement_scoring: {strategy: single} diff --git a/examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.vlm_lifecycle_smoke.yaml b/examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.vlm_lifecycle_smoke.yaml deleted file mode 100644 index 459d68b8e2f..00000000000 --- a/examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.vlm_lifecycle_smoke.yaml +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Two-GPU execution profile for the opt-in 4B lifecycle. The bounded MIP and -# evaluation stages use one GPU; the global-KD node derives its TP2 request -# from the experiment configuration. -execution: - defaults: - failure_policy: strict - halt_policy: drain - gpus_per_node: 2 - stages: - depth_importance: {strategy: single} - replacement_scoring: {strategy: single} diff --git a/examples/puzzletron/configs/orchestration/qwen_moe/execution.production.yaml b/examples/puzzletron/configs/orchestration/qwen_moe/execution.production.yaml deleted file mode 100644 index a563b7d368a..00000000000 --- a/examples/puzzletron/configs/orchestration/qwen_moe/execution.production.yaml +++ /dev/null @@ -1,41 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -execution: - defaults: - failure_policy: strict - halt_policy: drain - gpus_per_node: 8 - stages: - convert: {strategy: single, instances: 1} - tokenize_data: {strategy: single, instances: 1} - # One node (8 GPUs): pack 8 vLLM shards onto that single node. - vllm_stats: {strategy: sharded, instances: 8} - width_importance: {strategy: single, instances: 1} - # One gang-scheduled allocation: four 8-GPU workers, coordinator on node 0. - depth_importance: {strategy: persistent_pool, instances: 4} - sort: {strategy: single, instances: 1} - sort_sanity: {strategy: single, instances: 1} - width_sanity: {strategy: single, instances: 1} - slicing_sanity: {strategy: single, instances: 1} - bypass_sanity: {strategy: single, instances: 1} - bypass: {strategy: single, instances: 1} - build_library: - strategy: single - instances: 1 - gpus_per_node: 1 - parallel: - tp: 1 - cp: 1 - pp: 1 - ep: 1 - dp_shard: 1 - dp_replicate: 1 - # One gang-scheduled allocation: two 8-GPU workers, coordinator on node 0. - replacement_scoring: {strategy: persistent_pool, instances: 2} - mip: {strategy: single, instances: 1} - zero_shot_evaluation: {strategy: sharded, instances: 8} - aiperf: {strategy: sharded, instances: 8} - global_distillation_sanity: {strategy: single, instances: 1} - global_distillation: {strategy: single, instances: 1} - post_distillation_evaluation: {strategy: single, instances: 1} diff --git a/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml b/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml deleted file mode 100644 index 7363d863c92..00000000000 --- a/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml +++ /dev/null @@ -1,24 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -runner: - kind: slurm - slurm: - # Required placeholder. Replace with the Slurm account for your site. - account: REPLACE_WITH_SLURM_ACCOUNT - # Replace with one partition name or a list of eligible names for this site. - partition: REPLACE_WITH_SLURM_PARTITION - max_nodes: 20 - time_limit: "4:00:00" - execution_contract: - # These paths are provided by the Puzzletron worker image. - repository: /opt/puzzletron/src/modelopt - venv: /venv - # Replace with a registry reference or cluster-readable copy of that image. - container: REPLACE_WITH_PUZZLETRON_IMAGE - # Replace with the host and container paths required by the campaign. - container_mounts: "REPLACE_WITH_HOST_PATH:REPLACE_WITH_CONTAINER_PATH" - # Optional site setup, for example cache or authentication variables. - prerun_commands: [] - # Optional shell commands run when the stage payload exits. - postrun_commands: [] diff --git a/examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml b/examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml deleted file mode 100644 index 667e22c465b..00000000000 --- a/examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml +++ /dev/null @@ -1,25 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Example bare-metal runner config for the v2 orchestrator. - -runner: - kind: baremetal - inventory: - # Required placeholders. Replace both host values with SSH hosts reachable - # from the controller; the rendezvous host must match one inventory hostname. - rendezvous_host: REPLACE_WITH_PRIMARY_HOST - rendezvous_port_base: 29500 - hosts: - - hostname: REPLACE_WITH_PRIMARY_HOST - gpus: 8 - - hostname: REPLACE_WITH_SECONDARY_HOST - gpus: 8 - execution_contract: - # Required. Use the same checkout path visible on every SSH host. - repository: REPLACE_WITH_WORKER_VISIBLE_MODELOPT_CHECKOUT - # Required. Sourced as /bin/activate on every SSH host. - venv: REPLACE_WITH_WORKER_VISIBLE_MODELOPT_VENV - # Optional script sourced before virtualenv activation on every host. - # Leave null when the worker login environment needs no additional setup. - setup_env: diff --git a/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml b/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml deleted file mode 100644 index ed449d0f801..00000000000 --- a/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml +++ /dev/null @@ -1,40 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Example Slurm runner config for the v2 orchestrator. - -runner: - kind: slurm - slurm: - # Required placeholder. Replace with the Slurm account for your site. - account: REPLACE_WITH_SLURM_ACCOUNT - # Optional default for stages without a partition override. Use one name or - # a list of eligible names. Slurm uses the site default when this is omitted. - partition: - - REPLACE_WITH_PRIMARY_SLURM_PARTITION - - REPLACE_WITH_ALTERNATE_SLURM_PARTITION - # Bound zero-GPU stages instead of inheriting a site's full-node defaults. - cpu_cpus_per_task: 4 - cpu_memory_mb: 32768 - # Optional fallback for CPU-only stages without their own partition. - partition_cpu: - - REPLACE_WITH_PRIMARY_CPU_SLURM_PARTITION - - REPLACE_WITH_ALTERNATE_CPU_SLURM_PARTITION - time_limit: "4:00:00" - log_dir: logs - execution_contract: - # The Puzzletron worker image provides these paths. Change them only - # when using a custom worker environment. - repository: /opt/puzzletron/src/modelopt - venv: /venv - # Required for the worker image. Replace with a registry - # reference or image path accepted by the site's container plugin. - container: REPLACE_WITH_PUZZLETRON_IMAGE - # Optional and used only with a container. Use comma-separated - # /host/path:/container/path entries, for example /data:/data,/models:/models. - container_mounts: - # Optional shell commands run in order before virtualenv activation. - # Example: ["module load cuda", "export HF_HOME=/data/huggingface"] - prerun_commands: [] - # Optional shell commands run when the stage payload exits. - postrun_commands: [] diff --git a/examples/puzzletron/configs/recipes/qwen3p5_0p8b_text_smoke.yaml b/examples/puzzletron/configs/recipes/qwen3p5_0p8b_text_smoke.yaml new file mode 100644 index 00000000000..138448e5ca9 --- /dev/null +++ b/examples/puzzletron/configs/recipes/qwen3p5_0p8b_text_smoke.yaml @@ -0,0 +1,10 @@ +schema_version: 1 +name: qwen3p5-0p8b-text-smoke +model: qwen3.5-0.8b +workflow: text-pruning +mode: smoke +run_root: puzzle_runs/qwen3p5_0p8b_text_smoke +resource_profile: single-gpu +data: + path: REPLACE_WITH_PREPARED_TEXT_DATASET + revision: REPLACE_WITH_IMMUTABLE_DATASET_REVISION diff --git a/examples/puzzletron/configs/recipes/qwen3p5_0p8b_vlm_campaign.yaml b/examples/puzzletron/configs/recipes/qwen3p5_0p8b_vlm_campaign.yaml new file mode 100644 index 00000000000..89dc781e497 --- /dev/null +++ b/examples/puzzletron/configs/recipes/qwen3p5_0p8b_vlm_campaign.yaml @@ -0,0 +1,7 @@ +schema_version: 1 +name: qwen3p5-0p8b-vlm-campaign +model: qwen3.5-0.8b +workflow: vlm-pruning +mode: campaign +run_root: puzzle_runs/qwen3p5_0p8b_vlm_campaign +resource_profile: campaign diff --git a/examples/puzzletron/configs/recipes/qwen3p5_0p8b_vlm_smoke.yaml b/examples/puzzletron/configs/recipes/qwen3p5_0p8b_vlm_smoke.yaml new file mode 100644 index 00000000000..57fd7aa1dd1 --- /dev/null +++ b/examples/puzzletron/configs/recipes/qwen3p5_0p8b_vlm_smoke.yaml @@ -0,0 +1,7 @@ +schema_version: 1 +name: qwen3p5-0p8b-vlm-smoke +model: qwen3.5-0.8b +workflow: vlm-pruning +mode: smoke +run_root: puzzle_runs/qwen3p5_0p8b_vlm_smoke +resource_profile: smoke diff --git a/examples/puzzletron/configs/recipes/qwen3p5_4b_vlm_campaign.yaml b/examples/puzzletron/configs/recipes/qwen3p5_4b_vlm_campaign.yaml new file mode 100644 index 00000000000..3a61112ba23 --- /dev/null +++ b/examples/puzzletron/configs/recipes/qwen3p5_4b_vlm_campaign.yaml @@ -0,0 +1,10 @@ +schema_version: 1 +name: qwen3p5-4b-vlm-campaign +model: qwen3.5-4b +workflow: vlm-pruning +mode: campaign +run_root: puzzle_runs/qwen3p5_4b_vlm_campaign +resource_profile: campaign +data: + path: REPLACE_WITH_PREPARED_VLM_DATASET + revision: 51f4f4d219315c3283950994d4eb3d7fc30aa87b diff --git a/examples/puzzletron/configs/recipes/qwen3p5_4b_vlm_smoke.yaml b/examples/puzzletron/configs/recipes/qwen3p5_4b_vlm_smoke.yaml new file mode 100644 index 00000000000..298e197b758 --- /dev/null +++ b/examples/puzzletron/configs/recipes/qwen3p5_4b_vlm_smoke.yaml @@ -0,0 +1,10 @@ +schema_version: 1 +name: qwen3p5-4b-vlm-smoke +model: qwen3.5-4b +workflow: vlm-pruning +mode: smoke +run_root: puzzle_runs/qwen3p5_4b_vlm_smoke +resource_profile: smoke +data: + path: REPLACE_WITH_PREPARED_VLM_DATASET + revision: 51f4f4d219315c3283950994d4eb3d7fc30aa87b diff --git a/examples/puzzletron/configs/site.example.yaml b/examples/puzzletron/configs/site.example.yaml new file mode 100644 index 00000000000..e705a38801c --- /dev/null +++ b/examples/puzzletron/configs/site.example.yaml @@ -0,0 +1,36 @@ +schema_version: 1 +site: + kind: slurm + environment: + repository: REPLACE_WITH_WORKER_VISIBLE_MODELOPT_CHECKOUT + venv: REPLACE_WITH_WORKER_VISIBLE_MODELOPT_VENV + container: + container_mounts: + prerun_commands: [] + postrun_commands: [] + paths: + hf_home: REPLACE_WITH_SHARED_HF_HOME + slurm: + account: REPLACE_WITH_SLURM_ACCOUNT + partition: REPLACE_WITH_SLURM_PARTITION + partition_cpu: + cpu_cpus_per_task: 4 + cpu_memory_mb: 32768 + time_limit: "4:00:00" +resources: + single-gpu: + mode: per_attempt + gpus_per_node: 1 + max_nodes: 1 + smoke: + mode: reusable_allocation + gpus_per_node: 2 + max_nodes: 1 + campaign: + mode: reusable_allocation + gpus_per_node: 8 + max_nodes: 1 + multinode: + mode: per_attempt + gpus_per_node: 8 + max_nodes: 64 diff --git a/examples/puzzletron/docs/campaign_reports.md b/examples/puzzletron/docs/campaign_reports.md index ae19218c9cd..9ac6c27f97e 100644 --- a/examples/puzzletron/docs/campaign_reports.md +++ b/examples/puzzletron/docs/campaign_reports.md @@ -1,6 +1,6 @@ # Puzzletron Campaign Reports -After a campaign completes cleanly, `orchestrate.py` attempts to generate a +After a campaign completes cleanly, Puzzletron attempts to generate a cumulative HTML report through the configured runner. A report submission, polling, or artifact failure does not invalidate completed stages, but it is recorded in the run result and the command exits nonzero. Inspect the campaign @@ -45,18 +45,3 @@ reproduction status and unresolved findings below. | Current configuration references | Provide migration or reconstruction starting points, not frozen executed configurations. | | Reproduction status | No reproduction is recorded for the listed reports. | | Support status | Not established while reproduction and unresolved correctness findings remain open. | - -## Future campaign records - -The current entries were curated from retained reports. Future campaigns can -use the same catalog fields, but their evidence artifacts should be generated -by the campaign pipeline from the structured data used to render the HTML -report rather than assembled after the run. - -| Consideration | Future direction | -|---|---| -| Generation | Emit a versioned evidence artifact alongside the HTML report and report manifest. | -| Provenance | Record exact code, model, data, resolved configuration, and override identities. | -| Results | Reference canonical stage outcomes, metrics, warnings, and artifacts without copying them manually. | -| Validation | Make schema and semantic verification part of report publication. | -| Support | Keep model-support promotion as a separate reviewed decision after current-code reproduction and correctness gates pass. | diff --git a/examples/puzzletron/docs/configuration.md b/examples/puzzletron/docs/configuration.md new file mode 100644 index 00000000000..fad2336e55e --- /dev/null +++ b/examples/puzzletron/docs/configuration.md @@ -0,0 +1,99 @@ +# Configuration + +A maintained run has two user-owned files: + +- The recipe selects a maintained model, workflow, mode, run directory, and + named site resource. It may also provide required data identity and explicit + advanced changes. +- The site file owns the worker environment, shared cache, scheduler or + bare-metal hosts, and named resource capacities. Reuse it across runs. + +Run `puzzletron.py routes` to list supported combinations. Files under +`configs/families/` and the internal route catalog supply model facts and stage +defaults. They are implementation inputs, not additional public recipes. + +## Resolution + +Puzzletron validates the recipe and site, selects the internal facts for the +route, applies explicit recipe changes, selects the named site resource, and +compiles the full stage plan. Both public schemas are closed. Validation +rejects duplicate YAML keys, unknown or unused fields, unsupported routes, +no-op changes, settings for inactive stages, invalid meshes, and plans beyond +site capacity. + +Use `explain` to inspect the winning source and allocation for every value: + +```bash +python examples/puzzletron/puzzletron.py explain puzzletron.recipe.yaml \ + --site puzzletron.site.yaml +``` + +Relative `run_root` values resolve from the controller working directory. +Relative `data.path` values resolve inside that run directory. Use absolute +worker-visible paths for shared data stored elsewhere. + +## Advanced changes + +Normal recipes omit implementation settings. Use `advanced.experiment` only +to replace an existing internal value: + +```yaml +advanced: + experiment: + mip.runs.params-90.solver.num_solutions: 4 + pruning.eval_samples: 32 +``` + +Execution changes use the runtime compiler's stage schema: + +```yaml +advanced: + execution: + stages: + post.candidate-evaluation.screening_kd: + parallel: + tp: 8 + pp: 4 + cp: 1 + dp_shard: 1 + dp_replicate: 2 + ep: 1 +``` + +The model-instance GPU count is `TP * PP * CP * DP_SHARD * DP_REPLICATE`. +Expert parallelism overlays the sharded data-parallel dimension. Independent +`instances` multiply the total request. Site-owned partitions and capacities +cannot be changed from a recipe. Accepted changes and their sources are saved +in `provenance.json`. + +## Sealed run bundle + +`dry-run` and `launch` write a content-addressed bundle under +`/orchestration/resolved_bundles/`. It contains the normalized +recipe, runtime and audit experiment snapshots, runner and execution contracts, +compiled plan, provenance, source identity, and file hashes. + +Generated YAML files begin with a `DO NOT EDIT` header. +`experiment.runtime.yaml` is the executable worker input; +`experiment.resolved.yaml` is an audit view. Launch binds a run directory to +one bundle. Resume verifies and reuses that bundle instead of reading current +recipes, site settings, or internal defaults. Change an authored input only for +a new run directory. + +## Custom models and existing configurations + +When no maintained recipe fits, run the existing custom-model wizard through: + +```bash +python examples/puzzletron/puzzletron.py setup +``` + +It creates self-contained smoke and production bundles and does not submit +jobs. Their generated README contains the launch commands. Existing +`puzzletron_setup.py`, `puzzletron_setup_v2.py`, and experiment/runner/execution +configurations remain supported for compatibility. They are not alternative +authoring paths for a new maintained-route run. + +The large-model topology tests validate schema and capacity calculations only; +they are not runtime validation for an unlisted model. Generated bundles live +under ignored run directories and should not be committed. diff --git a/examples/puzzletron/docs/configuration_overrides.md b/examples/puzzletron/docs/configuration_overrides.md deleted file mode 100644 index c6d66e4eb0c..00000000000 --- a/examples/puzzletron/docs/configuration_overrides.md +++ /dev/null @@ -1,52 +0,0 @@ -# Configuration and experiment overrides - -Puzzletron builds experiment settings from reusable YAML files in this -directory: - -```text -examples/puzzletron/configs/ -├── base.yaml # pipeline-wide defaults -└── families/ - └── / - ├── family.yaml # descriptors, hooks, and family axes - └── / - ├── model.yaml # checkpoint metadata and legal domains - └── runs/.yaml # exact named campaign run -``` - -Choose where a built-in campaign stores its outputs without editing the YAML: - -```bash -export PUZZLETRON_RUN_ROOT=/shared/puzzle_runs/my_campaign -``` - -Built-in experiment YAMLs use `PUZZLETRON_RUN_ROOT` as their `puzzle_dir`. -Generated bundles write the selected `puzzle_dir` directly. This directory -contains campaign outputs, manifests, resume information, and logs unless -`runner.slurm.log_dir` sends job logs elsewhere. - -Run `orchestrate.py` with `--dry-run` after any configuration change. It -resolves and validates the experiment, runner, and execution files before job -submission, so misspelled or misplaced fields fail at the command boundary. - -## Command-line overrides - -Use command-line overrides for temporary experiment value changes. Append a -repeatable `--override KEY=VALUE` to the campaign command and inspect the -result with `--dry-run` before launch: - -```bash ---override mip.runs.params-90.solver.num_solutions=4 \ ---override ++runtime_annotations.reason=capacity-check \ ---dry-run -``` - -Plain `KEY=VALUE` and explicit `++KEY=VALUE` both add or replace experiment -values. The `orchestrate.py` command and GPU jobs interpret these forms -identically. -Single-plus add (`+KEY=VALUE`) and delete (`~KEY`) operators are not -supported. Put structural changes in a copied run config so they remain easy to -review. - -Overrides apply only to the experiment config. Edit or copy the runner and -execution files when changing site or scheduler settings. diff --git a/examples/puzzletron/docs/environment_setup.md b/examples/puzzletron/docs/environment_setup.md index e8b566a2b8d..c46671716de 100644 --- a/examples/puzzletron/docs/environment_setup.md +++ b/examples/puzzletron/docs/environment_setup.md @@ -2,18 +2,19 @@ Puzzletron uses two environments: -- one lightweight local Python environment for the setup wizard and campaign +- one lightweight local Python environment for the recipe command and campaign commands; and - a GPU worker environment for ModelOpt, the patched vLLM fork, AutoModel, and AIPerf. -The runner file connects them. `runner.execution_contract.venv` selects the -worker virtual environment, and `runner.execution_contract.container` selects -an optional Slurm container. +The site file connects them. `site.environment.repository` selects the worker +checkout, `site.environment.venv` selects its virtual environment, and +`site.environment.container` selects an optional Slurm container. The resolver +records the worker checkout revision separately from the controller checkout. ## Local Puzzletron environment -The setup wizard and `orchestrate.py` do not import PyTorch or initialize CUDA. +The recipe command does not import PyTorch or initialize CUDA. Use Python 3.10 through 3.14 to create one environment for both. Upgrade the venv's bundled `pip` before resolving the controller dependencies: @@ -25,10 +26,8 @@ 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. +Maintained model entries pin a Hugging Face revision. The resolved run bundle +records that model identity together with the current code revision. Only one local virtual environment is needed for a first campaign. `requirements-setup.txt` includes the packages required to generate, launch, @@ -65,16 +64,57 @@ docker run --gpus all --ipc=host --rm -it \ "${image}" ``` -Inside the image, the runner contract is: +Inside the image, the corresponding `site.environment` values are: - `repository: /opt/puzzletron/src/modelopt` - `venv: /venv` - `container: ` +Puzzletron normally detects the exact Git revision from `repository`. If that +worker-visible path cannot be inspected on the controller, set the optional +`source_revision` field to the worker checkout's full commit SHA. This fallback +is resolver metadata rather than a second runtime setting. The generated runner +checks the expected revision and source state before each per-attempt or SSH +execution and before every logical stage in a reusable allocation. + Add site-specific data, model, cache, and result mounts through `container_mounts`. A registry upload or conversion to a cluster container format changes how the image is delivered, not how its Python environment is created. +Container mounts use the comma-separated syntax expected by Slurm, for example: + +```yaml +site: + environment: + container: /shared/images/puzzletron.sqsh + container_mounts: /shared:/shared,/datasets:/datasets +``` + +For SSH-managed hosts, use the same recipe with a bare-metal site instead. The +executor uses the listed GPU capacity and optional rendezvous address: + +```yaml +schema_version: 1 +site: + kind: baremetal + environment: + repository: /shared/src/Model-Optimizer + venv: /shared/venvs/puzzletron + paths: + hf_home: /shared/cache/huggingface + baremetal: + hosts: + - {hostname: worker-a, gpus: 8} + - {hostname: worker-b, gpus: 8} + rendezvous_host: worker-a + rendezvous_port_base: 29500 +resources: + campaign: + mode: per_attempt + gpus_per_node: 8 + max_nodes: 2 +``` + CI jobs that need the Puzzletron worker stack should use this image and its `/venv`; they should not reinstall a separate environment. diff --git a/examples/puzzletron/docs/maintained_recipes.md b/examples/puzzletron/docs/maintained_recipes.md new file mode 100644 index 00000000000..7ecd7a5a58a --- /dev/null +++ b/examples/puzzletron/docs/maintained_recipes.md @@ -0,0 +1,100 @@ +# Maintained recipes + +Puzzletron provides five maintained recipes. Every recipe uses the same +`puzzletron.py` commands and a reusable site file. Start with a smoke for the +model and modality before running a campaign. + +| Recipe | Purpose | Site resource | +|---|---|---| +| `qwen3p5_0p8b_text_smoke.yaml` | Text pruning lifecycle with bounded evaluation, serving, and distillation | `single-gpu` | +| `qwen3p5_0p8b_vlm_smoke.yaml` | Image-text pruning lifecycle with bounded evaluation, serving, and distillation | `smoke` | +| `qwen3p5_0p8b_vlm_campaign.yaml` | Larger multi-axis image-text integration example | `campaign` | +| `qwen3p5_4b_vlm_smoke.yaml` | 4B FFN-pruning lifecycle with bounded evaluation and distillation | `smoke` | +| `qwen3p5_4b_vlm_campaign.yaml` | 4B FFN-width search and screening example | `campaign` | + +Smoke workloads validate that the configured lifecycle runs, produces usable +checkpoints, and resumes. Their scores and throughput are not model-quality or +production-performance results. The campaigns are scheduled integration +examples, not recommended pruning or training policies. + +## Run a recipe + +Follow the [quickstart](../README.md#quickstart) for the common site, validation, +launch, and resume commands. Set a new worker-visible `run_root` for every run +and fill in `data.path` and `data.revision` when the recipe contains them. + +## Route-specific requirements + +### Qwen 3.5 0.8B text smoke + +Provide a prepared Puzzle-KD dataset and immutable source revision in the +recipe. The materializer writes the required Hugging Face dataset with a +`messages` column: + +```bash +python examples/puzzletron/materialize_dataset.py puzzle_kd_v2 \ + --output /shared/data/puzzle-kd-v2 \ + --train-samples 8 \ + --validation-samples 2 \ + --revision REPLACE_WITH_IMMUTABLE_DATASET_REVISION +``` + +Set `data.path` to that output directory. Workers must have the pinned text +evaluator available. Each scheduled GPU attempt uses one GPU. The smoke checks +FFN-width pruning, checkpoint materialization and reload, two-sample IFEval, a +small serving measurement, two distillation steps, final evaluation, reporting, +and no-work resume. + +### Qwen 3.5 0.8B VLM smoke and campaign + +Workers need the pinned model, Nemotron-VLM data, evaluator datasets and media, +and a shared Hugging Face cache. The smoke prepares eight image conversations +and fixed evaluation rows, then checks pruning, MIP selection, checkpoint +materialization and reload, image-text evaluation, serving, two distillation +steps, reporting, and resume in one reusable two-GPU allocation. + +The campaign uses an eight-GPU allocation and increases the data, candidate, +evaluation, and distillation budgets. It explores hidden width, heterogeneous +FFN width, depth, grouped attention, and GDN geometry. Run it only after the +smoke succeeds. Its selection rule and 128-step distillation budget are example +settings. + +For offline workers, prepare the caches described in [VLM checkpoint +evaluation](vlm_checkpoint_evaluation.md#cache-benchmark-data) before launch. + +### Qwen 3.5 4B VLM smoke and campaign + +Both 4B recipes require a prepared normalized VLM dataset and immutable +revision. Prepare eight samples for the smoke: + +```bash +python examples/puzzletron/materialize_dataset.py nemotron_vlm_v2 \ + --output /shared/data/qwen3p5-vlm \ + --revision 51f4f4d219315c3283950994d4eb3d7fc30aa87b \ + --subsets sparsetables plotqa_cot wiki_en \ + --num-samples 8 \ + --max-shards-per-subset 1 +``` + +Set `data.path` to `/shared/data/qwen3p5-vlm`. The campaign needs at least 64 +samples, so prepare a separate directory with the same command and +`--num-samples 64`, then point the campaign recipe to it. + +The smoke searches aligned FFN widths and checks materialization, checkpoint +reload, bounded RealWorldQA evaluation, two TP2 distillation steps, final +image-text evaluation, reporting, and resume on two colocated GPUs. + +The campaign compares roughly 10%, 15%, and 20% FFN-pruning bands using matched +image-text loss, serving, screening distillation, and evaluation. Its ranking +policy and training budgets are examples. Compare candidates only when their +model, data, evaluator, sampling, generation, teacher, and runtime identities +match. + +## Inspect results + +Use `puzzletron.py inspect ` and the generated campaign report. For a +smoke, confirm that every planned stage completed, materialized checkpoints +reload, reported metrics are finite, and resume submits no completed work. +Use [run and recovery](orchestration_operations.md) for failure handling and +[evaluate saved checkpoints](post_mip_pipeline.md#evaluate-saved-checkpoints) +for evaluator details. diff --git a/examples/puzzletron/docs/orchestration_operations.md b/examples/puzzletron/docs/orchestration_operations.md index 40a72955297..90de948ed20 100644 --- a/examples/puzzletron/docs/orchestration_operations.md +++ b/examples/puzzletron/docs/orchestration_operations.md @@ -1,38 +1,16 @@ # Run and recover campaigns -Run one stage with the same experiment, runner, and execution files used for a -full campaign: - -```bash -PUZZLETRON_BUNDLE=/path/to/generated/campaign/production - -python examples/puzzletron/orchestrate.py \ - --experiment "$PUZZLETRON_BUNDLE/experiment.yaml" \ - --runner "$PUZZLETRON_BUNDLE/runner.yaml" \ - --execution "$PUZZLETRON_BUNDLE/execution.yaml" \ - --stage width_importance -``` - -`--stage ` runs only that stage and requires its parent artifacts to -be complete; it does not run missing prerequisites. Use `--stage full` for the -normal dependency-ordered campaign and whole-campaign resume. - -For repeated manual stage-by-stage work, `per_attempt` is usually the better -fit because Slurm reserves only the resources requested by that attempt. A -selected stage also works in `reusable_allocation` mode, but the complete outer -node remains reserved until that stage finishes or fails, even while some of -its GPUs are unused. Puzzletron then exits the outer job; it does not keep an -idle allocation between separate commands. Site policies that reclaim jobs -with low GPU utilization may still favor `per_attempt`. - -Run the same command with `--dry-run` after changing any input file or updating -the checkout. For Slurm runners, each dry-run submission includes a +The public `launch` and `resume` commands execute the complete sealed dependency +plan. This keeps the recorded plan identical to the launched plan. + +Use the public `dry-run` subcommand after changing the recipe or site file or +updating the checkout. For Slurm sites, each dry-run submission includes a submission-equivalent `sbatch` script with its nested `srun` command. CPU/GPU requests, task launchers, partitions, containers, mounts, and worker commands match a launch. The preview uses a deterministic attempt ID, job name, and log path; a real launch replaces those three identifiers. -When the execution file sets `mode: reusable_allocation`, dry-run instead shows +When the selected resource profile uses `reusable_allocation`, dry-run instead shows one outer Slurm script plus the logical attempts that share it. The public launch command is unchanged. Puzzletron persists the outer handle, reattaches when the same command finds it active, and runs the ordinary controller inside @@ -45,7 +23,7 @@ and exits when the selected plan completes or fails. In reusable mode, the foreground process watches the outer allocation; the inner controller writes its logical-stage progress to the allocation log whose path the watcher prints. -`--once` recovers and polls existing attempts, submits currently ready work, +`launch --once` recovers and polls existing attempts, submits currently ready work, and exits after one scheduling iteration. Submitted jobs keep running. Invoke the same `--once` command again for the next recovery and scheduling iteration. In reusable mode, `--once` submits or reattaches to the one outer allocation and @@ -83,15 +61,12 @@ detach while leaving jobs running, or continue. Non-interactive Ctrl-C and SIGTERM cancel active work and quit. Detaching preserves saved job information, so running the same command recovers the active jobs. -The reusable-allocation watcher has a narrower interruption contract: Ctrl-C detaches and leaves the outer Slurm job running. Rerun the identical command to reattach. Cancel the outer job with the site's normal Slurm command only when you intend to stop all work in that allocation. After Slurm reports the outer job cancelled or failed, rerunning starts a replacement unless a compatible worker result records clean completion or an uncancelled terminal stage failure. Completed logical stages remain skipped, and in-flight local attempts from that allocation are recorded as cancelled before unfinished work is resubmitted. A recorded cancellation and a final-report failure remain retryable. Change the configuration or use a fresh run root before rerunning a campaign with a terminal stage failure. +The reusable-allocation watcher has a narrower interruption contract: Ctrl-C detaches and leaves the outer Slurm job running. Rerun the identical command to reattach. Cancel the outer job with the site's normal Slurm command only when you intend to stop all work in that allocation. After Slurm reports the outer job cancelled or failed, rerunning starts a replacement unless a compatible worker result records clean completion or an uncancelled terminal stage failure. Completed logical stages remain skipped, and in-flight local attempts from that allocation are recorded as cancelled before unfinished work is resubmitted. A recorded cancellation and a final-report failure remain retryable. Use a new run root before rerunning a public recipe with a terminal stage failure. A legacy three-file run can instead change its configuration before rerunning. -Do not edit the experiment, runner, or execution files while a compatible -reusable allocation is active. Let it finish or cancel it, make the change, -and run a fresh dry-run before launching again. A changed configuration has a -new plan identity. In the same run root, Puzzletron can still reuse completed -stages whose recorded outputs remain compatible and valid; affected or -unfinished stages run again. Choose a new run root for a clean independent -campaign. +For a public recipe run, do not edit the generated experiment, runner, or +execution files. Change the recipe or site file and use a new run root; the +existing run root remains bound to its sealed bundle and is resumed with +`puzzletron.py resume`. Redirect stderr before piping through `tee` (for example, append `2>&1 | tee run.log`) so progress output is captured. Use `--color always` for @@ -100,17 +75,18 @@ to change the default five-second poll interval. ## State and execution records -The experiment file calls the campaign output directory `puzzle_dir`. Resume -information is written under `/orchestration/`. The command supports -`single`, `sharded`, and `persistent_pool` strategies, Slurm and SSH executors, -attempt recovery, and semantic stage validation. See the -[`configs/orchestration/`](../configs/orchestration/) directory for starter -runner and execution files. +The recipe calls the campaign output directory `run_root`. Resume information +and the sealed resolved bundle are written under `/orchestration/`. +`resume ` verifies and reuses the sealed experiment, runner, +execution, plan, manifest, and provenance records. The runtime +supports `single`, `sharded`, and `persistent_pool` strategies, Slurm and SSH +executors, attempt recovery, and semantic stage validation. Accepted rank-zero stage results also write checksum-validated execution -records under `/manifests/executions/`. Puzzletron validates these +records under `/manifests/executions/`. Puzzletron validates these records when resuming a stage. They identify existing outputs but do not copy -or make those outputs immutable. +or make those outputs immutable. The configuration bundle is separately +immutable and hash-checked. Dataset preparation additionally records the expected files. Routine resume checks their paths, types, sizes, and timestamps without rereading every image diff --git a/examples/puzzletron/docs/post_mip_pipeline.md b/examples/puzzletron/docs/post_mip_pipeline.md index 6e351a61597..01019d668c7 100644 --- a/examples/puzzletron/docs/post_mip_pipeline.md +++ b/examples/puzzletron/docs/post_mip_pipeline.md @@ -5,10 +5,12 @@ from one MIP run and consists of named, single-input nodes. A node can branch fr any earlier node. Node IDs must be unique across the campaign because they are also stable metric namespaces. -When at least one flow is configured, Puzzletron replaces the -legacy fixed post-MIP stages with these dynamic nodes. Run such campaigns through -`examples/puzzletron/orchestrate.py`; the simple `main.py` stage runner does not -schedule dynamic/manual nodes. +When at least one flow is configured, Puzzletron replaces the legacy fixed +post-MIP stages with these dynamic nodes. Maintained public flows run through a +typed `puzzletron.py` recipe route. Adding another public flow requires adding +such a route. Existing external custom experiment, runner, and execution files +remain supported through `orchestrate.py`; the simple `main.py` stage runner +does not schedule dynamic/manual nodes. ```yaml post_mip: diff --git a/examples/puzzletron/docs/qwen3p5_0p8b_smoke.md b/examples/puzzletron/docs/qwen3p5_0p8b_smoke.md deleted file mode 100644 index 5559202c327..00000000000 --- a/examples/puzzletron/docs/qwen3p5_0p8b_smoke.md +++ /dev/null @@ -1,95 +0,0 @@ -# Qwen 3.5 0.8B text pruning smoke - -The `full_smoke` recipe runs a small end-to-end test of text-only -pruning for Qwen 3.5 0.8B. It searches the FFN intermediate sizes -`[3072, 2048]`, evaluates the candidates, saves the two strongest candidates as -physical checkpoints, and reloads each saved directory through vLLM for two -IFEval samples. It then measures both checkpoints with AIPerf, distills the -candidate with higher measured output-token throughput for two steps, evaluates -the resulting checkpoint with another two IFEval samples and the internal -two-sample LM-loss check, and selects the final checkpoint. The recipe pins the -public checkpoint revision so repeated runs use the same starting model. See -[evaluate saved checkpoints](post_mip_pipeline.md#evaluate-saved-checkpoints) -for how both Hugging Face directories are loaded without an AnyModel-to-AutoModel -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 -3.5 0.8B. Its model-specific defaults generate both smoke and production -bundles covering conversion, pruning, search, MIP selection, materialization, -serving measurement, short distillation, final selection, and a pinned -student-versus-teacher quality comparison. Inspect the generated -`dry-run-plan.txt` and materialize the site-specific runner settings before -launching. The comparison is measurement-only and does not accept the current -scores as a quality baseline. - -The tracked recipes below remain useful as reviewable reference configurations -and for reproducing the bounded GPU comparison. - -## Run and resume the text workflow - -Use the `full_smoke.yaml` experiment with the shared `execution.single_gpu.yaml` -profile. The `runner.slurm.yaml` file is a portable template, not a -runnable site configuration. Copy it to a site-specific location and replace its -`REPLACE_WITH_` values before launching. Dry-run accepts the portable template -for plan inspection, but the orchestrator rejects unresolved placeholders -before submitting work. - -This is the per-attempt Slurm route. Each scheduled GPU attempt requests one -GPU and may run in a separate allocation, so it is not the reusable single-node -VLM smoke. - -```bash -EXPERIMENT=examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/full_smoke.yaml -EXECUTION=examples/puzzletron/configs/orchestration/execution.single_gpu.yaml -RUNNER=/path/to/site-specific/runner.slurm.yaml -export PUZZLETRON_RUN_ROOT=/path/to/qwen3p5_0p8b_full_smoke -``` - -Inspect the complete plan, which uses one GPU per scheduled GPU attempt, -without submitting work: - -```bash -python examples/puzzletron/orchestrate.py \ - --experiment "$EXPERIMENT" \ - --runner "$RUNNER" \ - --execution "$EXECUTION" \ - --stage full --dry-run -``` - -After reviewing the dry-run, launch with the same three inputs and omit only -`--dry-run`: - -```bash -python examples/puzzletron/orchestrate.py \ - --experiment "$EXPERIMENT" \ - --runner "$RUNNER" \ - --execution "$EXECUTION" \ - --stage full -``` - -Resume by rerunning that exact launch command with the same experiment, -materialized runner, execution config, and output root. - -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. - -After completion, inspect the `checkpoint_eval` and `post_kd_checkpoint_eval` -nodes under `artifacts/post_mip/nodes`. Their summaries must name the corresponding -pre-KD and post-KD checkpoints, report two effective IFEval samples, and contain -finite metrics. Also verify the cumulative report and confirm that resuming -submits no work for completed stages. diff --git a/examples/puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md b/examples/puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md deleted file mode 100644 index 1db2c9d185e..00000000000 --- a/examples/puzzletron/docs/qwen3p5_0p8b_vlm_smoke.md +++ /dev/null @@ -1,209 +0,0 @@ -# Qwen 3.5 0.8B VLM pruning example - -The Qwen 3.5 0.8B VLM example has two experiment files with distinct purposes: - -| Experiment | Purpose | Execution profile | -| --- | --- | --- | -| `full_vlm_smoke.yaml` | Check the complete lifecycle in one reusable two-GPU allocation on a single node | `qwen3p5_0p8b/execution.vlm_smoke.yaml` | -| `vlm_campaign.yaml` | Run a longer integration example in one reusable eight-GPU allocation on a single node | `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, -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 intended for scheduled -integration validation, not routine development or presubmit use. It is not a -recommended pruning recipe, training duration, or candidate-selection policy. -The qualifying one-node run completed in about 75 minutes; queue, cache, and -runtime differences can change the elapsed time at another site. - -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, -evaluator, and runtime dependencies. - -## Prerequisites - -Prepare the setup and worker environments described in -[environment setup](environment_setup.md). Workers need access to: - -- the pinned `Qwen/Qwen3.5-0.8B` revision; -- the pinned Nemotron-VLM data revision; -- the evaluator datasets and media described in - [VLM checkpoint evaluation](vlm_checkpoint_evaluation.md); -- a shared campaign output directory. - -If workers cannot access the network, populate those caches before launch and -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. - -The accelerated profiles are qualified for a Slurm environment with eight-GPU -nodes. The smoke reserves two GPUs on one node, while the representative -campaign reserves all eight GPUs on one node. Other Slurm sites may need -adaptation; see [Slurm -configuration](slurm_configuration.md#reusable-single-node-allocations). - -## Run the lifecycle smoke - -Set the shared cache, source identity, and run paths: - -```bash -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" -``` - -`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 its reusable two-GPU, single-node execution -profile and a site-specific runner: - -```bash -EXPERIMENT=examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/full_vlm_smoke.yaml -EXECUTION=examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_smoke.yaml -RUNNER=/path/to/site-specific/runner.slurm.yaml - -python examples/puzzletron/orchestrate.py \ - --experiment "$EXPERIMENT" \ - --runner "$RUNNER" \ - --execution "$EXECUTION" \ - --stage full --dry-run -``` - -Inspect the outer one-node allocation plus every logical attempt, resource, and -stage. Then omit `--dry-run` to launch. Rerun that same launch command to -reattach or resume compatible completed work. - -After completion, inspect the campaign report and verify that: - -- image-text evaluation processed real image tensors; -- the selected checkpoint was physically materialized and reloaded; -- the fixed 24-row RealWorldQA, MMMU, and MVBench evaluation completed before - and after KD; -- the serving requests and two KD steps completed with finite measurements; -- resuming submits no work for compatible completed stages. - -The smoke deliberately uses tiny workloads. Run a separate benchmark with -representative requests before drawing performance conclusions. - -## Run the longer illustrative campaign - -The campaign enables hidden width (the model's shared transformer hidden size), -heterogeneous FFN width (per-block feed-forward intermediate sizes), depth, -grouped-attention geometry, and GDN geometry. All retained candidates use the -same frozen evaluator rows, teacher checkpoint, and 128-step example KD budget. - -The flow is intentionally compact: - -1. Generate parameter-constrained multi-axis candidates. -2. Retain five searched candidates by image-text LM loss. -3. Materialize and evaluate every retained checkpoint before KD. -4. Run 128 KD steps and evaluate again on the same 344 rows. -5. Apply the configured aggregate-rank rule to the searched candidates and - measure serving for the selected result. - -The execution profile uses up to eight resident workers for replacement -scoring, distributed across the configured width scenarios. It shards initial -image evaluation across five tasks, then runs materialization, pre-KD -evaluation, KD, and post-KD evaluation for the five retained candidates across -disjoint GPUs. Each stage runs concurrently up to the configured capacity; -final serving measures only the selected candidate, and individual stages do -not all consume eight GPUs. - -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. - -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" \ - --runner "$RUNNER" \ - --execution "$EXECUTION" \ - --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. - -## Change the example - -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. Change campaign inputs before launch; do not patch -resolved or generated artifacts inside an existing run root. - -In this guide, hidden width means the model's transformer hidden size. -Heterogeneous FFN width is a separate axis that selects per-layer FFN -intermediate sizes; changing one does not implicitly change the other. - -Keep the model revision, frozen evaluator profile, KD exposure, and serving -workload fixed when comparing candidates. Changing any of them creates a -different experiment. - -The campaign evaluates saved checkpoints automatically and records the results -with each candidate. To check a checkpoint outside the campaign, follow -[VLM checkpoint evaluation](vlm_checkpoint_evaluation.md). That guide owns -environment and cache preflight, suite selection, exact commands, result files, -and diagnosis. - -Use guided setup when you need help resolving the model, dataset, and site -settings. The [setup wizard guide](setup_wizard.md) owns its invocation, -profiles, inputs, and generated files. Select Qwen 3.5 0.8B and the -Nemotron-VLM v2 image-text dataset to generate this route with site-specific -settings. Generated bundles use the default `per_attempt` scheduling mode; -follow the setup guide to opt into a reusable single-node allocation. Hidden -(residual/embedding) width, attention, GDN, and depth are available through -guided customization. Among the tracked routes, attention and GDN reductions -are enabled only by `vlm_campaign.yaml`. -Inspect every customized plan with `--dry-run` before launch. See -[configuration and overrides](configuration_overrides.md) for persistent and -temporary changes. diff --git a/examples/puzzletron/docs/qwen3p5_4b_vlm_example.md b/examples/puzzletron/docs/qwen3p5_4b_vlm_example.md deleted file mode 100644 index d675819ca4c..00000000000 --- a/examples/puzzletron/docs/qwen3p5_4b_vlm_example.md +++ /dev/null @@ -1,96 +0,0 @@ -# Qwen 3.5 4B VLM FFN example - -This example defines a bounded Puzzletron route for `Qwen/Qwen3.5-4B`. It pins the public checkpoint revision, spans seven aligned FFN intermediate-size targets from light pruning through roughly 20% whole-model pruning, and separates the default pruning plan from checkpoint evaluation and distillation. - -The default `mip_vlm_smoke.yaml` route: - -- changes the language-model FFN intermediate size from 9216 to one of seven aligned widths from 8704 through 4608; -- uses eight image-conversation samples for width importance and two for sort sanity, width sanity, and replacement scoring; -- compiles the complete teacher-plus-seven-width FFN candidate grid and searches both an approximately 20%-pruned parameter target and an analytical weight-plus-KV serving-memory target; -- stops at MIP without materializing checkpoints, running benchmark evaluation, or starting KD. - -The opt-in `vlm_lifecycle_smoke.yaml` route continues the selected candidate through image-text evaluation, physical checkpoint materialization, a fresh-process RealWorldQA smoke evaluation, two TP2 VLM KD steps, checkpoint reload, and final image-text evaluation. These limits check integration behavior; they do not establish model quality or performance. The old `full_vlm_smoke.yaml` name remains only as a deprecated compatibility alias. - -The extended `ffn_width_10to20pct_kd_search.yaml` route compares roughly 10%, 15%, and 20% FFN-pruning bands. It evaluates the heterogeneous and homogeneous MIP candidates with the same image-text loss, serving, 64-step screening KD, and repeated RealWorldQA/MMMU contract. Aggregate ranking selects one student for an exploratory fresh 256-step KD run initialized from its materialized pre-KD checkpoint, followed by a final matched comparison with the teacher. The screening KD weights are used only for ranking. The 256-step value is a configured follow-up budget, not an established convergence threshold or completion requirement. The old `vlm_campaign.yaml` name remains only as a deprecated compatibility alias. - -## Prepare the inputs - -Prepare the [Puzzletron worker environment](environment_setup.md), then choose worker-visible paths for the normalized dataset and run output. The dataset revision must be an immutable Hugging Face commit SHA. - -```bash -export PUZZLETRON_DATASET_PATH=/path/to/qwen3p5-vlm-smoke-data -export PUZZLETRON_DATASET_REVISION=51f4f4d219315c3283950994d4eb3d7fc30aa87b -export PUZZLETRON_RUN_ROOT=/path/to/qwen3p5_4b_vlm_smoke - -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 -``` - -The source checkpoint is `Qwen/Qwen3.5-4B` at the immutable revision recorded in `configs/families/qwen3_5/qwen3p5_4b/model.yaml`. Cache that revision before launch when workers cannot access Hugging Face. - -## Compile the default plan - -Copy the runner template and replace every `REPLACE_WITH_` value with the reviewed scheduler, repository, environment, container, and mount settings for the target site. - -```bash -EXPERIMENT=examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/mip_vlm_smoke.yaml -EXECUTION=examples/puzzletron/configs/orchestration/execution.single_gpu.yaml -RUNNER_TEMPLATE=examples/puzzletron/configs/orchestration/qwen3p5_4b/runner.slurm.yaml -RUNNER="$PUZZLETRON_RUN_ROOT/runner.slurm.yaml" - -mkdir -p "$PUZZLETRON_RUN_ROOT" -cp "$RUNNER_TEMPLATE" "$RUNNER" -${EDITOR:-vi} "$RUNNER" - -python examples/puzzletron/orchestrate.py \ - --experiment "$EXPERIMENT" \ - --runner "$RUNNER" \ - --execution "$EXECUTION" \ - --stage full --dry-run -``` - -The compiled default plan should end at `mip`. Conversion, library construction, -and MIP request no GPUs; every other enabled stage in this plan requests one -GPU. Inspect the complete plan, paths, and resource settings before removing -`--dry-run`. - -## Compile the opt-in lifecycle - -Use the same dataset and runner with the lifecycle experiment. Switch to the 4B -execution profile because the KD stage needs two colocated GPUs: - -```bash -EXPERIMENT=examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/vlm_lifecycle_smoke.yaml -EXECUTION=examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.vlm_lifecycle_smoke.yaml - -python examples/puzzletron/orchestrate.py \ - --experiment "$EXPERIMENT" \ - --runner "$RUNNER" \ - --execution "$EXECUTION" \ - --stage full --dry-run -``` - -Confirm that the post-MIP order is `image_eval`, `best_vlm_loss`, `materialized`, `checkpoint_eval`, `short_vlm_kd`, `post_kd_checkpoint_eval`, `final_image_eval`, and `best`. `short_vlm_kd` requests two GPUs and uses activation checkpointing, a frozen vision tower, and 64-token KD chunks. `checkpoint_eval` verifies that the physically sliced Hugging Face checkpoint can be loaded by the bounded RealWorldQA evaluator. `post_kd_checkpoint_eval` performs the same reload check on the consolidated KD checkpoint. - -The checked-in CPU tests validate configuration resolution, the full FFN candidate grid, resource counts, and the compiled DAGs. A real run is still required to establish checkpoint compatibility, memory use, finite metrics, and end-to-end behavior on the target GPU and worker image. - -## Compile the extended campaign - -Use the same prepared dataset and reviewed run packet, then compile the longer campaign separately: - -```bash -EXPERIMENT=examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml -EXECUTION=examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.ffn_width_10to20pct_kd_search.yaml - -python examples/puzzletron/orchestrate.py \ - --experiment "$EXPERIMENT" \ - --runner "$RUNNER" \ - --execution "$EXECUTION" \ - --stage full --dry-run -``` - -The campaign deliberately evaluates every emitted MIP candidate before selecting a student. Its benchmark scores are comparable only within runs that use the same pinned evaluator, dataset revisions, sampling policy, generation settings, and teacher checkpoint. diff --git a/examples/puzzletron/docs/setup_wizard.md b/examples/puzzletron/docs/setup_wizard.md deleted file mode 100644 index 968fc8948d2..00000000000 --- a/examples/puzzletron/docs/setup_wizard.md +++ /dev/null @@ -1,142 +0,0 @@ -# Setup wizard - -The Puzzletron setup wizard inspects a local checkpoint configuration or a -Hugging Face model configuration and generates self-contained smoke and -production bundles. It reads configuration metadata, not model weights, and -does not submit jobs. Use it after the maintained lifecycle smoke when you need -to select another model, dataset, search profile, or execution environment. - -## Profiles - -The guided flow offers three profiles: - -- **Quick smoke** creates the smallest campaign for checking campaign shape. -- **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 -model family's `setup_v2_defaults.yaml`. The wizard then asks for the model, -dataset, worker environment, and cluster settings. - -Model-specific defaults can also extend the recommended post-MIP flow. For -Qwen 3.5 0.8B, the detected dataset modality selects the complete post-MIP -route. Text bundles finish with pinned IFEval and GSM8K subsets. Multimodal -bundles use image-aware serving and VLM distillation, then compare the final -student and teacher on pinned RealWorldQA and MMMU subsets. Both comparisons -record measurements and samples without enforcing a quality threshold. - -Adding another size from the same model family does not require another wizard -implementation. Add a model inventory match and model-specific defaults for -its pruning domains, resources, and pinned evaluator, then validate the -generated smoke and production plans for that checkpoint. - -## Models and datasets - -At the **Model** prompt, provide an existing local checkpoint or configuration -path, or a Hugging Face model URL or repository ID. - -At the **Dataset** prompt, provide an existing local dataset path, a Hugging -Face dataset URL, or a repository ID. For a hosted dataset, setup records a -worker-visible output path. The generated campaign `README.md` contains the -exact acquisition command. Run that command from the worker environment before -launching the campaign. A local dataset is referenced directly. - -## Defaults and advanced mode - -Start the guided flow with the example defaults: - -```bash -python examples/puzzletron/puzzletron_setup_v2.py \ - --defaults examples/puzzletron/configs/setup/defaults.example.yaml -``` - -The example uses repository-relative values. Copy it and add site-specific -data, scheduler, and container settings before selecting it. The defaults file -is loaded only when passed explicitly and takes precedence over the selected -profile. If the primary Slurm partition requires GPUs, set -`infrastructure.runner.slurm.partition_cpu` to an eligible CPU partition; the -generated execution routes conversion, tokenization, solve-only MIP, -materialization, and other zero-GPU stages there. Setup-generated named MIP is -always solve-only and runs on CPU; any checkpoint materialization or validation -runs as an explicit later stage. Inspect the generated dry-run for the -authoritative resource and launcher choice. - -Use the full flow to expose every section and nested setting: - -```bash -python examples/puzzletron/puzzletron_setup_v2.py --full -``` - -Automation can use the same setup entry point without answering prompts. The -defaults file must provide every required value that has no resolved default: - -```bash -python examples/puzzletron/puzzletron_setup_v2.py \ - --defaults /path/to/setup-v2-defaults.yaml \ - --campaign-dir /path/to/campaign \ - --profile smoke \ - --non-interactive -``` - -Non-interactive setup fails instead of guessing when a required answer has no -resolved default. It generates and validates the same smoke and production -bundles as the interactive wizard. - -## Navigation and resume - -Press **Esc** to return from any prompt. Selection prompts include a visible -**← Back** action, and text or numeric prompts accept `:back`. - -The wizard saves accepted answers and navigation state in `answers_v2.yaml`. -Resume an interrupted setup with: - -```bash -python examples/puzzletron/puzzletron_setup_v2.py --resume /path/to/campaign -``` - -## Generated files - -The final review writes `resolved_defaults.yaml`, one campaign `README.md`, and -validated `smoke/` and `production/` bundles. The generated README presents -these as **Validate setup** and **Run campaign**. Each bundle contains -experiment, runner, and execution YAML plus a `dry-run-plan.txt`; users do not -need to construct a smoke configuration themselves. The wizard does not submit -either bundle, and the campaign is not automatically gated on validation. - -Generated execution files use Puzzletron's default `per_attempt` mode, which -submits each logical attempt as its own Slurm job. The wizard does not select -`reusable_allocation` automatically because it cannot infer whether the site -supports the required single-node container and shared-storage contract. To use -the accelerated route, set `execution.mode: reusable_allocation` in the -generated execution YAML and run a fresh dry-run. Compilation rejects a stage -whose workers and parallel mesh do not fit within -`execution.defaults.gpus_per_node`. See [reusable single-node -allocations](slurm_configuration.md#reusable-single-node-allocations) before -opting in. - -For named width/depth MIP, setup reads the teacher hidden size and layer count -from the inspected model configuration. It always retains the teacher hidden -size as the full-width scenario, even when the selected search only changes -FFN dimensions. Missing or contradictory teacher geometry stops setup with the -model field that must be fixed. Do not add these values by hand to generated -YAML. - -`dry-run-plan.txt` is a snapshot from the checkout that generated the bundle. -After updating ModelOpt, rerun the generated README's `--dry-run` command before -launching. Incompatible execution schemas and named-MIP bundles fail before -submission and report the setting to change or the setup-resume command that -regenerates both bundles. Named-MIP checks cover the teacher-width scenario, -depth domain, and realized-model validation. Bundles that omit an explicit -resource use the registered CPU-stage defaults. - -The generated configuration can include multiple deployment measurements, independent optimization goals, stage resource profiles, and editable downstream flows. See [experiment overrides](configuration_overrides.md), [Slurm configuration](slurm_configuration.md), and [post-MIP pipelines](post_mip_pipeline.md) for those controls. - -## Generated worker counts - -The wizard writes `instances` for each stage in the generated execution YAML. An instance is an execution worker, not a requested student candidate. Each MIP run requests a solution pool through `solver.num_solutions`; the solver may return fewer solutions, and the post-MIP source and filters determine which distinct candidates continue through the flow. - -Smoke bundles use one instance per stage. Production bundles default non-single stages to the configured `infrastructure.gpus_per_node`, unless a model profile or an advanced stage setting supplies another value. Setup does not generally reduce `instances` to match `num_solutions` or an upstream `top_k`. At runtime, post-MIP `evaluation` and `downstream_evaluation` stages use fewer workers when fewer candidate artifacts actually exist; other GPU candidate stages retain their configured worker count. Lower the advanced stage setting when an upstream filter bounds the input to a smaller set. - -See [stage instances](slurm_configuration.md#stage-instances) for the execution -and recovery behavior of generated worker counts. diff --git a/examples/puzzletron/docs/slurm_configuration.md b/examples/puzzletron/docs/slurm_configuration.md index e7f6975692e..73a12e47ac7 100644 --- a/examples/puzzletron/docs/slurm_configuration.md +++ b/examples/puzzletron/docs/slurm_configuration.md @@ -1,19 +1,25 @@ # Slurm configuration -Use the runner file for site-wide Slurm settings and the execution file for -stage-specific choices. +Use one site file for worker and Slurm settings. Its named resource profiles +provide GPU capacity and scheduling mode. Maintained route defaults own normal +stage behavior; rare stage changes belong under the recipe's explicit +`advanced.execution` section. + +Every Slurm site must set `site.slurm.account`. Container mounts are accepted +only with a container, and bare-metal sites reject container settings because +their executor does not consume them. ## Reusable single-node allocations -The default `execution.mode: per_attempt` submits one Slurm job for each -logical attempt. Set `execution.mode: reusable_allocation` to submit one outer +The resource profile mode `per_attempt` submits one Slurm job for each logical +attempt. Select `mode: reusable_allocation` to submit one outer Slurm job and run the existing dependency-aware controller inside it. This is additive: stage records, attempts, logs, artifacts, failure policies, and resume validation remain separate even though they share one container and allocation. Reusable mode currently supports one node. Compilation rejects any stage whose -resolved topology requires more than one. `execution.defaults.gpus_per_node` +resolved topology requires more than one. The selected profile's `gpus_per_node` sets the outer allocation capacity; logical attempts lease disjoint subsets of the GPUs visible inside that allocation. Stage `instances` and task topology control how much of that capacity each ready stage can use. Capacity is not @@ -26,7 +32,7 @@ This mode assumes a Slurm site where one containerized task can see the full node allocation, shared campaign paths are mounted identically, and runtime caches are writable for the lifetime of the outer job. The outer job uses the site's CPU and memory defaults for its GPU request, which must be sufficient for -the concurrent workers. Other sites may need to adapt the runner's account, +the concurrent workers. Other sites may need to adapt the site's account, partition, container integration, mounts, time limit, GPU capacity, and cache hooks. Use per-attempt mode when those assumptions do not hold. @@ -42,7 +48,12 @@ The meaning depends on the stage strategy: Puzzletron does not generally derive `instances` from MIP `num_solutions` or post-MIP `top_k`. Set both values intentionally when a filter leaves fewer candidates than a downstream GPU stage's configured worker count. -Each instance consumes the stage's resolved GPUs per instance. In reusable mode, compilation rejects a stage that cannot fit within `execution.defaults.gpus_per_node`; the controller starts ready stages only when enough of that capacity is free. For example, five one-GPU candidate instances use at most five GPUs concurrently. Lower `instances` to reduce concurrency or accommodate a larger per-instance parallel mesh. +Each instance consumes the stage's resolved GPUs per instance. In reusable mode, +compilation rejects a stage that cannot fit within the selected resource +profile's `gpus_per_node`; the controller starts ready stages only when enough +of that capacity is free. For example, five one-GPU candidate instances use at +most five GPUs concurrently. Lower `instances` to reduce concurrency or +accommodate a larger per-instance parallel mesh. ### Interruption and resume @@ -50,28 +61,33 @@ If the launching terminal exits while the outer Slurm job is still active, the j If the allocated node fails or the outer job reaches its Slurm time limit, all subprocesses inside that allocation stop. Run the same command again after Slurm reports the job as terminal. Unless the compatible worker result already records clean completion or an uncancelled terminal stage failure, Puzzletron submits a replacement allocation, skips stages with complete validated artifacts, and retries incomplete work. A recorded cancellation and a final report failure remain retryable. An incomplete stage resumes from its own checkpoints only when that stage supports native resume; otherwise its unfinished attempt runs again. -The generated job uses Slurm's no-requeue behavior, so replacement requires a new Puzzletron invocation or external automation. The runner's `runner.slurm.time_limit` covers the complete reusable campaign, not each stage separately. Allow enough headroom for the full campaign, or use per-attempt mode when independent scheduler failure domains are more important than avoiding repeated startup overhead. +The generated job uses Slurm's no-requeue behavior, so replacement requires a +new Puzzletron invocation or external automation. `site.slurm.time_limit` +covers the complete reusable campaign, not each stage separately. Allow enough +headroom for the full campaign, or use per-attempt mode when independent +scheduler failure domains are more important than avoiding repeated startup +overhead. ## Partitions and logs -`runner.slurm.partition` sets the default allocation partition. It accepts one +`site.slurm.partition` sets the default allocation partition. It accepts one partition name or a list of eligible names. Omit it to use the site's Slurm -default. In per-attempt mode, a stage can set -`execution.stages..partition` to one name or its own eligible list. A -reusable allocation has only one GPU partition; compilation rejects differing -GPU-stage overrides, while CPU-stage partition overrides and `partition_cpu` -do not apply to the zero-GPU subprocesses inside that allocation. +default. A selected resource profile may replace this partition for the whole +concise-recipe plan. Stage-specific partitions belong to external or +wizard-generated execution files; `advanced.execution` cannot override +site-owned partition fields. In reusable mode, `site.slurm.partition_cpu` does +not apply to zero-GPU subprocesses inside the outer allocation. -`runner.slurm.log_dir` sets the directory used for every attempt log, including +`site.slurm.log_dir` sets the directory used for every attempt log, including the final-report attempt. When omitted, logs are written below -`/logs`. Relative values are resolved from `puzzle_dir`; absolute +`/logs`. Relative values are resolved from the run root; absolute paths are used as written. -The runner loader accepts `partition_interactive`, `partition_batch`, and +The site loader accepts `partition_interactive`, `partition_batch`, and `interactive_max_nodes` as compatibility fields. They infer stage routing from -role names and node count, which assumes a particular site layout and duplicates -execution-stage settings. Maintained configs use `runner.slurm.partition`, the -supported `partition_cpu` fallback, and stage overrides instead. +role names and node count, which assumes a particular site layout. Maintained +concise-route sites use `site.slurm.partition`, the supported `partition_cpu` +fallback, and resource-profile partitions instead. The production examples also avoid literal `interactive` and `batch` stage overrides because those partition names are not portable between Slurm sites. @@ -97,13 +113,12 @@ Choosing `resource: cpu` for code that actually calls CUDA leaves GPUs hidden an can fail at runtime; choosing `resource: gpu` for CPU work reserves a GPU that the stage may not use. -In per-attempt mode, set `runner.slurm.partition_cpu` when CPU work must use a +In per-attempt mode, set `site.slurm.partition_cpu` when CPU work must use a different partition. Without it, CPU-routed stages request no GPUs on the -runner's default partition. A stage-specific partition remains available when -one CPU stage needs different routing: +site's default partition: ```yaml -runner: +site: kind: slurm slurm: partition: @@ -112,25 +127,15 @@ runner: partition_cpu: - cpu-general - cpu-overflow - -execution: - stages: - convert: - strategy: single - partition: - - cpu-large-memory - width_importance: - strategy: single ``` Slurm selects one partition from each eligible list. The -[`runner.slurm.example.yaml`](../configs/orchestration/runner.slurm.example.yaml) -and [`execution.example.yaml`](../configs/orchestration/execution.example.yaml) -files show the runner default and per-stage CPU routing together. The CPU-only -`final_report` task accepts only a `partition` override. +[`site.example.yaml`](../configs/site.example.yaml) shows the concise public +site and resource-profile contract. Wizard-generated and external execution +files retain their existing stage-partition controls. -For per-attempt CPU stages, `runner.slurm.cpu_cpus_per_task` sets the requested -CPU count and `runner.slurm.cpu_memory_mb` sets memory in MiB. Omit them to use +For per-attempt CPU stages, `site.slurm.cpu_cpus_per_task` sets the requested +CPU count and `site.slurm.cpu_memory_mb` sets memory in MiB. Omit them to use the site's Slurm defaults. These settings are part of resume identity, so changing either causes Puzzletron to submit the stage with the new allocation instead of treating an older active attempt as the same work. They do not @@ -155,18 +160,17 @@ environment or source an access-controlled `setup_env` file. ## Scheduler settings and model settings Do not put `sequence_parallel` under -`execution.stages..parallel`. That mapping controls scheduler allocation +`advanced.execution.stages..parallel`. That mapping controls scheduler allocation and accepts mesh dimensions such as `tp`, `pp`, and `dp_replicate`. -`sequence_parallel` changes model execution and belongs in the experiment's -model-parallel profile. Setup-generated execution files omit it for this -reason. +`sequence_parallel` changes model execution and belongs in the maintained +model or workflow profile, not in scheduler allocation. -Runner and execution files reject unknown fields and suggest the closest valid -name when possible. +Recipe and site files reject unknown fields and suggest the closest valid name +when possible. ## Worker setup hooks -`runner.execution_contract.prerun_commands` and `postrun_commands` are copied +`site.environment.prerun_commands` and `postrun_commands` are copied into generated worker scripts and appear in dry-run output. Puzzletron rejects obvious literal assignments to credential-like variables so those values are not persisted. Inherit credentials from the launch environment, require an @@ -174,10 +178,7 @@ 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. +For containerized workers, use `prerun_commands` when the site needs to place +`TMPDIR` or runtime caches in a short, worker-local writable directory. vLLM +uses Unix-domain sockets with a platform path limit, and a read-only container +home prevents runtime caches from being initialized. diff --git a/examples/puzzletron/docs/v2_architecture.md b/examples/puzzletron/docs/v2_architecture.md index 6464fab6ae5..60f6b3c0ee7 100644 --- a/examples/puzzletron/docs/v2_architecture.md +++ b/examples/puzzletron/docs/v2_architecture.md @@ -1,447 +1,95 @@ -# Puzzletron v2 Architecture +# Puzzletron architecture -> **Scope:** This document maps the Puzzletron v2 design to the checked-in -> implementation. Architecture components and schema capabilities are not model -> support claims. See the [campaign report catalog](campaign_reports.md) for -> retained observations and their reproduction status. +Puzzletron separates a dependency-light controller from GPU-heavy model work. +A recipe and site file select a maintained route. The resolver validates and +seals the full configuration, then the orchestrator compiles and executes a +resumable stage graph. -## Executive summary +Architecture support does not by itself establish that a model, pruning axis, +or topology has been validated end to end. See the [campaign report +catalog](campaign_reports.md) for recorded runs and their evidence status. -Puzzletron v2 turns pruning from a sequence of model-specific scripts into a -distributed, validated, end-to-end campaign. The design is driven by three -goals: - -1. **Scalability**, run large and long-context models with stage-specific - tensor, context, pipeline, expert, data, and sequence parallelism; avoid - repeatedly loading or materializing checkpoints; and distribute independent - work through persistent or sharded workers. -2. **Semantic correctness**, describe every model and pruning axis explicitly, - collect all compatible importance statistics together, and gate the campaign - with sorting, ranking, slicing, bypass, and distillation checks. -3. **End-to-end automation**, generate a campaign from a setup wizard, execute - its dependency graph, search for architectures, run downstream processing, - and continuously assemble a durable HTML report. - -| Concern | Puzzletron v1 | Puzzletron v2 | -|---|---|---| -| Model execution | Hugging Face plus naive pipeline parallelism | AutoModel recipes with stage-specific `TP`, `CP`, `PP`, `EP`, `DP/FSDP/HSDP`, and sequence parallelism | -| Inputs | Fixed-size text samples | Fixed, padded, or packed text and multimodal batches with valid-token/media semantics | -| Width analysis | Primarily FFN width, often one pass per axis | Descriptor-owned hooks for multiple layer types and axes, collected in combined passes | -| Checkpoints | Conversion copies/splits weights; many candidate checkpoints | Conversion links weights; the teacher is sorted once; most candidates are sliced dynamically | -| Bypass | One training run per layer configuration | Nested/Matryoshka sampling across the search space with block or subblock granularity | -| Depth | Search tends to retain no-op choices | Conditional, iterative depth ranking followed by separate MIP scenarios | -| Execution | Stage-specific launch scripts | One scheduler-neutral DAG controller over Slurm, SSH bare metal, or local execution | -| After MIP | Separate manual workflow | Configurable post-MIP node graph for filtering, materialization, evaluation, AIPerf, and global KD | -| Validation | Stage-local checks | Campaign-wide semantic gates with durable manifests and visible warnings/failures | -| Reporting | Results assembled after the fact | A cumulative HTML report regenerated from canonical artifacts throughout the campaign | - -## 1. Component diagram - -This view separates the lightweight control plane from GPU-heavy model work. -The experiment YAML owns algorithm semantics, the runner YAML owns the -environment, and the execution YAML owns placement and failure policy. The -orchestrator binds those contracts without importing the model runtime on a -login node. - -```mermaid -flowchart TB - user["Campaign owner"] - - subgraph experience["Campaign authoring"] - wizard["Setup wizard
model inspection and guided choices"] - contracts["Campaign contracts
experiment + runner + execution YAML"] - end - - subgraph control["Scheduler-neutral control plane"] - registry["Canonical stage registry
dependencies, enablement, artifacts"] - compiler["Plan compiler
mesh validation, instances, packing"] - controller["Durable campaign controller
schedule, poll, retry, resume"] - adapters["Execution adapters
single, sharded, persistent pool"] - executors["Executors
Slurm, SSH bare metal, local"] - end - - subgraph compute["Model and search data plane"] - model["Model abstraction and data
descriptors, capabilities, layouts, modalities"] - importance["Importance and correctness
hooks, sorting, sanity checks"] - candidate["Candidate execution
dynamic slicing, nested bypass, replacement scoring"] - cost["Cost models
parameters, memory, vLLM runtime"] - search["Architecture search
conditional depth + MIP profiles"] - post["Post-MIP processing
filter, materialize, evaluate, AIPerf, global KD"] - end - - subgraph systems["Execution systems"] - automodel["AutoModel
distributed forward, hooks, bypass, KD"] - hf["Hugging Face / AnyModel
model semantics and physical checkpoints"] - vllm["vLLM
exact serving measurements"] - aiperf["AIPerf and evaluators
candidate validation"] - cluster["CPU/GPU resources
Slurm or SSH-managed hosts"] - end - - subgraph artifacts["Durable campaign state"] - store["Canonical artifacts
linked/sorted checkpoints, datasets, scores,
candidate libraries, MIP solutions, observations"] - manifests["Identity and recovery
hashes, manifests, shards, checkpoints, lineage"] - report["Incremental HTML report
DAG state, sanity gates, quality, cost, results"] - end - - user --> wizard --> contracts - contracts --> compiler - registry --> compiler --> controller - controller --> adapters --> executors --> cluster - - adapters --> model - model --> importance - importance --> candidate - candidate --> search - cost --> search - search --> post - - model <--> hf - importance <--> automodel - candidate <--> automodel - cost <--> vllm - post <--> automodel - post <--> aiperf - - importance --> store - candidate --> store - cost --> store - search --> store - post --> store - controller <--> manifests - store --> manifests - manifests --> report - controller --> report - report --> user - - classDef author fill:#e8f1ff,stroke:#4472c4,color:#111; - classDef controlPlane fill:#e8f7ee,stroke:#2e8b57,color:#111; - classDef dataPlane fill:#fff3dc,stroke:#c98716,color:#111; - classDef external fill:#f3eafa,stroke:#7c4d9e,color:#111; - classDef durable fill:#f2f2f2,stroke:#666,color:#111; - - class user,wizard,contracts author; - class registry,compiler,controller,adapters,executors controlPlane; - class model,importance,candidate,cost,search,post dataPlane; - class automodel,hf,vllm,aiperf,cluster external; - class store,manifests,report durable; -``` - -### Key design decisions - -- **One sorted teacher, many logical candidates.** Activation hooks rank every - supported axis, then Puzzletron permutes the teacher once so each prefix is a - valid candidate. Dynamic slicing or masking supports fast experiments, while - physical materialization remains the correctness ground truth. -- **Parallelism belongs to the stage.** Width collection is one coordinated - distributed model; depth uses multiple persistent model instances; vLLM - statistics use independent sharded instances. `EP` overlays the FSDP shard - axis rather than multiplying the GPU allocation. -- **Granularity is explicit.** Bypass, replacement scoring, vLLM statistics, - and depth decisions can operate at block or subblock granularity. -- **Correctness and quality are separate.** Sort and slice equivalence are - correctness gates. Ranking against original-order and reverse controls is a - quality check whose warning remains visible. -- **Artifacts are APIs.** Stages communicate through versioned, hashed, - transactionally published artifacts rather than in-memory coupling. This - enables resume, parallel execution, and report regeneration. - -## 2. Code structure mapped to the component diagram - -The following diagram preserves the component boundaries above while replacing -each conceptual name with its primary implementation locations. Paths are -relative to the repository root. +## Components ```mermaid -flowchart TB - subgraph entry["Campaign authoring"] - setup["Setup wizard
puzzletron_setup/v2/
examples/puzzletron/puzzletron_setup_v2.py"] - config["Config and bundle generation
examples/puzzletron/configs/
puzzletron_setup/v2/bundle.py"] - public["Campaign entry boundary
Public: examples/puzzletron/orchestrate.py
Internal stage worker: examples/puzzletron/main.py"] - end - - subgraph cp["Scheduler-neutral control plane"] - stageGraph["Stage contract
modelopt/torch/puzzletron/stages/graph.py"] - compile["Plan compilation and schemas
orchestration/compiler.py
orchestration/schema.py
orchestration/mesh.py"] - drive["Durable control loop
orchestration/controller.py
orchestration/state.py
orchestration/progress.py"] - adapt["Work decomposition
orchestration/adapters/"] - execute["Resource execution
orchestration/executors/
orchestration/task_launcher.py"] - light["Dependency-light facade
puzzletron_orchestrator/"] - end - - subgraph semantics["Model semantics and input"] - desc["AnyModel descriptors and capabilities
anymodel/model_descriptor/
anymodel/models/
anymodel/capabilities.py"] - autoDesc["AutoModel-native descriptors
anymodel/automodel/"] - data["Fixed, padded, packed, and VLM data
dataset/
utils/data/
plugins/automodel/batch_adapter.py"] - end - - subgraph analysis["Importance, sorting, and sanity"] - hooks["Activation scoring
activation_scoring/
plugins/automodel/hooks/
plugins/automodel/scoring_recipe.py"] - sort["Sorted teacher
pruning/sorted_teacher.py"] - sanity["Semantic checks
stages/diagnostics.py
diagnostics/width_sanity.py
diagnostics/width_slice_equivalence.py"] - end - - subgraph candidates["Candidate execution and costs"] - slice["Dynamic and physical pruning
pruning/dynamic_block_prune.py
pruning/materialize.py
pruning/embedding_pruning.py"] - bypass["Nested bypass
bypass_distillation/
plugins/automodel/local_kd_recipe.py"] - depth["Conditional depth and worker pool
depth/
distributed_eval/"] - replace["Replacement library and scoring
replacement_library/
scoring.py"] - runtime["Parameters, memory, and serving runtime
subblock_stats/
export/vllm.py
utils/vllm_adapter.py"] - end - - subgraph optimize["Search and downstream processing"] - mip["MIP DSL and solver
mip/profiles.py
mip/run_puzzle.py
mip/mip_with_multi_layer_replacements*.py"] - post["Post-MIP node framework
post_mip/base.py
post_mip/builtin.py
post_mip/runner.py
post_mip/records.py"] - kd["Global KD and memory-efficient losses
distillation/global_automodel.py
distillation/flash_kld.py
utils/flash_kd.py"] - bench["Evaluation and performance
benchmarks/
examples/puzzletron/run_profile_*.py"] - end - - subgraph durable["Artifact and reporting plane"] - identity["Identity and transactions
identity.py
manifest.py
checkpoint_transactions.py"] - reports["Cumulative report
diagnostics/campaign_progress_report.py
diagnostics/html_report.py
post_mip/reporting.py"] - end - - setup --> config --> public - public --> light --> compile - stageGraph --> compile --> drive --> adapt --> execute - - adapt --> desc - desc --> autoDesc - desc --> data - data --> hooks --> sort --> sanity - sort --> slice - slice --> bypass - slice --> depth - slice --> replace - slice --> runtime - bypass --> replace - depth --> mip - replace --> mip - runtime --> mip - mip --> post - post --> kd - post --> bench - - drive <--> identity - hooks --> identity - sort --> identity - slice --> identity - bypass --> identity - depth --> identity - replace --> identity - runtime --> identity - mip --> identity - post --> identity - kd --> identity - bench --> identity - stageGraph --> reports - drive --> reports - identity --> reports - - classDef author fill:#e8f1ff,stroke:#4472c4,color:#111; - classDef controlPlane fill:#e8f7ee,stroke:#2e8b57,color:#111; - classDef dataPlane fill:#fff3dc,stroke:#c98716,color:#111; - classDef durablePlane fill:#f2f2f2,stroke:#666,color:#111; - - class setup,config,public author; - class stageGraph,compile,drive,adapt,execute,light controlPlane; - class desc,autoDesc,data,hooks,sort,sanity,slice,bypass,depth,replace,runtime,mip,post,kd,bench dataPlane; - class identity,reports durablePlane; +flowchart LR + inputs["Recipe + site"] --> resolver["Validate, resolve, seal"] + resolver --> compiler["Compile stage graph and resources"] + compiler --> controller["Schedule, monitor, retry, resume"] + controller --> workers["Slurm, bare-metal, or local workers"] + workers --> stages["Prune, search, evaluate, serve, distill"] + stages --> artifacts["Versioned artifacts and report"] + artifacts --> controller ``` -### Component-to-code map - | Component | Primary code | Responsibility | |---|---|---| -| Setup and campaign contracts | `puzzletron_setup/v2/`, `examples/puzzletron/configs/` | Inspect model capabilities, ask relevant questions, and emit validated smoke/production experiment, runner, and execution bundles | -| Stage contract | `modelopt/torch/puzzletron/stages/graph.py` | Define canonical stages, dependencies, completion artifacts, enablement, and distributed behavior | -| Orchestrator | `modelopt/torch/puzzletron/orchestration/` exposed through `puzzletron_orchestrator/` | Compile the DAG, validate meshes, pack instances, execute work, persist state, retry, resume, and monitor | -| Model abstraction | `modelopt/torch/puzzletron/anymodel/` | Describe model structure, supported axes, tensor bindings, export behavior, and HF versus AutoModel-native implementations | -| AutoModel integration | `modelopt/torch/puzzletron/plugins/automodel/` | Implement distributed forward, hook collection, dynamic candidate execution, local bypass/KD, solution scoring, and stage-specific recipes | -| Input layer | `modelopt/torch/puzzletron/dataset/`, `modelopt/torch/puzzletron/utils/data/` | Preserve fixed/padded/packed boundaries, valid tokens, multimodal inputs, and distributed sampling semantics | -| Width ranking and sorting | `activation_scoring/`, `pruning/sorted_teacher.py` | Collect multi-axis statistics and permute coupled tensors into importance order | -| Pruning and materialization | `pruning/` | Apply reversible dynamic slicing and create physically pruned checkpoints for export and correctness checks | -| Nested bypass | `bypass_distillation/` | Sample configurations, train nested candidates, checkpoint durably, and record per-architecture observations | -| Depth ranking | `depth/`, `distributed_eval/` | Iteratively score cumulative removals with a coordinator and persistent worker pool | -| Candidate quality | `replacement_library/`, `scoring.py` | Build valid replacement choices and estimate local quality impact at block or subblock granularity | -| Cost estimation | `subblock_stats/`, `export/vllm.py`, `utils/vllm_adapter.py` | Measure parameter, memory, and exact serving-runtime costs across widths and workloads | -| MIP search | `mip/` | Compile named profiles, constraints, objectives, depth selections, and homogeneous/heterogeneous solutions | -| Post-MIP | `post_mip/` | Execute typed filtering, materialization, evaluation, AIPerf, and global-KD nodes with candidate lineage | -| Global KD | `distillation/` | Run AutoModel teacher/student training, including online-softmax KLD/CE and MTP-aware loss terms | -| Artifacts and report | `identity.py`, `manifest.py`, `checkpoint_transactions.py`, `diagnostics/` | Publish resumable state and render the cumulative campaign report from canonical evidence | +| Public command | `examples/puzzletron/puzzletron.py` | Validate, explain, launch, resume, and inspect maintained recipes | +| Configuration | `modelopt/torch/puzzletron/orchestration/recipe_config.py` | Resolve recipe and site inputs and seal immutable run bundles | +| Stage graph | `modelopt/torch/puzzletron/stages/graph.py` | Define dependencies, enablement, and completion artifacts | +| Orchestrator | `modelopt/torch/puzzletron/orchestration/` | Compile resources, execute work, persist state, and recover | +| Model semantics | `modelopt/torch/puzzletron/anymodel/` | Describe model structure, supported axes, and tensor bindings | +| AutoModel integration | `modelopt/torch/puzzletron/plugins/automodel/` | Run distributed scoring, pruning, bypass, evaluation, and KD | +| Search and downstream work | `modelopt/torch/puzzletron/mip/` and `post_mip/` | Select candidates, materialize, evaluate, benchmark, and distill | +| Evidence | `identity.py`, `manifest.py`, `checkpoint_transactions.py`, and `diagnostics/` | Publish resumable artifacts and cumulative reports | -## 3. Campaign DAG +The compatibility command `orchestrate.py` consumes existing experiment, +runner, and execution files. Maintained recipes reach the same orchestrator +through sealed generated inputs. -The canonical graph exposes parallelism without changing algorithm semantics. -For example, vLLM statistics can start after conversion while depth and width -importance start after tokenization. MIP waits until quality, cost, and depth -evidence are all available. +## Campaign graph + +Routes enable a subset of this graph and may configure their post-MIP nodes +differently: ```mermaid flowchart LR - convert["Convert checkpoint
link teacher weights"] - tokenize["Tokenize / pack data"] - vllm["vLLM runtime statistics"] - depth["Iterative conditional
depth importance"] - width["Multi-axis width
importance hooks"] - sort["Sort teacher once"] - sortcheck["Sort equivalence"] - widthcheck["Ranking quality
sorted vs original vs reverse"] - slicecheck["Dynamic vs physical
slicing equivalence"] - bypasscheck["Bypass overfit checks"] - bypass["Nested bypass"] - library["Build candidate /
replacement library"] - replace["Replacement scoring"] - mip["MIP profiles
quality + cost + depth constraints"] - - subgraph after["Post-MIP candidate graph"] - select["Filter / Pareto selection"] - materialize["Physical materialization"] - eval["Zero-shot / downstream evaluation"] - perf["AIPerf"] - kdcheck["Global KD overfit check"] - kd["Global KD"] - posteval["Post-KD evaluation"] - end - - report["Cumulative campaign report"] - - convert --> tokenize - convert --> vllm - tokenize --> depth - tokenize --> width - width --> sort - sort --> sortcheck --> widthcheck --> slicecheck - sort --> bypasscheck --> bypass - bypass --> library - vllm --> library - library --> replace - vllm --> mip - depth --> mip - replace --> mip - - mip --> select - select --> materialize - materialize --> eval - materialize --> perf - materialize --> kdcheck --> kd --> posteval - - convert -.-> report - widthcheck -.-> report - slicecheck -.-> report - bypass -.-> report - depth -.-> report - vllm -.-> report - mip -.-> report - eval -.-> report - perf -.-> report - kd -.-> report - posteval -.-> report - - classDef required fill:#fff3dc,stroke:#c98716,color:#111; - classDef gate fill:#fde9e7,stroke:#c94c4c,color:#111; - classDef downstream fill:#e8f7ee,stroke:#2e8b57,color:#111; - classDef reportStyle fill:#f2f2f2,stroke:#666,color:#111; - - class convert,tokenize,vllm,depth,width,sort,bypass,library,replace,mip required; - class sortcheck,widthcheck,slicecheck,bypasscheck,kdcheck gate; - class select,materialize,eval,perf,kd,posteval downstream; - class report reportStyle; + convert["Convert checkpoint"] --> data["Prepare data"] + convert --> runtime["Measure runtime costs"] + data --> importance["Collect importance"] + importance --> sort["Sort and validate"] + sort --> candidates["Score candidates"] + runtime --> search["Search"] + candidates --> search + search --> materialize["Materialize"] + materialize --> evaluate["Evaluate"] + materialize --> serve["Measure serving"] + materialize --> kd["Distill"] + kd --> final["Final evaluation"] ``` -Solid arrows are execution dependencies. Dotted arrows are incremental report -updates. Post-MIP nodes are campaign-configurable: they may branch, refer to -metrics from earlier nodes, select an earlier model revision, and include a -durable manual decision gate. - -## Semantic validation gates - -| Gate | Question answered | Outcome | -|---|---|---| -| Capability validation | Does the model support every requested axis, backend, and parallel mode? | Invalid campaign configuration fails before expensive work | -| Sort sanity | Does full-width sorting or reverse sorting preserve teacher behavior? | Difference beyond dtype-aware tolerance is a correctness failure | -| Width sanity | Does the proposed ranking outperform original-order and reverse controls at reduced width? | Poor ranking is a quality finding; warning policy determines whether it also fails the stage | -| Slicing sanity | Does dynamic slicing agree with physical materialization? | Disagreement is a correctness failure | -| Bypass sanity | Can a fixed small candidate and a sampled nested search overfit one batch? | Validates boundaries, gradients, and sampling mechanics | -| Depth evaluation | Are removal scores recomputed after every selected removal? | Produces a conditional trajectory rather than independent linear scores | -| Global KD sanity | Can the student overfit with the configured CE/KLD/MTP loss path? | Validates forward/backward and loss semantics before a long run | -| Artifact completion | Are all expected identities, shards, candidates, and outputs present? | Partial work remains resumable progress, not a completed stage | - -On filesystems with delayed visibility, the controller allows completed work a -bounded interval to publish valid stage artifacts. Advanced deployments can set -`execution.defaults.artifact_settling_timeout_seconds` in the execution config; -the default is 300 seconds. This changes only the post-completion settling -window, not scheduler or stage timeouts. - -These gates answer different questions. Width ranking compares the quality of -different reduced candidates at the same target geometry. Sort and slicing -equivalence compare routes that are supposed to represent the same model -operation. A poor ranking can show that the importance heuristic is not useful -for a case even when every candidate is structurally valid. An equivalence -failure shows that a permutation or runtime slice does not reproduce its -reference implementation, so later measurements cannot be trusted as evidence -for the physical checkpoint. - -The implementation has two stage-completion policies. Correctness failures -always fail their stage. Ranking-quality findings remain warnings by default -and fail the stage when `sanity.fail_on_warnings` is enabled. Scientific, -customer, or release qualification is a third layer and is not currently a -separate Puzzletron verdict. A qualification plan must declare its required -controls, metrics, sample count, tolerances, axis and target coverage, and -aggregation rule. It may reject a campaign for a ranking-quality warning -without reclassifying that warning as a correctness error. The -[sanity validation guide](sanity_validation.md) provides definitions, the -slicing mental model, and a worked example. - -## Current implementation versus design direction - -The list below describes implementation components. It does not imply that every -component, model, axis, modality, or topology combination is end-to-end -demonstrated. Current support requires new reproduced evidence; the -[campaign report catalog](campaign_reports.md) records retained observations -and their evidence status. - -### Implemented in the current v2 code - -- A canonical stage registry and dependency graph. -- A dependency-light setup wizard that emits smoke and production bundles. -- A scheduler-neutral orchestrator with `single`, `sharded`, and - `persistent_pool` execution strategies. -- Slurm, SSH bare-metal, and local executors. -- AutoModel-backed recipes for activation scoring, candidate scoring, bypass, - evaluation, and global KD. -- Descriptor-owned pruning capabilities, multi-axis hooks, sorted teachers, - dynamic slicing, and physical materialization. -- Conditional depth evaluation with persistent distributed workers. -- Named MIP profiles and typed post-MIP nodes for filtering, materialization, - evaluation, AIPerf, and global KD. -- Durable identity, manifest, transaction, resume, lineage, and cumulative - reporting infrastructure. - -### Deliberate next-step convergence - -- Make every pre- and post-MIP operation use one uniform node interface with - dependencies, parallelization policy, artifact contracts, and - `generate_partial_report()` behavior. -- Generate report-only work on the control plane instead of consuming GPU - allocations. -- Expand post-MIP nodes to RL and speculative-decoding workflows such as EAGLE - or DFlash when their contracts are stable. -- Improve embedding-width search so ranking, replacement scoring, and bypass do - not require a full independent sweep for every embedding size. -- Systematically exercise valid parallelism, input-layout, and modality - combinations, then promote only the campaign slices with documented evidence. - -## Manager takeaway - -Puzzletron v2 is not only a larger set of pruning algorithms. It is an -architecture for running pruning as a reproducible distributed campaign: - -- **AutoModel and persistent workers provide scale.** -- **Descriptors, physical materialization, and sanity gates provide semantic - confidence.** -- **The DAG, artifact contracts, post-MIP nodes, and cumulative report provide - end-to-end automation.** - -The remaining architectural opportunity is consolidation: represent all work -as typed nodes so the orchestrator, artifact system, and report generator share -one extensible contract from setup through final model selection. +Stages communicate through versioned, hashed artifacts rather than shared +in-memory state. This permits retries, parallel execution, resume, and report +regeneration. The compiler validates each stage's tensor, pipeline, context, +data, and expert parallel dimensions plus its number of independent instances. + +## Validation gates + +| Gate | What it checks | +|---|---| +| Capability | The model supports every requested axis, backend, and parallel mode | +| Sort sanity | Full-width sorting preserves teacher behavior | +| Width sanity | A reduced-width ranking is compared with original and reverse controls | +| Slice equivalence | Dynamic slicing matches physical materialization | +| Bypass sanity | Nested candidates, gradients, and sampling work on a bounded batch | +| Depth evaluation | Removal scores are recomputed after each selected removal | +| Distillation sanity | The configured student and loss path can train on a bounded batch | +| Artifact completion | Expected identities, shards, candidates, and outputs are complete | + +Correctness failures fail their stage. Ranking-quality findings are warnings by +default and may be configured to fail. Neither result is a release or model +quality verdict. A qualification decision must define its own controls, +metrics, sample counts, tolerances, and aggregation rule. See [sanity +validation](sanity_validation.md) for the detailed checks. + +## Execution properties + +- Puzzletron sorts a teacher once so prefix slicing can represent many logical + candidates; physical materialization remains the checkpoint ground truth. +- Parallelism belongs to each stage. Independent instances multiply a stage's + resource request; expert parallelism overlays the sharded data-parallel + dimension. +- Slurm, SSH-managed bare metal, and local execution share the same compiled + plan contract. +- The controller treats partial work as resumable progress and binds completed + work to artifact identity rather than filenames alone. diff --git a/examples/puzzletron/puzzletron.py b/examples/puzzletron/puzzletron.py new file mode 100644 index 00000000000..31da680d80f --- /dev/null +++ b/examples/puzzletron/puzzletron.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""One public command for Puzzletron configuration and campaign lifecycle.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import yaml + +REPOSITORY_ROOT = Path(__file__).resolve().parents[2] +if str(REPOSITORY_ROOT) not in sys.path: + sys.path.insert(0, str(REPOSITORY_ROOT)) + +from examples.puzzletron import orchestrate # noqa: E402 +from puzzletron_orchestrator.recipe_config import ( # noqa: E402 + ROUTES, + bundle_for_run_root, + explain_resolved_run, + materialize_resolved_bundle, + resolve_recipe_run, +) + + +def _add_recipe_inputs(parser: argparse.ArgumentParser) -> None: + parser.add_argument("recipe", type=Path, help="Concise Puzzletron recipe YAML.") + parser.add_argument("--site", type=Path, required=True, help="Reusable site YAML.") + parser.add_argument( + "--run-root", + type=Path, + help="Override recipe.run_root, normally to create a distinct immutable run.", + ) + + +def _add_run_options(parser: argparse.ArgumentParser, *, dry_run: bool = False) -> None: + if not dry_run: + parser.add_argument( + "--once", + action="store_true", + help="Recover, poll, and submit ready work once, then exit.", + ) + parser.add_argument( + "--max-iterations", + type=int, + help="Stop the controller after this many polling iterations.", + ) + parser.add_argument( + "--expect", + type=Path, + help="Verify completed artifacts against this expectation contract.", + ) + parser.add_argument( + "--color", + choices=("auto", "always", "never"), + default="auto", + help="Colorize progress logs (default: auto).", + ) + parser.add_argument( + "--poll-interval", + type=float, + default=5.0, + help="Seconds between scheduler polls (default: 5).", + ) + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Validate, explain, run, resume, and inspect Puzzletron from recipe + site." + ) + commands = parser.add_subparsers(dest="command", required=True) + + routes = commands.add_parser( + "routes", help="List maintained model, workflow, mode, and intent choices." + ) + routes.add_argument("--json", action="store_true", help="Print one JSON document.") + + commands.add_parser( + "setup", + help="Set up a custom model when no maintained recipe fits.", + ) + + validate = commands.add_parser("validate", help="Validate without writing run artifacts.") + _add_recipe_inputs(validate) + + explain = commands.add_parser( + "explain", help="Show winning sources and the resolved stage mesh." + ) + _add_recipe_inputs(explain) + explain.add_argument("--json", action="store_true", help="Print machine-readable output.") + + dry_run = commands.add_parser("dry-run", help="Seal a preview bundle and render submissions.") + _add_recipe_inputs(dry_run) + _add_run_options(dry_run, dry_run=True) + + launch = commands.add_parser("launch", help="Seal and launch one immutable run bundle.") + _add_recipe_inputs(launch) + _add_run_options(launch) + + resume = commands.add_parser("resume", help="Resume the bundle already bound to a run root.") + resume.add_argument("run_root", type=Path, help="Run root containing the active sealed bundle.") + _add_run_options(resume) + + inspect = commands.add_parser("inspect", help="Inspect a run's sealed inputs and plan.") + inspect.add_argument( + "run_root", type=Path, help="Run root containing the active sealed bundle." + ) + inspect.add_argument("--json", action="store_true", help="Print one JSON document.") + + return parser + + +def _orchestrator_args(bundle: Path, args: argparse.Namespace, *, dry_run: bool) -> list[str]: + command = [ + "--experiment", + str(bundle / "experiment.runtime.yaml"), + "--runner", + str(bundle / "runner.yaml"), + "--execution", + str(bundle / "execution.yaml"), + "--stage", + "full", + "--color", + args.color, + "--poll-interval", + str(args.poll_interval), + ] + if dry_run: + command.append("--dry-run") + if not dry_run: + if args.once: + command.append("--once") + if args.max_iterations is not None: + command.extend(("--max-iterations", str(args.max_iterations))) + if args.expect is not None: + command.extend(("--expect", str(args.expect))) + return command + + +def _inspect(bundle: Path, *, as_json: bool) -> None: + names = ("manifest.json", "plan.json", "provenance.json") + payload = { + name.removesuffix(".json"): json.loads((bundle / name).read_text()) for name in names + } + if as_json: + print(json.dumps({"bundle": str(bundle), **payload}, indent=2)) + return + manifest = payload["manifest"] + plan = payload["plan"] + print(f"bundle: {manifest['bundle_id']}") + print(f"path: {bundle}") + print("runtime input: experiment.runtime.yaml (sealed; do not edit)") + print("audit view: experiment.resolved.yaml (resolved output; do not edit)") + print(f"controller revision: {manifest['code']['controller']['revision']}") + print(f"worker revision: {manifest['code']['worker']['revision']}") + print(f"execution contract: {plan['contract_hash']}") + print(f"stages: {len(plan['stages'])}") + for stage in plan["stages"]: + print( + f" {stage['stage_id']}: {stage['nodes']} node(s), " + f"{stage['total_gpus']} GPU(s), mesh={stage['mesh']}" + ) + + +def main(argv: list[str] | None = None) -> int: + command_line = list(sys.argv[1:] if argv is None else argv) + if command_line[:1] == ["setup"]: + from puzzletron_setup.v2.cli import main as setup_main + + return setup_main(command_line[1:]) + args = _build_parser().parse_args(command_line) + try: + if args.command == "routes": + rows = [ + { + "model": route.model, + "workflow": route.workflow, + "mode": route.mode, + "intent": { + "search": route.search, + "evaluation": route.evaluation, + "distillation": route.distillation, + }, + } + for route in ROUTES + ] + if args.json: + print(json.dumps({"routes": rows}, indent=2)) + else: + for row in rows: + intent = ", ".join(f"{key}={value}" for key, value in row["intent"].items()) + print(f"{row['model']} {row['workflow']} {row['mode']} {intent}") + return 0 + if args.command == "resume": + bundle = bundle_for_run_root(args.run_root) + return orchestrate.main(_orchestrator_args(bundle, args, dry_run=False)) + if args.command == "inspect": + _inspect(bundle_for_run_root(args.run_root), as_json=args.json) + return 0 + + resolved = resolve_recipe_run(args.recipe, args.site, run_root=args.run_root) + if args.command == "validate": + print( + f"valid: {resolved.recipe.name} -> {resolved.route.route_id} " + f"({len(resolved.plan['stages'])} stages, {resolved.bundle_id})" + ) + return 0 + if args.command == "explain": + if args.json: + print( + json.dumps( + { + "bundle_id": resolved.bundle_id, + "plan": resolved.plan, + "provenance": resolved.provenance, + }, + indent=2, + default=str, + ) + ) + else: + print(explain_resolved_run(resolved)) + return 0 + if args.command == "dry-run": + bundle = materialize_resolved_bundle(resolved, activate=False) + return orchestrate.main(_orchestrator_args(bundle, args, dry_run=True)) + if args.command == "launch": + bundle = materialize_resolved_bundle(resolved, activate=True) + return orchestrate.main(_orchestrator_args(bundle, args, dry_run=False)) + raise AssertionError(f"Unhandled command: {args.command}") + except ( + KeyError, + OSError, + RuntimeError, + TypeError, + ValueError, + yaml.YAMLError, + ) as error: + print(f"error: {error}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/puzzletron/puzzletron_setup.py b/examples/puzzletron/puzzletron_setup.py index 5af8b553cb4..41df5259fac 100644 --- a/examples/puzzletron/puzzletron_setup.py +++ b/examples/puzzletron/puzzletron_setup.py @@ -1,6 +1,18 @@ #!/usr/bin/env python3 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """Create a model-aware Puzzletron pruning campaign.""" diff --git a/examples/puzzletron/reports/qwen3p5_0p8b/vlm/qwen35_vlm_smoke/README.md b/examples/puzzletron/reports/qwen3p5_0p8b/vlm/qwen35_vlm_smoke/README.md index 6c8071c5a99..337cb3ec79a 100644 --- a/examples/puzzletron/reports/qwen3p5_0p8b/vlm/qwen35_vlm_smoke/README.md +++ b/examples/puzzletron/reports/qwen3p5_0p8b/vlm/qwen35_vlm_smoke/README.md @@ -10,9 +10,12 @@ See the [exclusive-node run](runs/exclusive_w32_v2/summary.md) for the measured values, study conditions, limitations, structured record, and recorded recipe. The maintained -[campaign config](../../../../configs/families/qwen3_5/qwen3p5_0p8b/runs/vlm_campaign.yaml) +[campaign recipe](../../../../configs/recipes/qwen3p5_0p8b_vlm_campaign.yaml) runs AIPerf only after candidate screening and final KD, with three repetitions of 32 warmup and 64 measured requests per serving cell. Follow the -[campaign guide](../../../../docs/qwen3p5_0p8b_vlm_smoke.md) to configure, -inspect, and run it. The maintained route does not reproduce this historical -pre-KD study, its eight-GPU placement swap, or its 256-request cells. +[maintained recipes](../../../../docs/maintained_recipes.md) to configure, +inspect, and run it. The internal +[route template](../../../../configs/families/qwen3_5/qwen3p5_0p8b/runs/vlm_campaign.yaml) +is retained as historical implementation context. The maintained route does +not reproduce this historical pre-KD study, its eight-GPU placement swap, or +its 256-request cells. diff --git a/examples/puzzletron/reports/qwen3p5_0p8b/vlm/qwen35_vlm_smoke/runs/exclusive_w32_v2/result_record.json b/examples/puzzletron/reports/qwen3p5_0p8b/vlm/qwen35_vlm_smoke/runs/exclusive_w32_v2/result_record.json index 35b7c23cfb2..d7aad156704 100644 --- a/examples/puzzletron/reports/qwen3p5_0p8b/vlm/qwen35_vlm_smoke/runs/exclusive_w32_v2/result_record.json +++ b/examples/puzzletron/reports/qwen3p5_0p8b/vlm/qwen35_vlm_smoke/runs/exclusive_w32_v2/result_record.json @@ -151,7 +151,7 @@ "reproduction": { "campaign_config": "../../../../../../configs/families/qwen3_5/qwen3p5_0p8b/runs/vlm_campaign.yaml", "campaign_config_relationship": "successor_not_execution_equivalent", - "campaign_guide": "../../../../../../docs/qwen3p5_0p8b_vlm_smoke.md", + "campaign_guide": "../../../../../../docs/maintained_recipes.md#qwen-35-08b-vlm-smoke-and-campaign", "serving_recipe": "recipe.json", "observations": "observations.csv", "exact_launcher_retained": false diff --git a/examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/README.md b/examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/README.md index cb39ddf36df..f9224a03e7b 100644 --- a/examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/README.md +++ b/examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/README.md @@ -16,36 +16,29 @@ repeated here. ## Reproduce Use the model and dataset versions recorded in the run's structured record. -Prepare the dataset where workers can read it, and replace every placeholder in -the runner template before launching. +Prepare the dataset where workers can read it, copy the maintained recipe, and +configure the reusable site file before launching. ```bash -export PUZZLETRON_DATASET_PATH=/path/to/qwen3p5-vlm-campaign-data -export PUZZLETRON_DATASET_REVISION=51f4f4d219315c3283950994d4eb3d7fc30aa87b -export PUZZLETRON_RUN_ROOT=/path/to/puzzle_runs/qwen3p5_4b_vlm_campaign - python examples/puzzletron/materialize_dataset.py nemotron_vlm_v2 \ - --output "$PUZZLETRON_DATASET_PATH" \ - --revision "$PUZZLETRON_DATASET_REVISION" \ + --output /path/to/qwen3p5-vlm-campaign-data \ + --revision 51f4f4d219315c3283950994d4eb3d7fc30aa87b \ --subsets sparsetables plotqa_cot wiki_en \ --num-samples 64 \ --max-shards-per-subset 1 -EXPERIMENT=examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml -EXECUTION=examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.ffn_width_10to20pct_kd_search.yaml -RUNNER=examples/puzzletron/configs/orchestration/qwen3p5_4b/runner.slurm.yaml +cp examples/puzzletron/configs/recipes/qwen3p5_4b_vlm_campaign.yaml campaign.recipe.yaml +cp examples/puzzletron/configs/site.example.yaml puzzletron.site.yaml +# Set recipe data.path and run_root, then fill in the site placeholders. -python examples/puzzletron/orchestrate.py \ - --experiment "$EXPERIMENT" \ - --runner "$RUNNER" \ - --execution "$EXECUTION" \ - --stage full --dry-run +python examples/puzzletron/puzzletron.py dry-run campaign.recipe.yaml \ + --site puzzletron.site.yaml ``` Run the dry run first and check the commands, paths, and requested GPUs. The -runner file is only a template until its placeholders are replaced. The full +site file is only a template until its placeholders are replaced. The full recipe also contains a fresh 256-step KD run and a final teacher comparison; run `2026-09-01-r2` stopped before those steps. -See the [Qwen 3.5 4B VLM example](../../../../docs/qwen3p5_4b_vlm_example.md) +See [maintained recipes](../../../../docs/maintained_recipes.md) for environment preparation and lifecycle details. diff --git a/examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.ffn_width_10to20pct_kd_search.yaml b/examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/configuration/execution.yaml similarity index 100% rename from examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.ffn_width_10to20pct_kd_search.yaml rename to examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/configuration/execution.yaml diff --git a/examples/puzzletron/configs/orchestration/qwen3p5_4b/runner.slurm.yaml b/examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/configuration/runner.slurm.yaml similarity index 100% rename from examples/puzzletron/configs/orchestration/qwen3p5_4b/runner.slurm.yaml rename to examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/configuration/runner.slurm.yaml diff --git a/examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/result_record.json b/examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/result_record.json index b971a6dc41d..c9b12ead8aa 100644 --- a/examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/result_record.json +++ b/examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/result_record.json @@ -59,16 +59,18 @@ "configuration": { "experiment": { "path": "examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml", - "packaged_sha256": "01384151e7cc8910e961e5eb238c1e206c9167f9e64c86b57c18eeb7294d906a", + "packaged_sha256": "dbb2016cce98a738b3f03e52d4279745a4b905aeecd3d4c122edbafd8f32b2f8", "execution_equivalent_sha256": "c4e6be62bf55980d23d20a12f35e11261d8dffbe59a1f457c7bccec6cc84ec06", "relationship": "canonical_profile_alias_and_comments_only" }, "execution": { - "path": "examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.ffn_width_10to20pct_kd_search.yaml", - "sha256": "07d460e2894259c43bc4cc446dedd2ad5a27020f92c9de6724ca502e88f823ce" + "path": "examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/configuration/execution.yaml", + "sha256": "62e294807c963f67533e1f6a0a1eeb7f00fdf8b46285d33fedb6fb40bf164095", + "executed_sha256": "07d460e2894259c43bc4cc446dedd2ad5a27020f92c9de6724ca502e88f823ce", + "relationship": "archived maintained successor; exact executed bytes were not retained" }, "runner_template": { - "path": "examples/puzzletron/configs/orchestration/qwen3p5_4b/runner.slurm.yaml", + "path": "examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/configuration/runner.slurm.yaml", "sha256": "f3450eaa14acdb28316ec7aeddcbe8a212bc0ced09c64cf16d97ca7a33ea8191", "relationship": "illustrative_placeholder" } @@ -82,9 +84,9 @@ "--experiment", "examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml", "--runner", - "examples/puzzletron/configs/orchestration/qwen3p5_4b/runner.slurm.yaml", + "examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/configuration/runner.slurm.yaml", "--execution", - "examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.ffn_width_10to20pct_kd_search.yaml", + "examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/configuration/execution.yaml", "--stage", "full", "--dry-run" @@ -95,9 +97,9 @@ "--experiment", "examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml", "--runner", - "examples/puzzletron/configs/orchestration/qwen3p5_4b/runner.slurm.yaml", + "examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/configuration/runner.slurm.yaml", "--execution", - "examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.ffn_width_10to20pct_kd_search.yaml", + "examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/configuration/execution.yaml", "--stage", "full" ], @@ -113,15 +115,16 @@ { "path": "examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml", "role": "validated_successor_experiment_recipe", - "sha256": "01384151e7cc8910e961e5eb238c1e206c9167f9e64c86b57c18eeb7294d906a" + "sha256": "dbb2016cce98a738b3f03e52d4279745a4b905aeecd3d4c122edbafd8f32b2f8" }, { - "path": "examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.ffn_width_10to20pct_kd_search.yaml", + "path": "examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/configuration/execution.yaml", "role": "execution_recipe", - "sha256": "07d460e2894259c43bc4cc446dedd2ad5a27020f92c9de6724ca502e88f823ce" + "sha256": "62e294807c963f67533e1f6a0a1eeb7f00fdf8b46285d33fedb6fb40bf164095", + "executed_sha256": "07d460e2894259c43bc4cc446dedd2ad5a27020f92c9de6724ca502e88f823ce" }, { - "path": "examples/puzzletron/configs/orchestration/qwen3p5_4b/runner.slurm.yaml", + "path": "examples/puzzletron/reports/qwen3p5_4b/vlm/ffn_width_10to20pct_kd_search/runs/2026-09-01-r2/configuration/runner.slurm.yaml", "role": "illustrative_runner_template", "sha256": "f3450eaa14acdb28316ec7aeddcbe8a212bc0ced09c64cf16d97ca7a33ea8191" } diff --git a/modelopt/torch/puzzletron/orchestration/_recipe_inputs.py b/modelopt/torch/puzzletron/orchestration/_recipe_inputs.py new file mode 100644 index 00000000000..fff4a6ffd0a --- /dev/null +++ b/modelopt/torch/puzzletron/orchestration/_recipe_inputs.py @@ -0,0 +1,775 @@ +# 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. + +"""Closed schemas and parsing for human-authored recipes and site files.""" + +from __future__ import annotations + +import re +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +from difflib import get_close_matches +from pathlib import Path +from typing import Any + +import yaml +from yaml.nodes import MappingNode, Node, ScalarNode, SequenceNode + +from ._route_catalog import ROUTES, ROUTES_BY_KEY +from .schema import ExecutionMode + +RECIPE_SCHEMA_VERSION = 1 +DATA_FIELDS = {"path", "revision"} + +_RECIPE_FIELDS = { + "schema_version", + "name", + "model", + "workflow", + "mode", + "run_root", + "resource_profile", + "data", + "advanced", +} +_ADVANCED_FIELDS = {"experiment", "execution"} +_ADVANCED_EXECUTION_FIELDS = {"defaults", "stages"} +_SITE_FIELDS = {"schema_version", "site", "resources"} +_SITE_BODY_FIELDS = {"kind", "environment", "paths", "slurm", "baremetal"} +_PATH_FIELDS = {"hf_home"} +_ENVIRONMENT_FIELDS = { + "repository", + "source_revision", + "venv", + "container", + "container_mounts", + "setup_env", + "prerun_commands", + "postrun_commands", +} +_SLURM_FIELDS = { + "account", + "job_name_prefix", + "partition", + "partition_interactive", + "partition_batch", + "partition_cpu", + "cpu_cpus_per_task", + "cpu_memory_mb", + "interactive_max_nodes", + "time_limit", + "qos", + "log_dir", +} +_BAREMETAL_FIELDS = {"hosts", "rendezvous_host", "rendezvous_port_base"} +_HOST_FIELDS = {"hostname", "gpus"} +_RESOURCE_FIELDS = {"mode", "gpus_per_node", "max_nodes", "partition"} +_ROUTE_OWNED_EXPERIMENT_PATHS = { + "data.revision", + "dataset_path", + "descriptor", + "experiment.dir", + "input_hf_model_path", + "model.descriptor_override", + "model.revision", + "model.source", + "prepare_dataset.evaluation_hf_home", + "prepare_dataset.revision", + "puzzle_dir", +} +_SITE_OWNED_EXECUTION_FIELDS = {"gpus_per_node", "partition"} + + +class _UniqueKeyLoader(yaml.SafeLoader): + """Safe YAML loader that refuses duplicate mapping keys.""" + + +def _construct_unique_mapping( + loader: _UniqueKeyLoader, node: MappingNode, deep: bool = False +) -> dict[Any, Any]: + loader.flatten_mapping(node) + result: dict[Any, Any] = {} + for key_node, value_node in node.value: + key = loader.construct_object(key_node, deep=deep) + if not isinstance(key, (str, int, float, bool, type(None))): + raise TypeError( + f"YAML mapping key at line {key_node.start_mark.line + 1} must be scalar" + ) + if key in result: + raise ValueError(f"Duplicate YAML key {key!r} at line {key_node.start_mark.line + 1}") + result[key] = loader.construct_object(value_node, deep=deep) + return result + + +_UniqueKeyLoader.add_constructor( + yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, + _construct_unique_mapping, +) + + +def _collect_locations( + node: Node, + *, + path: str, + source: Path, + output: dict[str, str], +) -> None: + if isinstance(node, MappingNode): + for key_node, value_node in node.value: + if not isinstance(key_node, ScalarNode): + continue + key = str(key_node.value) + child = f"{path}.{key}" if path else key + output[child] = f"{source}:{key_node.start_mark.line + 1}" + _collect_locations(value_node, path=child, source=source, output=output) + elif isinstance(node, SequenceNode): + for index, value_node in enumerate(node.value): + child = f"{path}[{index}]" + output[child] = f"{source}:{value_node.start_mark.line + 1}" + _collect_locations(value_node, path=child, source=source, output=output) + + +def _load_yaml(path: str | Path) -> tuple[dict[str, Any], dict[str, str]]: + source = Path(path).resolve() + text = source.read_text() + loader = _UniqueKeyLoader(text) + try: + payload = loader.get_single_data() + finally: + loader.dispose() + if not isinstance(payload, Mapping): + raise TypeError(f"YAML root must be a mapping: {source}") + locations: dict[str, str] = {} + root = yaml.compose(text, Loader=yaml.SafeLoader) + if root is not None: + _collect_locations(root, path="", source=source, output=locations) + return dict(payload), locations + + +def location(locations: Mapping[str, str], path: str, fallback: Path) -> str: + current = path + while current: + if current in locations: + return locations[current] + current = current.rpartition(".")[0] + return str(fallback) + + +def _reject_unknown( + payload: Mapping[str, Any], + allowed: set[str], + *, + path: str, + locations: Mapping[str, str], + source: Path, +) -> None: + for key in payload: + if not isinstance(key, str): + raise TypeError(f"{path or 'config'} keys must be strings; got {key!r}") + if key in allowed: + continue + dotted = f"{path}.{key}" if path else key + suggestion = get_close_matches(key, sorted(allowed), n=1) + suffix = f"; did you mean {suggestion[0]!r}?" if suggestion else "" + raise ValueError(f"Unknown field {dotted} at {location(locations, dotted, source)}{suffix}") + + +def required_string(value: Any, *, path: str, location: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise TypeError(f"{path} at {location} must be a non-empty string") + return value.strip() + + +def _positive_int(value: Any, *, path: str, location: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{path} at {location} must be a positive integer") + if value < 1: + raise ValueError(f"{path} at {location} must be at least 1") + return value + + +def _optional_string(value: Any, *, path: str, location: str) -> str | None: + if value is None: + return None + return required_string(value, path=path, location=location) + + +def _string_sequence(value: Any, *, path: str, location: str) -> list[str]: + if value is None: + return [] + if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): + raise TypeError(f"{path} at {location} must be a sequence of command strings") + return [ + required_string(item, path=f"{path}[{index}]", location=location) + for index, item in enumerate(value) + ] + + +def _partition_selector(value: Any, *, path: str, location: str) -> str | list[str] | None: + if value is None: + return None + if isinstance(value, str): + return required_string(value, path=path, location=location) + if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): + raise TypeError(f"{path} at {location} must be a string or sequence of strings") + selected = [ + required_string(item, path=f"{path}[{index}]", location=location) + for index, item in enumerate(value) + ] + if not selected: + raise ValueError(f"{path} at {location} must not be empty") + return selected + + +def mapping(value: Any, *, path: str, location: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"{path} at {location} must be a mapping") + return dict(value) + + +@dataclass(frozen=True) +class Recipe: + source: Path + locations: Mapping[str, str] + name: str + model: str + workflow: str + mode: str + run_root: Path + run_root_source: str + resource_profile: str + data: Mapping[str, str] + advanced_experiment: Mapping[str, Any] + advanced_execution: Mapping[str, Any] + + def as_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "schema_version": RECIPE_SCHEMA_VERSION, + "name": self.name, + "model": self.model, + "workflow": self.workflow, + "mode": self.mode, + "run_root": str(self.run_root), + "resource_profile": self.resource_profile, + } + if self.data: + payload["data"] = dict(self.data) + advanced: dict[str, Any] = {} + if self.advanced_experiment: + advanced["experiment"] = dict(self.advanced_experiment) + if self.advanced_execution: + advanced["execution"] = deepcopy(dict(self.advanced_execution)) + if advanced: + payload["advanced"] = advanced + return payload + + +@dataclass(frozen=True) +class Site: + source: Path + locations: Mapping[str, str] + body: Mapping[str, Any] + resources: Mapping[str, Mapping[str, Any]] + + +def parse_recipe(path: str | Path, *, run_root: str | Path | None) -> Recipe: + source = Path(path).resolve() + payload, locations = _load_yaml(source) + _reject_unknown(payload, _RECIPE_FIELDS, path="", locations=locations, source=source) + version = payload.get("schema_version") + if version != RECIPE_SCHEMA_VERSION: + raise ValueError( + f"schema_version at {location(locations, 'schema_version', source)} must be " + f"{RECIPE_SCHEMA_VERSION}; got {version!r}" + ) + values = { + key: required_string(payload.get(key), path=key, location=location(locations, key, source)) + for key in ("name", "model", "workflow", "mode", "resource_profile") + } + selected_run_root = str(run_root) if run_root is not None else payload.get("run_root") + run_root_text = required_string( + selected_run_root, + path="run_root", + location=( + "command line" if run_root is not None else location(locations, "run_root", source) + ), + ) + if "$" in run_root_text: + raise ValueError("run_root does not support environment interpolation; use --run-root") + data = mapping( + payload.get("data", {}), + path="data", + location=location(locations, "data", source), + ) + _reject_unknown(data, DATA_FIELDS, path="data", locations=locations, source=source) + normalized_data = { + key: required_string( + value, + path=f"data.{key}", + location=location(locations, f"data.{key}", source), + ) + for key, value in data.items() + } + if "$" in normalized_data.get("path", ""): + raise ValueError("data.path does not support environment interpolation") + resolved_run_root = Path(run_root_text).expanduser().resolve() + if "path" in normalized_data: + data_path = Path(normalized_data["path"]).expanduser() + if not data_path.is_absolute(): + data_path = resolved_run_root / data_path + normalized_data["path"] = str(data_path) + advanced = mapping( + payload.get("advanced") or {}, + path="advanced", + location=location(locations, "advanced", source), + ) + _reject_unknown(advanced, _ADVANCED_FIELDS, path="advanced", locations=locations, source=source) + advanced_experiment = mapping( + advanced.get("experiment", {}), + path="advanced.experiment", + location=location(locations, "advanced.experiment", source), + ) + for dotted in advanced_experiment: + if not isinstance(dotted, str) or not dotted or any(not part for part in dotted.split(".")): + raise ValueError( + f"advanced.experiment keys must be non-empty dotted paths; got {dotted!r}" + ) + overlaps_owned_path = any( + dotted == owned or dotted.startswith(f"{owned}.") or owned.startswith(f"{dotted}.") + for owned in _ROUTE_OWNED_EXPERIMENT_PATHS + ) + if ( + overlaps_owned_path + or dotted == "model_info" + or dotted.startswith("model_info.") + or dotted.endswith(".evaluator_revision") + ): + raise ValueError( + f"advanced.experiment.{dotted} is owned by the selected route, recipe, " + "or site and cannot be overridden" + ) + advanced_execution = mapping( + advanced.get("execution", {}), + path="advanced.execution", + location=location(locations, "advanced.execution", source), + ) + _reject_unknown( + advanced_execution, + _ADVANCED_EXECUTION_FIELDS, + path="advanced.execution", + locations=locations, + source=source, + ) + for execution_field in _ADVANCED_EXECUTION_FIELDS: + if execution_field not in advanced_execution: + continue + execution_values = mapping( + advanced_execution[execution_field], + path=f"advanced.execution.{execution_field}", + location=location(locations, f"advanced.execution.{execution_field}", source), + ) + if execution_field == "defaults": + reserved = sorted(_SITE_OWNED_EXECUTION_FIELDS.intersection(execution_values)) + if reserved: + raise ValueError( + "advanced.execution.defaults cannot override site-owned fields: " + + ", ".join(reserved) + ) + continue + for stage_id, raw_stage in execution_values.items(): + if not isinstance(raw_stage, Mapping): + continue + reserved = sorted(_SITE_OWNED_EXECUTION_FIELDS.intersection(raw_stage)) + if reserved: + raise ValueError( + f"advanced.execution.stages.{stage_id} cannot override site-owned fields: " + + ", ".join(reserved) + ) + return Recipe( + source=source, + locations=locations, + name=values["name"], + model=values["model"], + workflow=values["workflow"], + mode=values["mode"], + run_root=resolved_run_root, + run_root_source=( + "command line --run-root" + if run_root is not None + else location(locations, "run_root", source) + ), + resource_profile=values["resource_profile"], + data=normalized_data, + advanced_experiment=advanced_experiment, + advanced_execution=advanced_execution, + ) + + +def parse_site(path: str | Path) -> Site: + source = Path(path).resolve() + payload, locations = _load_yaml(source) + _reject_unknown(payload, _SITE_FIELDS, path="", locations=locations, source=source) + if payload.get("schema_version") != RECIPE_SCHEMA_VERSION: + raise ValueError( + f"schema_version at {location(locations, 'schema_version', source)} must be " + f"{RECIPE_SCHEMA_VERSION}; got {payload.get('schema_version')!r}" + ) + site = mapping(payload.get("site"), path="site", location=location(locations, "site", source)) + _reject_unknown(site, _SITE_BODY_FIELDS, path="site", locations=locations, source=source) + kind = required_string( + site.get("kind"), path="site.kind", location=location(locations, "site.kind", source) + ) + if kind not in {"slurm", "baremetal"}: + raise ValueError(f"site.kind must be 'slurm' or 'baremetal'; got {kind!r}") + environment = mapping( + site.get("environment"), + path="site.environment", + location=location(locations, "site.environment", source), + ) + _reject_unknown( + environment, + _ENVIRONMENT_FIELDS, + path="site.environment", + locations=locations, + source=source, + ) + normalized_environment = { + required: required_string( + environment.get(required), + path=f"site.environment.{required}", + location=location(locations, f"site.environment.{required}", source), + ) + for required in ("repository", "venv") + } + for optional in ("source_revision", "container", "container_mounts", "setup_env"): + if optional in environment: + normalized_environment[optional] = _optional_string( + environment[optional], + path=f"site.environment.{optional}", + location=location(locations, f"site.environment.{optional}", source), + ) + for commands in ("prerun_commands", "postrun_commands"): + if commands in environment: + normalized_environment[commands] = _string_sequence( + environment[commands], + path=f"site.environment.{commands}", + location=location(locations, f"site.environment.{commands}", source), + ) + environment = normalized_environment + if environment.get("container_mounts") and not environment.get("container"): + raise ValueError( + "site.environment.container_mounts is unused without site.environment.container" + ) + if kind == "baremetal" and ( + environment.get("container") or environment.get("container_mounts") + ): + raise ValueError( + "site.environment container settings are unused when site.kind is 'baremetal'" + ) + paths = mapping( + site.get("paths", {}), + path="site.paths", + location=location(locations, "site.paths", source), + ) + _reject_unknown(paths, _PATH_FIELDS, path="site.paths", locations=locations, source=source) + paths = { + name: required_string( + value, + path=f"site.paths.{name}", + location=location(locations, f"site.paths.{name}", source), + ) + for name, value in paths.items() + } + if "hf_home" not in paths: + raise ValueError("site.paths.hf_home is required") + hf_home_assignment = re.compile(r"(?:^|\s)(?:export\s+)?HF_HOME\s*=") + for index, command in enumerate(environment.get("prerun_commands", ())): + if hf_home_assignment.search(command): + raise ValueError( + f"site.environment.prerun_commands[{index}] assigns HF_HOME, which conflicts " + "with site.paths.hf_home" + ) + slurm = mapping( + site.get("slurm", {}), + path="site.slurm", + location=location(locations, "site.slurm", source), + ) + baremetal = mapping( + site.get("baremetal", {}), + path="site.baremetal", + location=location(locations, "site.baremetal", source), + ) + _reject_unknown(slurm, _SLURM_FIELDS, path="site.slurm", locations=locations, source=source) + _reject_unknown( + baremetal, + _BAREMETAL_FIELDS, + path="site.baremetal", + locations=locations, + source=source, + ) + if kind == "slurm" and baremetal: + raise ValueError("site.baremetal is unused when site.kind is 'slurm'") + if kind == "baremetal" and slurm: + raise ValueError("site.slurm is unused when site.kind is 'baremetal'") + if kind == "baremetal": + hosts = baremetal.get("hosts") + if isinstance(hosts, (str, bytes)) or not isinstance(hosts, Sequence) or not hosts: + raise TypeError("site.baremetal.hosts must be a non-empty sequence") + for index, raw_host in enumerate(hosts): + host = mapping( + raw_host, + path=f"site.baremetal.hosts[{index}]", + location=location(locations, f"site.baremetal.hosts[{index}]", source), + ) + _reject_unknown( + host, + _HOST_FIELDS, + path=f"site.baremetal.hosts[{index}]", + locations=locations, + source=source, + ) + required_string( + host.get("hostname"), + path=f"site.baremetal.hosts[{index}].hostname", + location=location(locations, f"site.baremetal.hosts[{index}].hostname", source), + ) + _positive_int( + host.get("gpus", 8), + path=f"site.baremetal.hosts[{index}].gpus", + location=location(locations, f"site.baremetal.hosts[{index}].gpus", source), + ) + if "rendezvous_host" in baremetal: + baremetal["rendezvous_host"] = _optional_string( + baremetal["rendezvous_host"], + path="site.baremetal.rendezvous_host", + location=location(locations, "site.baremetal.rendezvous_host", source), + ) + if "rendezvous_port_base" in baremetal: + baremetal["rendezvous_port_base"] = _positive_int( + baremetal["rendezvous_port_base"], + path="site.baremetal.rendezvous_port_base", + location=location(locations, "site.baremetal.rendezvous_port_base", source), + ) + normalized_slurm: dict[str, Any] = {} + if kind == "slurm": + normalized_slurm["account"] = required_string( + slurm.get("account"), + path="site.slurm.account", + location=location(locations, "site.slurm.account", source), + ) + for name in ("job_name_prefix", "time_limit"): + if name in slurm: + normalized_slurm[name] = required_string( + slurm[name], + path=f"site.slurm.{name}", + location=location(locations, f"site.slurm.{name}", source), + ) + for name in ("qos", "log_dir"): + if name in slurm: + normalized_slurm[name] = _optional_string( + slurm[name], + path=f"site.slurm.{name}", + location=location(locations, f"site.slurm.{name}", source), + ) + for name in ("partition", "partition_interactive", "partition_batch", "partition_cpu"): + if name in slurm: + normalized_slurm[name] = _partition_selector( + slurm[name], + path=f"site.slurm.{name}", + location=location(locations, f"site.slurm.{name}", source), + ) + for name in ("cpu_cpus_per_task", "cpu_memory_mb", "interactive_max_nodes"): + if name in slurm: + normalized_slurm[name] = _positive_int( + slurm[name], + path=f"site.slurm.{name}", + location=location(locations, f"site.slurm.{name}", source), + ) + slurm = normalized_slurm + resources = mapping( + payload.get("resources"), + path="resources", + location=location(locations, "resources", source), + ) + if not resources: + raise ValueError("resources must define at least one named resource profile") + normalized_resources: dict[str, dict[str, Any]] = {} + for name, raw_resource in resources.items(): + if not isinstance(name, str) or not name.strip(): + raise TypeError("resource profile names must be non-empty strings") + resource = mapping( + raw_resource, + path=f"resources.{name}", + location=location(locations, f"resources.{name}", source), + ) + _reject_unknown( + resource, + _RESOURCE_FIELDS, + path=f"resources.{name}", + locations=locations, + source=source, + ) + mode = required_string( + resource.get("mode", "per_attempt"), + path=f"resources.{name}.mode", + location=location(locations, f"resources.{name}.mode", source), + ) + if mode not in {item.value for item in ExecutionMode}: + raise ValueError(f"resources.{name}.mode has unsupported value {mode!r}") + gpus_per_node = _positive_int( + resource.get("gpus_per_node"), + path=f"resources.{name}.gpus_per_node", + location=location(locations, f"resources.{name}.gpus_per_node", source), + ) + max_nodes = _positive_int( + resource.get("max_nodes"), + path=f"resources.{name}.max_nodes", + location=location(locations, f"resources.{name}.max_nodes", source), + ) + if kind == "baremetal": + if mode == ExecutionMode.REUSABLE_ALLOCATION.value: + raise ValueError( + f"resources.{name}.mode reusable_allocation is only available for Slurm" + ) + if "partition" in resource: + raise ValueError(f"resources.{name}.partition is unused for bare-metal sites") + hosts = list(baremetal["hosts"]) + if max_nodes > len(hosts): + raise ValueError( + f"resources.{name}.max_nodes={max_nodes} exceeds the {len(hosts)} " + "bare-metal hosts" + ) + undersized = [ + str(host["hostname"]) + for host in hosts[:max_nodes] + if int(host.get("gpus", 8)) < gpus_per_node + ] + if undersized: + raise ValueError( + f"resources.{name}.gpus_per_node={gpus_per_node} exceeds host capacity: " + + ", ".join(undersized) + ) + normalized_resources[name] = { + "mode": mode, + "gpus_per_node": gpus_per_node, + "max_nodes": max_nodes, + **( + { + "partition": _partition_selector( + resource["partition"], + path=f"resources.{name}.partition", + location=location(locations, f"resources.{name}.partition", source), + ) + } + if "partition" in resource + else {} + ), + } + return Site( + source=source, + locations=locations, + body={ + **site, + "kind": kind, + "environment": environment, + "paths": paths, + "slurm": slurm, + "baremetal": baremetal, + }, + resources=normalized_resources, + ) + + +def recipe_template( + *, + model: str = "qwen3.5-0.8b", + workflow: str = "vlm-pruning", + mode: str = "smoke", + resource_profile: str = "smoke", + run_root: str = "puzzle_runs/my-puzzletron-run", + name: str = "my-puzzletron-run", +) -> dict[str, Any]: + """Return the exact recipe contract used by the resolver.""" + + route = ROUTES_BY_KEY.get((model, workflow, mode)) + if route is None: + choices = ", ".join(sorted(item.route_id for item in ROUTES)) + raise ValueError(f"Unknown route {(model, workflow, mode)!r}; choose one of: {choices}") + recipe = { + "schema_version": RECIPE_SCHEMA_VERSION, + "name": name, + "model": model, + "workflow": workflow, + "mode": mode, + "run_root": run_root, + "resource_profile": resource_profile, + } + if route.requires_data: + recipe["data"] = { + "path": "REPLACE_WITH_PREPARED_DATASET", + "revision": "REPLACE_WITH_IMMUTABLE_DATASET_REVISION", + } + return recipe + + +def site_template() -> dict[str, Any]: + """Return one portable site contract with small and multi-node profiles.""" + + return { + "schema_version": RECIPE_SCHEMA_VERSION, + "site": { + "kind": "slurm", + "environment": { + "repository": "REPLACE_WITH_WORKER_VISIBLE_MODELOPT_CHECKOUT", + "venv": "REPLACE_WITH_WORKER_VISIBLE_MODELOPT_VENV", + "container": None, + "container_mounts": None, + "prerun_commands": [], + "postrun_commands": [], + }, + "paths": {"hf_home": "REPLACE_WITH_SHARED_HF_HOME"}, + "slurm": { + "account": "REPLACE_WITH_SLURM_ACCOUNT", + "partition": "REPLACE_WITH_SLURM_PARTITION", + "partition_cpu": None, + "cpu_cpus_per_task": 4, + "cpu_memory_mb": 32768, + "time_limit": "4:00:00", + }, + }, + "resources": { + "single-gpu": { + "mode": "per_attempt", + "gpus_per_node": 1, + "max_nodes": 1, + }, + "smoke": { + "mode": "reusable_allocation", + "gpus_per_node": 2, + "max_nodes": 1, + }, + "campaign": { + "mode": "reusable_allocation", + "gpus_per_node": 8, + "max_nodes": 1, + }, + "multinode": { + "mode": "per_attempt", + "gpus_per_node": 8, + "max_nodes": 64, + }, + }, + } diff --git a/modelopt/torch/puzzletron/orchestration/_route_catalog.py b/modelopt/torch/puzzletron/orchestration/_route_catalog.py new file mode 100644 index 00000000000..4de6e3dbf3c --- /dev/null +++ b/modelopt/torch/puzzletron/orchestration/_route_catalog.py @@ -0,0 +1,133 @@ +# 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. + +"""Maintained model and workflow routes for concise Puzzletron recipes.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Mapping + +__all__ = ["MODEL_IDS", "ROUTES", "ROUTES_BY_KEY", "RouteProfile"] + + +@dataclass(frozen=True) +class RouteProfile: + """One maintained logical workflow; no site facts live here.""" + + model: str + workflow: str + mode: str + search: str + evaluation: str + distillation: str + experiment_template: str + execution_stages: Mapping[str, Mapping[str, Any]] + requires_data: bool = False + + @property + def route_id(self) -> str: + return f"{self.model}/{self.workflow}/{self.mode}" + + +_TEXT_SMOKE_STAGES: dict[str, dict[str, Any]] = { + "replacement_scoring": {"strategy": "single"}, +} + +_VLM_SMOKE_STAGES: dict[str, dict[str, Any]] = { + "replacement_scoring": {"strategy": "single"}, +} + +_VLM_CAMPAIGN_STAGES: dict[str, dict[str, Any]] = { + "depth_importance": {"strategy": "single"}, + "replacement_scoring": {"instances": 8}, + "post.candidates.image_eval": {"instances": 5}, + "post.candidates.materialized": {"instances": 5}, + "post.candidates.pre_kd_eval": {"instances": 5}, + "post.candidates.kd": {"instances": 5}, + "post.candidates.post_kd_eval": {"instances": 5}, +} + +_FOUR_B_CAMPAIGN_STAGES: dict[str, dict[str, Any]] = { + "replacement_scoring": {"strategy": "single"}, + "post.candidate-evaluation.online_eval": {"instances": 4}, + "post.candidate-evaluation.materialized": {"instances": 4}, + "post.candidate-evaluation.serving": {"instances": 4}, + "post.candidate-evaluation.screening_kd": {"instances": 4}, + "post.candidate-evaluation.screening_eval": {"instances": 4}, + "post.candidate-evaluation.quality_screen": {"instances": 4}, +} + +ROUTES = ( + RouteProfile( + model="qwen3.5-0.8b", + workflow="text-pruning", + mode="smoke", + search="bounded-ffn", + evaluation="smoke", + distillation="smoke", + experiment_template="families/qwen3_5/qwen3p5_0p8b/runs/full_smoke.yaml", + execution_stages=_TEXT_SMOKE_STAGES, + requires_data=True, + ), + RouteProfile( + model="qwen3.5-0.8b", + workflow="vlm-pruning", + mode="smoke", + search="bounded-ffn", + evaluation="smoke", + distillation="smoke", + experiment_template="families/qwen3_5/qwen3p5_0p8b/runs/vlm_smoke.yaml", + execution_stages=_VLM_SMOKE_STAGES, + ), + RouteProfile( + model="qwen3.5-0.8b", + workflow="vlm-pruning", + mode="campaign", + search="multi-axis", + evaluation="quality", + distillation="short", + experiment_template="families/qwen3_5/qwen3p5_0p8b/runs/vlm_campaign.yaml", + execution_stages=_VLM_CAMPAIGN_STAGES, + ), + RouteProfile( + model="qwen3.5-4b", + workflow="vlm-pruning", + mode="smoke", + search="bounded-ffn", + evaluation="smoke", + distillation="smoke", + experiment_template="families/qwen3_5/qwen3p5_4b/runs/vlm_smoke.yaml", + execution_stages={"replacement_scoring": {"strategy": "single"}}, + requires_data=True, + ), + RouteProfile( + model="qwen3.5-4b", + workflow="vlm-pruning", + mode="campaign", + search="ffn-10-to-20pct", + evaluation="quality", + distillation="screening", + experiment_template=("families/qwen3_5/qwen3p5_4b/runs/ffn_width_10to20pct_kd_search.yaml"), + execution_stages=_FOUR_B_CAMPAIGN_STAGES, + requires_data=True, + ), +) + +ROUTES_BY_KEY = {(route.model, route.workflow, route.mode): route for route in ROUTES} +MODEL_IDS = tuple(sorted({route.model for route in ROUTES})) diff --git a/modelopt/torch/puzzletron/orchestration/_source_identity.py b/modelopt/torch/puzzletron/orchestration/_source_identity.py new file mode 100644 index 00000000000..96f2ea5cbe7 --- /dev/null +++ b/modelopt/torch/puzzletron/orchestration/_source_identity.py @@ -0,0 +1,294 @@ +# 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. + +"""Source identity and worker checkout validation for resolved run bundles.""" + +from __future__ import annotations + +import asyncio +import hashlib +import os +import re +import shutil +from functools import lru_cache +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Mapping + + from ._recipe_inputs import Site + +REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +_SOURCE_PATHSPECS = (".", ":(exclude,attr:filter=lfs)") + +_STANDALONE_SOURCE_GUARD = r""" +import hashlib +import json +import os +import re +import shutil +import subprocess +import sys +from pathlib import Path + +repository = Path(sys.argv[1]).expanduser().resolve() +expected = json.loads(sys.argv[2]) +pathspecs = (".", ":(exclude,attr:filter=lfs)") + + +def git_output(*args): + executable = shutil.which("git") + if executable is None: + raise OSError("git executable not found") + result = subprocess.run( + [executable, *args], + cwd=repository, + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + return result.stdout + + +def packaged_revision(): + if repository.parent.name != "src": + return None + marker = repository.parent.parent / "modelopt_revision" + try: + revision = marker.read_text().strip().lower() + except OSError: + return None + if not re.fullmatch(r"(?:[0-9a-f]{40}|[0-9a-f]{64})", revision): + return None + return {"revision": revision, "dirty": False} + + +def working_tree_fingerprint(): + digest = hashlib.sha256() + digest.update(git_output("diff", "--binary", "HEAD", "--", *pathspecs)) + untracked = git_output("ls-files", "--others", "--exclude-standard", "-z").split(b"\0") + for encoded in sorted(item for item in untracked if item): + digest.update(b"\0path\0" + encoded + b"\0") + candidate = repository / os.fsdecode(encoded) + if candidate.is_symlink(): + digest.update(b"symlink\0" + os.fsencode(os.readlink(candidate))) + continue + with candidate.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def repository_revision(): + try: + revision = git_output("rev-parse", "HEAD").decode().strip() + dirty = bool(git_output("status", "--porcelain", "--", *pathspecs).strip()) + except (OSError, subprocess.CalledProcessError): + return packaged_revision() or {"revision": None, "dirty": None} + code = {"revision": revision, "dirty": dirty} + if dirty: + try: + code["working_tree_sha256"] = working_tree_fingerprint() + except (OSError, subprocess.CalledProcessError): + code["working_tree_sha256"] = None + return code + + +def _assert_worker_source(): + actual = repository_revision() + if actual.get("revision") != expected.get("revision"): + raise RuntimeError( + "Worker source revision changed after the run bundle was resolved: " + f"expected {expected.get('revision')}, got {actual.get('revision')}" + ) + expected_dirty = expected.get("dirty") + if expected_dirty is None: + expected_dirty = False + if actual.get("dirty") != expected_dirty: + raise RuntimeError( + "Worker source state changed after the run bundle was resolved: " + f"expected dirty={expected_dirty}, got dirty={actual.get('dirty')}" + ) + if expected_dirty and actual.get("working_tree_sha256") != expected.get( + "working_tree_sha256" + ): + raise RuntimeError("Worker source contents changed after the run bundle was resolved") + + +_assert_worker_source() +""".strip() + + +async def _git_output(repository: Path, *args: str) -> bytes: + """Run one fixed-argument Git query without involving a shell.""" + + executable = shutil.which("git") + if executable is None: + raise OSError("git executable not found") + process = await asyncio.create_subprocess_exec( + executable, + *args, + cwd=repository, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await process.communicate() + if process.returncode: + raise RuntimeError(stderr.decode(errors="replace").strip() or "Git query failed") + return stdout + + +def _run_git(repository: Path, *args: str) -> bytes: + """Run a Git query from the synchronous configuration API.""" + + return asyncio.run(_git_output(repository, *args)) + + +def standalone_source_guard() -> str: + """Return the dependency-free worker identity guard sealed into run bundles.""" + + return _STANDALONE_SOURCE_GUARD + + +def _packaged_revision(repository: Path) -> dict[str, Any] | None: + """Read the revision baked beside the source tree in Puzzletron worker images.""" + + if repository.parent.name != "src": + return None + marker = repository.parent.parent / "modelopt_revision" + try: + revision = marker.read_text().strip().lower() + except OSError: + return None + if not re.fullmatch(r"(?:[0-9a-f]{40}|[0-9a-f]{64})", revision): + return None + return {"revision": revision, "dirty": False} + + +def working_tree_fingerprint(repository: Path = REPOSITORY_ROOT) -> str: + """Hash tracked source changes and untracked files, excluding LFS materialization.""" + + digest = hashlib.sha256() + tracked = _run_git(repository, "diff", "--binary", "HEAD", "--", *_SOURCE_PATHSPECS) + digest.update(tracked) + untracked = _run_git(repository, "ls-files", "--others", "--exclude-standard", "-z").split( + b"\0" + ) + for encoded in sorted(item for item in untracked if item): + digest.update(b"\0path\0" + encoded + b"\0") + candidate = repository / os.fsdecode(encoded) + if candidate.is_symlink(): + digest.update(b"symlink\0" + os.fsencode(os.readlink(candidate))) + continue + with candidate.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def repository_revision(repository: Path) -> dict[str, Any]: + """Return an immutable revision plus a dirty-tree fingerprint when available.""" + + try: + revision = _run_git(repository, "rev-parse", "HEAD").decode().strip() + dirty = bool( + _run_git(repository, "status", "--porcelain", "--", *_SOURCE_PATHSPECS).strip() + ) + except (OSError, RuntimeError): + return _packaged_revision(repository) or {"revision": None, "dirty": None} + code = {"revision": revision, "dirty": dirty} + if dirty: + try: + code["working_tree_sha256"] = working_tree_fingerprint(repository) + except (OSError, RuntimeError): + code["working_tree_sha256"] = None + return code + + +@lru_cache(maxsize=1) +def code_revision() -> dict[str, Any]: + """Return the controller checkout identity.""" + + return repository_revision(REPOSITORY_ROOT) + + +def worker_code( + site: Site, + controller_code: Mapping[str, Any], + *, + detect_revision=None, +) -> dict[str, Any]: + """Resolve and validate the worker checkout identity selected by a site.""" + + environment = dict(site.body["environment"]) + repository = Path(str(environment["repository"])).expanduser().resolve() + explicit = environment.get("source_revision") + if explicit and not re.fullmatch(r"(?:[0-9a-fA-F]{40}|[0-9a-fA-F]{64})", str(explicit)): + raise ValueError("site.environment.source_revision must be a full immutable Git commit") + if explicit: + explicit = str(explicit).lower() + detector = detect_revision or repository_revision + detected = dict(controller_code) if repository == REPOSITORY_ROOT else detector(repository) + detected_revision = detected.get("revision") + if explicit and detected_revision and explicit != detected_revision: + raise ValueError( + "site.environment.source_revision does not match the configured worker " + f"repository HEAD ({detected_revision})" + ) + if detected_revision: + if detected.get("dirty") and not detected.get("working_tree_sha256"): + raise ValueError( + "The configured worker repository is dirty but could not be fingerprinted" + ) + return {**detected, "source": "detected from site.environment.repository"} + if explicit: + return { + "revision": explicit, + "dirty": None, + "source": "site.environment.source_revision", + } + raise ValueError( + "Cannot determine the worker source revision from site.environment.repository; " + "set site.environment.source_revision to its full Git commit" + ) + + +def assert_worker_source( + repository: str, + expected: Mapping[str, Any], + *, + detect_revision=None, +) -> None: + """Reject a worker checkout that no longer matches its sealed identity.""" + + detector = detect_revision or repository_revision + actual = detector(Path(repository).expanduser().resolve()) + if actual.get("revision") != expected.get("revision"): + raise RuntimeError( + "Worker source revision changed after the run bundle was resolved: " + f"expected {expected.get('revision')}, got {actual.get('revision')}" + ) + expected_dirty = expected.get("dirty") + actual_dirty = actual.get("dirty") + if expected_dirty is None: + expected_dirty = False + if actual_dirty != expected_dirty: + raise RuntimeError( + "Worker source state changed after the run bundle was resolved: " + f"expected dirty={expected_dirty}, got dirty={actual_dirty}" + ) + if expected_dirty and actual.get("working_tree_sha256") != expected.get("working_tree_sha256"): + raise RuntimeError("Worker source contents changed after the run bundle was resolved") diff --git a/modelopt/torch/puzzletron/orchestration/recipe_config.py b/modelopt/torch/puzzletron/orchestration/recipe_config.py new file mode 100644 index 00000000000..6004cf318b1 --- /dev/null +++ b/modelopt/torch/puzzletron/orchestration/recipe_config.py @@ -0,0 +1,931 @@ +# 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. + +"""Resolve human-authored recipes and seal immutable run bundles. + +The recipe contract deliberately has no inheritance. It selects one maintained +model/workflow route and one site-owned resource profile. The resolver then +materializes the established experiment, runner, and execution contracts for +the runtime, while retaining their complete provenance as generated evidence. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +import shlex +import shutil +import tempfile +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass, replace +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import yaml + +from . import _source_identity +from ._recipe_inputs import DATA_FIELDS as _DATA_FIELDS +from ._recipe_inputs import Recipe as _Recipe +from ._recipe_inputs import Site as _Site +from ._recipe_inputs import location as _location +from ._recipe_inputs import parse_recipe as _parse_recipe +from ._recipe_inputs import parse_site as _parse_site +from ._recipe_inputs import recipe_template, site_template +from ._recipe_inputs import required_string as _required_string +from ._route_catalog import MODEL_IDS, ROUTES, ROUTES_BY_KEY, RouteProfile +from .compiler import compile_campaign_plan, load_execution_config, load_runner_config, plan_to_dict +from .config import _compose, _config_root +from .identity import stable_hash +from .stages import stage_ids + +if TYPE_CHECKING: + from .schema import CampaignPlan + +__all__ = [ + "MODEL_IDS", + "ROUTES", + "ResolvedRecipeRun", + "RouteProfile", + "bundle_for_run_root", + "explain_resolved_run", + "materialize_resolved_bundle", + "recipe_template", + "resolve_recipe_run", + "site_template", +] + +_REPOSITORY_ROOT = _source_identity.REPOSITORY_ROOT +_CONFIG_ROOT = _REPOSITORY_ROOT / "examples" / "puzzletron" / "configs" +_BUNDLE_SCHEMA_VERSION = 1 + + +@dataclass(frozen=True) +class ResolvedRecipeRun: + """Validated recipe inputs and generated runtime contracts.""" + + recipe: _Recipe + site: _Site + route: RouteProfile + resource: Mapping[str, Any] + experiment: Mapping[str, Any] + resolved_experiment: Mapping[str, Any] + runner: Mapping[str, Any] + execution: Mapping[str, Any] + plan: Mapping[str, Any] + provenance: Mapping[str, Any] + code: Mapping[str, Any] + bundle_id: str + + +def _route_for(recipe: _Recipe) -> RouteProfile: + key = (recipe.model, recipe.workflow, recipe.mode) + route = ROUTES_BY_KEY.get(key) + if route is None: + choices = ", ".join(sorted(item.route_id for item in ROUTES)) + raise ValueError(f"Unknown model/workflow/mode route {key!r}. Maintained routes: {choices}") + return route + + +def _lookup_dotted(payload: Mapping[str, Any], dotted: str) -> Any: + value: Any = payload + for part in dotted.split("."): + if not isinstance(value, Mapping) or part not in value: + raise KeyError(dotted) + value = value[part] + return value + + +def _set_dotted(payload: dict[str, Any], dotted: str, value: Any) -> Any: + parts = dotted.split(".") + target = payload + for part in parts[:-1]: + child = target.get(part) + if not isinstance(child, dict): + raise KeyError(dotted) + target = child + if parts[-1] not in target: + raise KeyError(dotted) + previous = target[parts[-1]] + target[parts[-1]] = deepcopy(value) + return previous + + +def _merge_explicit( + base: dict[str, Any], + update: Mapping[str, Any], + *, + path: str, + source: Path, + locations: Mapping[str, str], + shadowed: list[dict[str, Any]], +) -> None: + for key, value in update.items(): + dotted = f"{path}.{key}" if path else key + if isinstance(value, Mapping) and isinstance(base.get(key), Mapping): + child = dict(base[key]) + _merge_explicit( + child, + value, + path=dotted, + source=source, + locations=locations, + shadowed=shadowed, + ) + base[key] = child + continue + if key in base and base[key] == value: + raise ValueError( + f"Duplicate no-op advanced value {dotted} at " + f"{_location(locations, f'advanced.{dotted}', source)}" + ) + previous = deepcopy(base.get(key)) if key in base else None + base[key] = deepcopy(value) + shadowed.append( + { + "path": dotted, + "previous": previous, + "value": deepcopy(value), + "winning_source": _location(locations, f"advanced.{dotted}", source), + } + ) + + +def _source_guard_command(environment: Mapping[str, Any], worker_code: Mapping[str, Any]) -> str: + expected = {key: worker_code.get(key) for key in ("revision", "dirty", "working_tree_sha256")} + repository = str(environment["repository"]) + python = str(Path(str(environment["venv"])) / "bin" / "python") + argv = ( + python, + "-I", + "-c", + _source_identity.standalone_source_guard(), + repository, + json.dumps(expected, sort_keys=True, separators=(",", ":")), + ) + return " ".join(shlex.quote(part) for part in argv) + + +def _runner_payload( + site: _Site, + resource: Mapping[str, Any], + worker_code: Mapping[str, Any], +) -> dict[str, Any]: + body = site.body + kind = str(body["kind"]) + environment = deepcopy(dict(body["environment"])) + environment.pop("source_revision", None) + hf_home = dict(body.get("paths") or {}).get("hf_home") + if hf_home: + prerun = list(environment.get("prerun_commands") or ()) + prerun.insert(0, f"export HF_HOME={shlex.quote(str(hf_home))}") + else: + prerun = list(environment.get("prerun_commands") or ()) + prerun.append(_source_guard_command(environment, worker_code)) + environment["prerun_commands"] = prerun + runner: dict[str, Any] = { + "kind": kind, + "execution_contract": environment, + } + if kind == "slurm": + slurm = deepcopy(dict(body["slurm"])) + slurm["max_nodes"] = resource["max_nodes"] + if resource.get("partition") is not None: + slurm["partition"] = resource["partition"] + runner["slurm"] = slurm + else: + baremetal = deepcopy(dict(body["baremetal"])) + baremetal["hosts"] = list(baremetal["hosts"])[: int(resource["max_nodes"])] + runner["inventory"] = baremetal + return {"runner": runner} + + +def _execution_payload( + recipe: _Recipe, + route: RouteProfile, + resource: Mapping[str, Any], + *, + shadowed: list[dict[str, Any]], +) -> dict[str, Any]: + execution: dict[str, Any] = { + "schema_version": 1, + "mode": resource["mode"], + "defaults": { + "failure_policy": "strict", + "halt_policy": "drain", + "gpus_per_node": resource["gpus_per_node"], + **({"partition": resource["partition"]} if resource.get("partition") else {}), + }, + "stages": deepcopy(dict(route.execution_stages)), + } + if recipe.advanced_execution: + _merge_explicit( + execution, + recipe.advanced_execution, + path="execution", + source=recipe.source, + locations=recipe.locations, + shadowed=shadowed, + ) + return {"execution": execution} + + +def _write_yaml( + path: Path, + payload: Mapping[str, Any], + *, + header: Sequence[str] = (), +) -> None: + comments = "".join(f"# {line}\n" for line in header) + path.write_text(comments + yaml.safe_dump(dict(payload), sort_keys=False)) + + +def _compile_preview( + experiment: Mapping[str, Any], + runner: Mapping[str, Any], + execution: Mapping[str, Any], +) -> CampaignPlan: + with tempfile.TemporaryDirectory(prefix="puzzletron-recipe-config-") as temporary: + root = Path(temporary) + experiment_path = root / "experiment.yaml" + runner_path = root / "runner.yaml" + execution_path = root / "execution.yaml" + _write_yaml(experiment_path, experiment) + _write_yaml(runner_path, runner) + _write_yaml(execution_path, execution) + return compile_campaign_plan( + experiment_config_path=experiment_path, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + ) + + +def _validate_usage( + recipe: _Recipe, + plan: CampaignPlan, + resource: Mapping[str, Any], +) -> None: + enabled = {node.stage_id for node in plan.stages} + configured_execution = set(recipe.advanced_execution.get("stages", {})) + unused_execution = sorted(configured_execution - enabled - {"final_report"}) + if unused_execution: + raise ValueError( + "Execution values target inactive stages and would be unused: " + + ", ".join(unused_execution) + ) + known_stages = set(stage_ids()) + for dotted in recipe.advanced_experiment: + root = dotted.partition(".")[0] + if root in known_stages and root not in enabled: + if dotted == f"{root}.enabled" and recipe.advanced_experiment[dotted] is False: + continue + raise ValueError( + f"advanced.experiment.{dotted} targets inactive stage {root!r} and is unused" + ) + max_nodes = int(resource["max_nodes"]) + oversized = [ + f"{node.stage_id} ({node.nodes} nodes, {node.total_gpus} GPUs, mesh={dict(node.mesh)})" + for node in plan.stages + if node.nodes > max_nodes + ] + if oversized: + raise ValueError( + f"Resource profile {recipe.resource_profile!r} provides at most {max_nodes} node(s), " + "but the resolved topology requires more: " + "; ".join(oversized) + ) + + +def _plan_without_preview_path(plan: CampaignPlan) -> dict[str, Any]: + payload = plan_to_dict(plan) + payload["experiment_config_path"] = "" + return payload + + +def resolve_recipe_run( + recipe_path: str | Path, + site_path: str | Path, + *, + run_root: str | Path | None = None, +) -> ResolvedRecipeRun: + """Resolve and validate one recipe + site pair without writing its run root.""" + + recipe = _parse_recipe(recipe_path, run_root=run_root) + site = _parse_site(site_path) + route = _route_for(recipe) + if (route.requires_data or recipe.data) and set(recipe.data) != _DATA_FIELDS: + missing = ", ".join(sorted(_DATA_FIELDS - set(recipe.data))) + raise ValueError( + f"Route {route.route_id} requires explicit immutable data values; missing: {missing}" + ) + input_placeholders = _placeholder_paths( + { + "recipe": recipe.as_dict(), + "site": {"site": site.body}, + } + ) + if input_placeholders: + raise ValueError( + "Unresolved example placeholders must be replaced before validation: " + + ", ".join(input_placeholders) + ) + controller_code = _code_revision() + worker_code = _worker_code(site, controller_code) + code = {"controller": controller_code, "worker": worker_code} + try: + resource = deepcopy(dict(site.resources[recipe.resource_profile])) + except KeyError as error: + choices = ", ".join(sorted(site.resources)) + raise ValueError( + f"Unknown resource_profile {recipe.resource_profile!r} at " + f"{_location(recipe.locations, 'resource_profile', recipe.source)}; " + f"site profiles: {choices}" + ) from error + + template_path = _CONFIG_ROOT / route.experiment_template + experiment = _compose( + template_path, + root=_config_root(template_path), + stack=(), + ) + shadowed: list[dict[str, Any]] = [] + for section_name, raw_section in list(experiment.items()): + if not isinstance(raw_section, Mapping) or "evaluator_revision" not in raw_section: + continue + section = dict(raw_section) + previous = section["evaluator_revision"] + if isinstance(previous, str) and "PUZZLETRON_SOURCE_REVISION" in previous: + if worker_code.get("dirty"): + fingerprint = worker_code.get("working_tree_sha256") + revision = f"uncommitted:{str(fingerprint)[:12]}" if fingerprint else "uncommitted" + winning_source = str(worker_code["source"]) + else: + revision = str(worker_code["revision"]) + winning_source = str(worker_code["source"]) + section["evaluator_revision"] = revision + experiment[section_name] = section + shadowed.append( + { + "path": f"experiment.{section_name}.evaluator_revision", + "previous": previous, + "value": revision, + "winning_source": winning_source, + } + ) + experiment = _freeze_environment_defaults(experiment, path="", shadowed=shadowed) + if recipe.data.get("path"): + previous = experiment.get("dataset_path") + experiment["dataset_path"] = recipe.data["path"] + shadowed.append( + { + "path": "experiment.dataset_path", + "previous": previous, + "value": recipe.data["path"], + "winning_source": _location(recipe.locations, "data.path", recipe.source), + } + ) + if recipe.data.get("revision"): + if ( + isinstance(experiment.get("prepare_dataset"), Mapping) + and "revision" in experiment["prepare_dataset"] + ): + revision_section = "prepare_dataset" + elif isinstance(experiment.get("data"), Mapping) and "revision" in experiment["data"]: + revision_section = "data" + else: + raise ValueError( + f"Route {route.route_id} cannot consume data.revision; " + "the internal route template must define a revision field" + ) + section_payload = dict(experiment[revision_section]) + previous = section_payload["revision"] + section_payload["revision"] = recipe.data["revision"] + experiment[revision_section] = section_payload + shadowed.append( + { + "path": f"experiment.{revision_section}.revision", + "previous": previous, + "value": recipe.data["revision"], + "winning_source": _location(recipe.locations, "data.revision", recipe.source), + } + ) + hf_home = dict(site.body.get("paths") or {}).get("hf_home") + if hf_home and isinstance(experiment.get("prepare_dataset"), Mapping): + prepare_dataset = dict(experiment["prepare_dataset"]) + if "evaluation_hf_home" in prepare_dataset: + previous = prepare_dataset["evaluation_hf_home"] + prepare_dataset["evaluation_hf_home"] = str(hf_home) + experiment["prepare_dataset"] = prepare_dataset + shadowed.append( + { + "path": "experiment.prepare_dataset.evaluation_hf_home", + "previous": previous, + "value": str(hf_home), + "winning_source": _location(site.locations, "site.paths.hf_home", site.source), + } + ) + previous_root = experiment.get("puzzle_dir") + experiment["puzzle_dir"] = str(recipe.run_root) + shadowed.append( + { + "path": "experiment.puzzle_dir", + "previous": previous_root, + "value": str(recipe.run_root), + "winning_source": recipe.run_root_source, + } + ) + for dotted, value in recipe.advanced_experiment.items(): + try: + previous = _lookup_dotted(experiment, dotted) + except KeyError as error: + raise ValueError( + f"Unknown advanced experiment path {dotted!r} at " + f"{_location(recipe.locations, f'advanced.experiment.{dotted}', recipe.source)}" + ) from error + if previous == value: + raise ValueError( + f"Duplicate no-op advanced value {dotted!r} at " + f"{_location(recipe.locations, f'advanced.experiment.{dotted}', recipe.source)}" + ) + _set_dotted(experiment, dotted, value) + shadowed.append( + { + "path": f"experiment.{dotted}", + "previous": previous, + "value": deepcopy(value), + "winning_source": _location( + recipe.locations, f"advanced.experiment.{dotted}", recipe.source + ), + } + ) + + runner = _runner_payload(site, resource, worker_code) + site_partition = dict(site.body.get("slurm") or {}).get("partition") + if resource.get("partition") is not None and resource.get("partition") != site_partition: + shadowed.append( + { + "path": "runner.slurm.partition", + "previous": site_partition, + "value": resource["partition"], + "winning_source": _location( + site.locations, + f"resources.{recipe.resource_profile}.partition", + site.source, + ), + } + ) + execution = _execution_payload(recipe, route, resource, shadowed=shadowed) + hidden_environment = _string_paths_containing(experiment, "${oc.env:") + if hidden_environment: + raise ValueError( + "Internal route left environment-dependent values unresolved: " + + ", ".join(hidden_environment) + + ". Move the choice into the recipe or site contract." + ) + placeholders = _placeholder_paths( + { + "experiment": experiment, + "runner": runner, + } + ) + if placeholders: + raise ValueError( + "Unresolved example placeholders must be replaced before validation: " + + ", ".join(placeholders) + ) + compiled_plan = _compile_preview(experiment, runner, execution) + _validate_usage(recipe, compiled_plan, resource) + enabled_stages = {node.stage_id for node in compiled_plan.stages} + execution_body = deepcopy(dict(execution["execution"])) + execution_body["stages"] = { + stage_id: value + for stage_id, value in dict(execution_body.get("stages", {})).items() + if stage_id in enabled_stages or stage_id == "final_report" + } + execution = {"execution": execution_body} + compiled_plan = _compile_preview(experiment, runner, execution) + resolved_experiment = { + key: value for key, value in compiled_plan.experiment_config.items() if key != "_runtime" + } + + provenance = { + "schema_version": 1, + "route": { + "value": route.route_id, + "source": f"internal route catalog ({route.experiment_template})", + }, + "recipe_values": { + key: { + "value": value, + "source": ( + recipe.run_root_source + if key == "run_root" + else _location(recipe.locations, key, recipe.source) + ), + } + for key, value in recipe.as_dict().items() + if key != "advanced" + }, + "site": { + "value": str(site.source), + "resource_profile": recipe.resource_profile, + "source": _location( + site.locations, f"resources.{recipe.resource_profile}", site.source + ), + }, + "generated_values": { + "experiment": {"source": f"internal route catalog ({route.experiment_template})"}, + "runner.execution_contract": { + "source": _location(site.locations, "site.environment", site.source) + }, + ("runner.slurm" if site.body["kind"] == "slurm" else "runner.inventory"): { + "source": _location( + site.locations, + "site.slurm" if site.body["kind"] == "slurm" else "site.baremetal", + site.source, + ) + }, + "execution.mode": { + "source": _location( + site.locations, + f"resources.{recipe.resource_profile}.mode", + site.source, + ) + }, + "execution.defaults.gpus_per_node": { + "source": _location( + site.locations, + f"resources.{recipe.resource_profile}.gpus_per_node", + site.source, + ) + }, + "execution.stages": {"source": f"internal route catalog ({route.route_id})"}, + }, + "shadowed": shadowed, + } + identity_payload = { + "schema_version": _BUNDLE_SCHEMA_VERSION, + "recipe": recipe.as_dict(), + "selected_site": runner, + "route": route.route_id, + "experiment_runtime": experiment, + "experiment_resolved": resolved_experiment, + "execution": execution, + "provenance": provenance, + "code": code, + } + bundle_id = stable_hash(identity_payload, prefix="resolved_bundle") + return ResolvedRecipeRun( + recipe=recipe, + site=site, + route=route, + resource=resource, + experiment=experiment, + resolved_experiment=resolved_experiment, + runner=runner, + execution=execution, + plan=_plan_without_preview_path(compiled_plan), + provenance=provenance, + code=code, + bundle_id=bundle_id, + ) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +_working_tree_fingerprint = _source_identity.working_tree_fingerprint +_repository_revision = _source_identity.repository_revision + + +def _code_revision() -> dict[str, Any]: + return _source_identity.code_revision() + + +def _worker_code(site: _Site, controller_code: Mapping[str, Any]) -> dict[str, Any]: + return _source_identity.worker_code(site, controller_code, detect_revision=_repository_revision) + + +def _assert_worker_source(repository: str, expected: Mapping[str, Any]) -> None: + _source_identity.assert_worker_source( + repository, expected, detect_revision=_repository_revision + ) + + +def _verify_existing_bundle( + path: Path, bundle_id: str, *, manifest_sha256: str | None = None +) -> None: + manifest_path = path / "manifest.json" + if not manifest_path.is_file(): + raise RuntimeError(f"Resolved bundle is incomplete: {path}") + if manifest_sha256 is not None and _sha256(manifest_path) != manifest_sha256: + raise RuntimeError(f"Resolved bundle manifest changed after activation: {manifest_path}") + manifest = json.loads(manifest_path.read_text()) + if manifest.get("bundle_id") != bundle_id: + raise RuntimeError(f"Resolved bundle identity mismatch: {path}") + for name, digest in dict(manifest.get("files") or {}).items(): + candidate = path / name + if not candidate.is_file() or _sha256(candidate) != digest: + raise RuntimeError(f"Resolved bundle file changed after sealing: {candidate}") + recipe = yaml.safe_load((path / "recipe.yaml").read_text()) + runner = yaml.safe_load((path / "runner.yaml").read_text()) + experiment_runtime = yaml.safe_load((path / "experiment.runtime.yaml").read_text()) + experiment_resolved = yaml.safe_load((path / "experiment.resolved.yaml").read_text()) + execution = yaml.safe_load((path / "execution.yaml").read_text()) + provenance = json.loads((path / "provenance.json").read_text()) + recomputed = stable_hash( + { + "schema_version": _BUNDLE_SCHEMA_VERSION, + "recipe": recipe, + "selected_site": runner, + "route": provenance["route"]["value"], + "experiment_runtime": experiment_runtime, + "experiment_resolved": experiment_resolved, + "execution": execution, + "provenance": provenance, + "code": manifest["code"], + }, + prefix="resolved_bundle", + ) + if recomputed != bundle_id: + raise RuntimeError(f"Resolved bundle contents do not match its identity: {path}") + + +def _string_paths_containing(value: Any, needle: str, *, path: str = "") -> list[str]: + if isinstance(value, str): + return [path] if needle in value else [] + if isinstance(value, Mapping): + return [ + candidate + for key, item in value.items() + for candidate in _string_paths_containing( + item, needle, path=f"{path}.{key}" if path else str(key) + ) + ] + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return [ + candidate + for index, item in enumerate(value) + for candidate in _string_paths_containing(item, needle, path=f"{path}[{index}]") + ] + return [] + + +_ENVIRONMENT_DEFAULT = re.compile(r"^\$\{oc\.env:([^,}]+),(.+)\}$") + + +def _freeze_environment_defaults( + value: Any, + *, + path: str, + shadowed: list[dict[str, Any]], +) -> Any: + """Use a template's declared fallback without consulting process state.""" + + if isinstance(value, Mapping): + return { + key: _freeze_environment_defaults( + item, + path=f"{path}.{key}" if path else str(key), + shadowed=shadowed, + ) + for key, item in value.items() + } + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return [ + _freeze_environment_defaults( + item, + path=f"{path}[{index}]", + shadowed=shadowed, + ) + for index, item in enumerate(value) + ] + if not isinstance(value, str) or not (match := _ENVIRONMENT_DEFAULT.fullmatch(value)): + return value + default = yaml.safe_load(match.group(2)) + shadowed.append( + { + "path": f"experiment.{path}", + "previous": value, + "value": deepcopy(default), + "winning_source": f"declared fallback for {match.group(1)}", + } + ) + return default + + +def _placeholder_paths(value: Any, *, path: str = "") -> list[str]: + return _string_paths_containing(value, "REPLACE_WITH_", path=path) + + +def _replace_plan_path(plan: CampaignPlan, experiment_path: Path) -> CampaignPlan: + return replace(plan, experiment_config_path=str(experiment_path)) + + +def materialize_resolved_bundle( + resolved: ResolvedRecipeRun, + *, + activate: bool, +) -> Path: + """Write one content-addressed bundle and optionally bind the run root to it.""" + + orchestration_root = resolved.recipe.run_root / "orchestration" + bundles_root = orchestration_root / "resolved_bundles" + final = bundles_root / resolved.bundle_id + current_path = orchestration_root / "current_bundle.json" + if final.exists(): + _verify_existing_bundle(final, resolved.bundle_id) + else: + bundles_root.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=f".{resolved.bundle_id}.", dir=str(bundles_root))) + try: + generated_header = ( + "GENERATED PUZZLETRON RUN ARTIFACT: DO NOT EDIT.", + f"Resolved bundle: {resolved.bundle_id}", + "Change the user recipe or site config, then launch under a new run root.", + ) + _write_yaml( + staging / "recipe.yaml", + resolved.recipe.as_dict(), + header=generated_header + ("Normalized snapshot of the human-authored recipe.",), + ) + _write_yaml( + staging / "experiment.runtime.yaml", + resolved.experiment, + header=generated_header + + ( + "Executable worker input; Hydra expressions are intentionally preserved.", + "Resume uses this sealed file. Inspect it, but do not use it as authoring input.", + ), + ) + _write_yaml( + staging / "experiment.resolved.yaml", + resolved.resolved_experiment, + header=generated_header + + ( + "Fully resolved audit view for understanding the winning values.", + "This evidence-only view is not executable worker input.", + ), + ) + _write_yaml( + staging / "runner.yaml", + resolved.runner, + header=generated_header + ("Resolved site and runner snapshot.",), + ) + _write_yaml( + staging / "execution.yaml", + resolved.execution, + header=generated_header + ("Resolved execution-policy snapshot.",), + ) + plan = compile_campaign_plan( + experiment_config_path=staging / "experiment.runtime.yaml", + runner=load_runner_config(staging / "runner.yaml"), + execution=load_execution_config(staging / "execution.yaml"), + ) + plan = _replace_plan_path(plan, final / "experiment.runtime.yaml") + (staging / "plan.json").write_text(json.dumps(plan_to_dict(plan), indent=2) + "\n") + (staging / "provenance.json").write_text( + json.dumps(resolved.provenance, indent=2, default=str) + "\n" + ) + files = { + path.name: _sha256(path) for path in sorted(staging.iterdir()) if path.is_file() + } + (staging / "manifest.json").write_text( + json.dumps( + { + "schema_version": _BUNDLE_SCHEMA_VERSION, + "bundle_id": resolved.bundle_id, + "code": dict(resolved.code), + "files": files, + }, + indent=2, + ) + + "\n" + ) + try: + os.rename(staging, final) + except OSError: + if not final.is_dir(): + raise + _verify_existing_bundle(final, resolved.bundle_id) + shutil.rmtree(staging) + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + if activate: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + prefix=f".{current_path.name}.", + suffix=".tmp", + dir=orchestration_root, + delete=False, + ) as stream: + json.dump( + { + "schema_version": _BUNDLE_SCHEMA_VERSION, + "bundle_id": resolved.bundle_id, + "manifest_sha256": _sha256(final / "manifest.json"), + }, + stream, + indent=2, + ) + stream.write("\n") + temporary = Path(stream.name) + try: + os.link(temporary, current_path) + except FileExistsError: + current = json.loads(current_path.read_text()) + if current.get("bundle_id") != resolved.bundle_id: + raise ValueError( + f"Run root {resolved.recipe.run_root} is already bound to bundle " + f"{current.get('bundle_id')}; choose a new run_root or resume it" + ) + _verify_existing_bundle( + final, + resolved.bundle_id, + manifest_sha256=current.get("manifest_sha256"), + ) + finally: + temporary.unlink(missing_ok=True) + return final + + +def bundle_for_run_root(run_root: str | Path) -> Path: + """Return and integrity-check the active immutable bundle for one run.""" + + root = Path(run_root).expanduser().resolve() + current_path = root / "orchestration" / "current_bundle.json" + if not current_path.is_file(): + raise FileNotFoundError(f"No active resolved bundle exists under {root}") + current = json.loads(current_path.read_text()) + bundle_id = _required_string( + current.get("bundle_id"), path="bundle_id", location=str(current_path) + ) + bundle = root / "orchestration" / "resolved_bundles" / bundle_id + _verify_existing_bundle( + bundle, + bundle_id, + manifest_sha256=current.get("manifest_sha256"), + ) + return bundle + + +def explain_resolved_run(resolved: ResolvedRecipeRun) -> str: + """Render a compact, human-readable source and allocation explanation.""" + + lines = [ + (f"route: {resolved.route.route_id} <- {resolved.provenance['route']['source']}"), + ( + "intent: " + f"search={resolved.route.search}, evaluation={resolved.route.evaluation}, " + f"distillation={resolved.route.distillation}" + ), + f"bundle: {resolved.bundle_id}", + "recipe values:", + ] + for key, item in resolved.provenance["recipe_values"].items(): + lines.append(f" {key}: {item['value']!r} <- {item['source']}") + lines.extend( + [ + ( + f"site resource: {resolved.recipe.resource_profile} " + f"({resolved.resource['mode']}, {resolved.resource['gpus_per_node']} GPUs/node, " + f"up to {resolved.resource['max_nodes']} node(s)) <- " + f"{resolved.provenance['site']['source']}" + ), + "stages:", + ] + ) + for node in resolved.plan["stages"]: + mesh = node["mesh"] + mesh_text = " ".join( + f"{key}={mesh[key]}" for key in ("tp", "pp", "cp", "dp_shard", "dp_replicate", "ep") + ) + lines.append( + f" {node['stage_id']}: {node['nodes']} node(s), {node['total_gpus']} GPU(s), " + f"{node['instances']} instance(s), {mesh_text}" + ) + shadows = list(resolved.provenance.get("shadowed") or ()) + lines.append(f"explicit resolved overrides: {len(shadows)}") + for item in shadows: + lines.append(f" {item['path']} <- {item['winning_source']}") + return "\n".join(lines) diff --git a/modelopt/torch/puzzletron/plugins/automodel/config.py b/modelopt/torch/puzzletron/plugins/automodel/config.py index 3704151a2ee..a8fc021a6d9 100644 --- a/modelopt/torch/puzzletron/plugins/automodel/config.py +++ b/modelopt/torch/puzzletron/plugins/automodel/config.py @@ -64,6 +64,7 @@ _FROM_PRETRAINED_TARGET = "nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained" _FROM_PRETRAINED_VLM_TARGET = "nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained" _DEFAULT_TEACHER_SUBDIR = "ckpts/teacher" +_RUNTIME_CHECKPOINT_SUBDIR = ".runtime/nemo_automodel/checkpoints" def _int_or_default(value, default: int = 1) -> int: @@ -88,6 +89,16 @@ def _teacher_path(hydra_cfg) -> str: return f"{hydra_cfg.puzzle_dir}/{_DEFAULT_TEACHER_SUBDIR}" +def _route_runtime_files_to_run_root(recipe: dict, hydra_cfg) -> dict: + """Keep NeMo's disabled-checkpoint bookkeeping out of the source checkout.""" + checkpoint = dict(recipe.get("checkpoint") or {}) + checkpoint.setdefault( + "checkpoint_dir", str(Path(str(hydra_cfg.puzzle_dir)) / _RUNTIME_CHECKPOINT_SUBDIR) + ) + recipe["checkpoint"] = checkpoint + return recipe + + def _trust_remote_code(hydra_cfg) -> bool: """Resolve the shared model and processor remote-code policy.""" model = _as_dict(hydra_cfg.get("model", None)) @@ -547,6 +558,7 @@ def build_recipe_config(hydra_cfg) -> dict: """ automodel_cfg = hydra_cfg.pruning.get("automodel", None) recipe = _inject_authored_model_backend(build_stage_recipe_config(automodel_cfg), hydra_cfg) + recipe = _route_runtime_files_to_run_root(recipe, hydra_cfg) model = dict(recipe.get("model", {})) model.setdefault("_target_", _FROM_PRETRAINED_TARGET) @@ -645,6 +657,7 @@ def build_solution_recipe_config(hydra_cfg, model_path) -> dict: """ automodel_cfg = hydra_cfg.scoring.get("automodel", None) recipe = _inject_authored_model_backend(build_stage_recipe_config(automodel_cfg), hydra_cfg) + recipe = _route_runtime_files_to_run_root(recipe, hydra_cfg) force_hf = bool(_as_dict(automodel_cfg).get("force_hf", False)) runtime_cfg = hydra_cfg.get("_runtime", {}) or {} descriptor = hydra_cfg.get("descriptor", None) or runtime_cfg.get("descriptor", None) diff --git a/modelopt/torch/puzzletron/stages/pipeline.py b/modelopt/torch/puzzletron/stages/pipeline.py index c37841ba024..73b367ef464 100644 --- a/modelopt/torch/puzzletron/stages/pipeline.py +++ b/modelopt/torch/puzzletron/stages/pipeline.py @@ -220,6 +220,10 @@ def _calculate_static_workload_stats(config: dict[str, Any], hydra_cfg: Any) -> """Append one analytical memory profile for every configured MIP workload.""" from ..subblock_stats.calc_subblock_stats import launch_calc_subblock_stats + # Hidden-width scenarios own the width list. Keep the static cost + # inventory aligned without requiring users to repeat it under + # ``vllm_stats`` when runtime measurements are disabled. + configure_vllm_stats_widths(config, hydra_cfg) workloads = dict((config.get("mip") or {}).get("workloads") or {}) if not workloads: workloads = { diff --git a/noxfile.py b/noxfile.py index fe8f982818a..6ed49e4936c 100644 --- a/noxfile.py +++ b/noxfile.py @@ -316,7 +316,7 @@ def gpu_puzzletron(session): "addopts=", ( "tests/gpu/torch/puzzletron/test_puzzletron.py::" - "test_tiny_qwen_campaign_uses_current_public_route" + "test_tiny_qwen_full_lifecycle_uses_public_bundle_contract" ), *_cov_args(), ) diff --git a/puzzletron_setup/profiles.py b/puzzletron_setup/profiles.py index 6876be4dfe4..f79e4ee7526 100644 --- a/puzzletron_setup/profiles.py +++ b/puzzletron_setup/profiles.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. """Declarative, dependency-free model capabilities for Puzzletron setup.""" @@ -204,8 +216,7 @@ def __init__(self, config: Mapping[str, Any]): message = ( "Unsupported model family. Detected model types " f"{list(model_types)} and architectures {list(architectures)}. " - "Use .agents/skills/running-puzzletron/SKILL.md to add the runtime " - "descriptor and a lightweight puzzletron_setup profile for this family." + "Add a runtime descriptor and setup profile for this family before rerunning setup." ) super().__init__(message) self.model_types = model_types diff --git a/puzzletron_setup/v2/bundle.py b/puzzletron_setup/v2/bundle.py index ef93b1e1760..a848fe7ae45 100644 --- a/puzzletron_setup/v2/bundle.py +++ b/puzzletron_setup/v2/bundle.py @@ -406,7 +406,7 @@ def _bundle_readme( lines = [ "# Puzzletron campaign", "", - "Generated by `puzzletron_setup_v2.py`. The wizard did not launch any jobs.", + "Generated by `puzzletron.py setup`. 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'}`." @@ -573,7 +573,8 @@ def _bundle_readme( shlex.join( [ "python", - str(Path(repository) / "examples/puzzletron/puzzletron_setup_v2.py"), + str(Path(repository) / "examples/puzzletron/puzzletron.py"), + "setup", "--resume", str(campaign_dir), ] diff --git a/puzzletron_setup/v2/cli.py b/puzzletron_setup/v2/cli.py index b3d7848739b..38659ac1bd5 100644 --- a/puzzletron_setup/v2/cli.py +++ b/puzzletron_setup/v2/cli.py @@ -100,7 +100,7 @@ def main(argv: Sequence[str] | None = None) -> int: target = args.resume or "" print( "\nSetup interrupted. Resume with: " - f"python examples/puzzletron/puzzletron_setup_v2.py --resume {target}" + f"python examples/puzzletron/puzzletron.py setup --resume {target}" ) return 130 except SetupError as error: diff --git a/tests/_test_utils/torch/puzzletron/configs/tiny_qwen.recipe.yaml b/tests/_test_utils/torch/puzzletron/configs/tiny_qwen.recipe.yaml new file mode 100644 index 00000000000..7b3c270f5d8 --- /dev/null +++ b/tests/_test_utils/torch/puzzletron/configs/tiny_qwen.recipe.yaml @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# The fixture replaces only the temporary run and data paths. +schema_version: 1 +name: tiny-qwen-fast-e2e +model: qwen3.5-0.8b +workflow: text-pruning +mode: smoke +run_root: REPLACE_WITH_TEST_RUN_ROOT +resource_profile: fast-e2e +data: + path: REPLACE_WITH_TEST_DATASET + revision: hermetic-tiny-qwen-v1 diff --git a/tests/_test_utils/torch/puzzletron/configs/tiny_qwen.site.yaml b/tests/_test_utils/torch/puzzletron/configs/tiny_qwen.site.yaml new file mode 100644 index 00000000000..743bceff8cd --- /dev/null +++ b/tests/_test_utils/torch/puzzletron/configs/tiny_qwen.site.yaml @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# The fixture replaces only paths created for the current test process. +schema_version: 1 +site: + kind: slurm + environment: + repository: REPLACE_WITH_TEST_CHECKOUT + venv: REPLACE_WITH_TEST_VENV + container: + container_mounts: + prerun_commands: [] + postrun_commands: [] + paths: + hf_home: REPLACE_WITH_TEST_HF_HOME + slurm: + account: test + job_name_prefix: test-puzzletron +resources: + fast-e2e: + mode: per_attempt + gpus_per_node: 1 + max_nodes: 1 diff --git a/tests/_test_utils/torch/puzzletron/configs/tiny_qwen_lifecycle.overlay.yaml b/tests/_test_utils/torch/puzzletron/configs/tiny_qwen_lifecycle.overlay.yaml new file mode 100644 index 00000000000..46b579d5995 --- /dev/null +++ b/tests/_test_utils/torch/puzzletron/configs/tiny_qwen_lifecycle.overlay.yaml @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Test-only size and runtime overrides for the maintained 0.8B text lifecycle. +# The fixture applies this overlay to the real full_smoke.yaml route. +model: + descriptor_override: qwen3_5_text + revision: +model_info: + model_type: qwen3_5_text + architectures: [Qwen3_5ForCausalLM] + num_hidden_layers: 2 + layer_counts: {full_attention: 2} + hidden_size: 512 + intermediate_size: 768 + num_attention_heads: 4 + num_key_value_heads: 2 + head_dim: 16 + vocab_size: 128 + tie_word_embeddings: false + max_position_embeddings: 128 + +data: + max_sample_length: 32 +tokenize_data: + workers: 1 + caches: + - output: ${train_token_cache_path} + split: train + num_samples: 2 + seq_length: 32 + shuffle_seed: 444 + - output: ${validation_token_cache_path} + split: validation + num_samples: 1 + seq_length: 32 + shuffle_seed: 445 + +embedding_pruning: + widths: [512, 256] +pruning: + eval_samples: 2 + block_size: 32 + intermediate_size_list: [768, 512, 256] + # The maintained route targets the VLM checkpoint layout. This hermetic + # fixture substitutes a text-only Qwen model, so its scoring descriptor must + # follow the text model's `model.layers` module paths. + activation_passes: + - name: ffn_iterative + pruning_mixin: + _target_: modelopt.torch.puzzletron.pruning.ffn_intermediate_pruning_mixin.FFNIntermediatePruningMixIn + layer_descriptor: + _target_: modelopt.torch.puzzletron.anymodel.models.qwen3_5.qwen3_5_model_descriptor.Qwen3P5TextFFNIntermediateLayerDescriptor + hook_class: ${get_object:modelopt.torch.prune.importance_hooks.base_hooks.IterativeChannelContributionHook} + activation_hooks_kwargs: + method: iterative + target_layer: mlp.down_proj + layer_input_descriptors_path: + +search_space: + axes: + ffn_intermediate: {enabled: true, teacher_value: 768, values: [768, 512, 256]} + +sort_sanity: {enabled: false} +width_sanity: {enabled: false} +slicing_sanity: {enabled: false} + +replacement_scoring: + granularity: block + eval_samples: 1 + block_size: 32 + automodel: + lm_head_backend: streaming + +mip: + single_block_replacement_validation_dir: ${to_path:${replacement_scoring.output_dir}} + score_granularity: block + canonical_solutions_path: ${to_path:${replacement_scoring.block_solutions_path}} + runs: + params-90: + search_space: {depth: [0], embedding: [256]} + solver: + num_solutions: 1 + min_hamming_distance: 1 + homogeneous: {keep: 1} + +# External benchmark datasets are deliberately excluded from this hermetic GPU test. +post_mip: + flows: + params-90: + source: {run: params-90, variants: all, objectives: all} + nodes: + online_eval: + type: evaluation + config: {eval_samples: 1, block_size: 32, micro_batch_size: 1} + best_lm: + type: filter + input: online_eval + mode: top_k + metric: online_eval.lm_loss + direction: minimize + top_k: 1 + materialized: {type: materialize, input: best_lm} + serving: + type: aiperf + input: materialized + config: + input_tokens: 16 + output_tokens: 4 + concurrency: [1] + request_count: 1 + use_server_token_count: true + readiness_timeout: 120 + benchmark_timeout: 120 + allow_aiperf_v011_online_tokenizer_resolution: true + topology: + extra_vllm_args: + - -cc.cudagraph_mode=NONE + - --no-enable-flashinfer-autotune + - --max-num-batched-tokens + - "64" + - --max-num-seqs + - "2" + - --gpu-memory-utilization + - "0.5" + fastest: + type: filter + input: serving + mode: top_k + metric: serving.output_token_throughput + direction: maximize + top_k: 1 + short_kd: + type: global_kd + input: fastest + config: + max_steps: 1 + global_batch_size: 1 + local_batch_size: 1 + checkpoint_every_steps: 1 + final_eval: + type: evaluation + input: short_kd + config: {eval_samples: 1, block_size: 32, micro_batch_size: 1} + best: + type: filter + input: final_eval + mode: top_k + metric: final_eval.lm_loss + direction: minimize + top_k: 1 diff --git a/tests/_test_utils/torch/puzzletron/tiny_qwen_experiment_overlay.yaml b/tests/_test_utils/torch/puzzletron/tiny_qwen_experiment_overlay.yaml deleted file mode 100644 index 5c12eb9307c..00000000000 --- a/tests/_test_utils/torch/puzzletron/tiny_qwen_experiment_overlay.yaml +++ /dev/null @@ -1,45 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Keep the generated smoke campaign small while exercising its complete DAG. -tokenize_data: - workers: 1 - -replacement_scoring: - automodel: - lm_head_backend: streaming - -post_mip: - flows: - params-90: - nodes: - online_eval: - config: - eval_samples: 2 - best_lm: - top_k: 3 - serving: - config: - request_count: 4 - allow_aiperf_v011_online_tokenizer_resolution: true - topology: - extra_vllm_args: - - -cc.cudagraph_mode=NONE - - --no-enable-flashinfer-autotune - - --max-num-batched-tokens - - "128" - - --max-num-seqs - - "4" - - --gpu-memory-utilization - - "0.5" - fastest: - top_k: 2 - short_kd: - config: - max_steps: 2 - global_batch_size: 1 - local_batch_size: 1 - checkpoint_every_steps: 2 - final_eval: - config: - eval_samples: 2 diff --git a/tests/_test_utils/torch/puzzletron/tiny_qwen_fixture.py b/tests/_test_utils/torch/puzzletron/tiny_qwen_fixture.py index 1c4feb4131a..fbcf0ec9dbe 100644 --- a/tests/_test_utils/torch/puzzletron/tiny_qwen_fixture.py +++ b/tests/_test_utils/torch/puzzletron/tiny_qwen_fixture.py @@ -21,7 +21,7 @@ import os import subprocess import sys -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from typing import TYPE_CHECKING, Any @@ -29,60 +29,24 @@ from _test_utils.torch.transformers_models import create_tiny_qwen3_5_dir from datasets import Dataset, DatasetDict +import puzzletron_orchestrator.recipe_config as recipe_config from modelopt.torch.puzzletron.pipeline_config import pipeline_config_from_path from puzzletron_orchestrator.compiler import ( compile_campaign_plan, load_execution_config, load_runner_config, ) -from puzzletron_setup.v2.wizard import _DEFAULT_DATA_SOURCE, _DEFAULT_MODEL_SOURCE, run_wizard_v2 +from puzzletron_orchestrator.config import _compose, _config_root, _merge __all__ = ["TinyQwenCampaign", "build_tiny_qwen_campaign"] if TYPE_CHECKING: - from collections.abc import Sequence - from puzzletron_orchestrator.schema import CampaignPlan - from puzzletron_setup.v2.prompts import PromptChoice - - -_SETUP_DEFAULTS = Path(__file__).with_name("tiny_qwen_setup_defaults.yaml") -_EXPERIMENT_OVERLAY = Path(__file__).with_name("tiny_qwen_experiment_overlay.yaml") - - -class _DefaultsBackend: - """Select resolved guided defaults while supplying the campaign directory.""" - - def __init__(self, campaign_dir: Path) -> None: - self.campaign_dir = campaign_dir - - def text(self, message: str, default: str) -> Any: - if message == "Campaign directory:": - return str(self.campaign_dir) - return default - - def select( - self, - message: str, - choices: Sequence[PromptChoice], - default: Any, - ) -> Any: - if message == "Model:": - return _DEFAULT_MODEL_SOURCE - if message == "Dataset:": - return _DEFAULT_DATA_SOURCE - if default is not None: - return default - return next(choice.value for choice in choices if choice.disabled is None) - def checkbox( - self, - message: str, - choices: Sequence[PromptChoice], - defaults: Sequence[Any], - ) -> Any: - del message, choices - return list(defaults) +_CONFIG_DIR = Path(__file__).with_name("configs") +_RECIPE = _CONFIG_DIR / "tiny_qwen.recipe.yaml" +_SITE = _CONFIG_DIR / "tiny_qwen.site.yaml" +_EXPERIMENT_OVERLAY = _CONFIG_DIR / "tiny_qwen_lifecycle.overlay.yaml" @dataclass(frozen=True) @@ -98,14 +62,14 @@ class TinyQwenCampaign: config: dict[str, Any] compiled_plan: CampaignPlan - def run(self, *, timeout: int = 2100) -> subprocess.CompletedProcess[str]: + def run(self, *, timeout: int = 720) -> subprocess.CompletedProcess[str]: """Run or resume the full campaign through the public local orchestrator.""" command = [ sys.executable, str(self.project_root / "examples/puzzletron/orchestrate.py"), "--experiment", - str(self.smoke_bundle / "experiment.yaml"), + str(self.smoke_bundle / "experiment.runtime.yaml"), "--runner", str(self.smoke_bundle / "runner.yaml"), "--execution", @@ -120,15 +84,32 @@ def run(self, *, timeout: int = 2100) -> subprocess.CompletedProcess[str]: ] for override in self.overrides: command.extend(("--override", override)) - return subprocess.run( - command, - cwd=self.project_root, - env=self.environment, - capture_output=True, - text=True, - timeout=timeout, - check=False, - ) + try: + return subprocess.run( + command, + cwd=self.project_root, + env=self.environment, + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + except subprocess.TimeoutExpired as error: + stderr = error.stderr or "" + if isinstance(stderr, bytes): + stderr = stderr.decode(errors="replace") + logs = sorted( + self.smoke_root.glob("logs/**/*.log"), + key=lambda path: path.stat().st_mtime_ns, + ) + log_tail = ( + logs[-1].read_text(errors="replace")[-12000:] if logs else "no task log found" + ) + raise AssertionError( + f"Tiny Qwen Puzzletron campaign timed out after {timeout}s.\n" + f"stderr tail:\n{stderr[-12000:]}\n" + f"latest task-log tail:\n{log_tail}" + ) from error def require_success(self, completed: subprocess.CompletedProcess[str]) -> dict[str, Any]: """Return the controller result or raise with the most useful task log.""" @@ -168,7 +149,7 @@ def _save_messages_dataset(path: Path) -> None: {"role": "user", "content": "What is model compression?"}, {"role": "assistant", "content": response}, ] - rows = [{"messages": messages}] * 8 + rows = [{"messages": messages}] * 2 DatasetDict( { "train": Dataset.from_list(rows), @@ -177,42 +158,31 @@ def _save_messages_dataset(path: Path) -> None: ).save_to_disk(str(path)) -def _merge_config(base: dict[str, Any], overlay: dict[str, Any]) -> dict[str, Any]: - """Recursively apply a declarative test overlay to a generated config.""" - - for key, value in overlay.items(): - if isinstance(value, dict) and isinstance(base.get(key), dict): - _merge_config(base[key], value) - else: - base[key] = value - return base - - -def _write_setup_defaults( - path: Path, - *, - model_dir: Path, - dataset_dir: Path, - result_root: Path, - project_root: Path, -) -> None: - """Materialize the checked-in setup config with per-test runtime paths.""" +def _write_tiny_route_template(project_root: Path, tmp_path: Path, model_dir: Path) -> Path: + """Compose the checked-in tiny overlay onto the maintained lifecycle.""" - defaults = yaml.safe_load(_SETUP_DEFAULTS.read_text()) - defaults["model"]["source"] = str(model_dir) - defaults["data"]["source"] = str(dataset_dir) - defaults["output"]["result_root"] = str(result_root) - contract = defaults["infrastructure"]["execution_contract"] - contract["repository"] = str(project_root) - contract["venv"] = sys.prefix - path.write_text(yaml.safe_dump(defaults, sort_keys=False)) - - -def build_tiny_qwen_campaign( - project_root: Path, - tmp_path: Path, -) -> TinyQwenCampaign: - """Generate the sole tiny-Qwen setup-to-resume Puzzletron E2E fixture.""" + source = ( + project_root + / "examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/full_smoke.yaml" + ) + experiment = _compose(source, root=_config_root(source), stack=()) + overlay = yaml.safe_load(_EXPERIMENT_OVERLAY.read_text()) + experiment = _merge(experiment, overlay) + # The hermetic test replaces external downstream benchmarks with the + # checked-in local-only flow rather than merging the two node mappings. + experiment["post_mip"]["flows"]["params-90"]["nodes"] = overlay["post_mip"]["flows"][ + "params-90" + ]["nodes"] + experiment["input_hf_model_path"] = str(model_dir) + experiment["model_info"]["hf_repo"] = str(model_dir) + catalog = tmp_path / "test-route-catalog" + catalog.mkdir() + (catalog / "tiny_qwen.yaml").write_text(yaml.safe_dump(experiment, sort_keys=False)) + return catalog + + +def build_tiny_qwen_campaign(project_root: Path, tmp_path: Path) -> TinyQwenCampaign: + """Generate the sole tiny-Qwen public-route E2E fixture.""" model_dir = create_tiny_qwen3_5_dir( tmp_path / "model", @@ -225,39 +195,40 @@ def build_tiny_qwen_campaign( layer_types=["full_attention"] * 2, ) dataset_dir = tmp_path / "dataset" - campaign_dir = tmp_path / "campaign" - result_root = tmp_path / "results" + result_root = tmp_path / "fast-e2e" cache_dir = tmp_path / "cache" - defaults_path = tmp_path / "defaults.yaml" + recipe_path = tmp_path / "puzzletron.recipe.yaml" + site_path = tmp_path / "puzzletron.site.yaml" _save_messages_dataset(dataset_dir) - _write_setup_defaults( - defaults_path, - model_dir=model_dir, - dataset_dir=dataset_dir, - result_root=result_root, - project_root=project_root, - ) - generated = run_wizard_v2( - resume=None, - defaults_path=defaults_path, - backend=_DefaultsBackend(campaign_dir), + recipe = yaml.safe_load(_RECIPE.read_text()) + recipe["run_root"] = str(result_root) + recipe["data"]["path"] = str(dataset_dir) + site = yaml.safe_load(_SITE.read_text()) + site["site"]["environment"].update({"repository": str(project_root), "venv": sys.prefix}) + site["site"]["paths"]["hf_home"] = str(cache_dir / "huggingface") + recipe_path.write_text(yaml.safe_dump(recipe, sort_keys=False)) + site_path.write_text(yaml.safe_dump(site, sort_keys=False)) + + catalog_root = _write_tiny_route_template(project_root, tmp_path, model_dir) + route_key = (recipe["model"], recipe["workflow"], recipe["mode"]) + original_root = recipe_config._CONFIG_ROOT + original_routes = recipe_config.ROUTES_BY_KEY + test_routes = dict(original_routes) + test_routes[route_key] = replace( + original_routes[route_key], + experiment_template="tiny_qwen.yaml", ) - smoke_bundle = generated / "smoke" - experiment_path = smoke_bundle / "experiment.yaml" - experiment = yaml.safe_load(experiment_path.read_text()) - flows = dict((experiment.get("post_mip") or {}).get("flows") or {}) - if len(flows) != 1: - raise AssertionError(f"expected one recommended post-MIP flow, found {sorted(flows)}") - flow_id = next(iter(flows)) - overlay = yaml.safe_load(_EXPERIMENT_OVERLAY.read_text()) - overlay_flows = dict(overlay["post_mip"]["flows"]) - if tuple(overlay_flows) != (flow_id,): - raise AssertionError( - f"tiny-Qwen overlay targets {sorted(overlay_flows)}, generated flow is {flow_id!r}" - ) - _merge_config(experiment, overlay) - experiment_path.write_text(yaml.safe_dump(experiment, sort_keys=False)) + try: + recipe_config._CONFIG_ROOT = catalog_root + recipe_config.ROUTES_BY_KEY = test_routes + resolved = recipe_config.resolve_recipe_run(recipe_path, site_path) + finally: + recipe_config._CONFIG_ROOT = original_root + recipe_config.ROUTES_BY_KEY = original_routes + + smoke_bundle = recipe_config.materialize_resolved_bundle(resolved, activate=True) + experiment_path = smoke_bundle / "experiment.runtime.yaml" overrides: tuple[str, ...] = () config = pipeline_config_from_path(experiment_path, overrides=overrides) compiled_plan = compile_campaign_plan( @@ -287,8 +258,8 @@ def build_tiny_qwen_campaign( return TinyQwenCampaign( project_root=project_root, smoke_bundle=smoke_bundle, - smoke_root=result_root / "smoke", - flow_id=flow_id, + smoke_root=result_root, + flow_id="params-90", overrides=overrides, environment=environment, config=config, diff --git a/tests/_test_utils/torch/puzzletron/tiny_qwen_setup_defaults.yaml b/tests/_test_utils/torch/puzzletron/tiny_qwen_setup_defaults.yaml deleted file mode 100644 index 2c0e67c8f6a..00000000000 --- a/tests/_test_utils/torch/puzzletron/tiny_qwen_setup_defaults.yaml +++ /dev/null @@ -1,78 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# The tiny model, dataset, result, checkout, and virtual-environment paths are -# created per test and filled in by tiny_qwen_fixture.py before setup runs. -schema_version: 1 - -model: - source: - trust_remote_code: false - force_hf: false - -data: - source: - modality: text - layout: fixed - sequence_length: 32 - -pruning: - depth_remove: 0 - depth_importance_samples: 2 - width_importance_samples: 2 - replacement_samples: 2 - sort_sanity: false - width_sanity: false - slicing_sanity: false - replacement_granularity: block - axes: - hidden_width: - values: [256] - kv_groups: - values: [2] - q_heads_per_group: - values: [2] - ffn_intermediate: - values: [768, 512, 256] - gdn_key_groups: - values: [2] - gdn_value_heads_per_group: - values: [2] - gdn_key_head_dim: - values: [8] - gdn_value_head_dim: - values: [8] - bypass: - enabled: false - -vllm: - enabled: false - prefill_seq_len: 32 - generation_seq_len: 8 - batch_size: 1 - max_num_seqs: 1 - -mip: - goal_metric: params - goal_value: 90% - num_solutions: 3 - -stages: - width_importance: - batch: 1 - replacement_scoring: - batch: 1 - instances: 1 - -output: - result_root: - -infrastructure: - gpus_per_node: 1 - execution_contract: - repository: - venv: - container: - container_mounts: - prerun_commands: [] - postrun_commands: [] diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py index fce62cb1b37..110a528adb8 100644 --- a/tests/gpu/torch/puzzletron/test_puzzletron.py +++ b/tests/gpu/torch/puzzletron/test_puzzletron.py @@ -39,8 +39,8 @@ from puzzletron_orchestrator.adapters.registry import adapter_for_stage -@pytest.mark.timeout(2400) -def test_tiny_qwen_campaign_uses_current_public_route( +@pytest.mark.timeout(1200) +def test_tiny_qwen_full_lifecycle_uses_public_bundle_contract( project_root_path: Path, tmp_path: Path, ) -> None: @@ -51,7 +51,7 @@ def test_tiny_qwen_campaign_uses_current_public_route( campaign = build_tiny_qwen_campaign(project_root_path, tmp_path) _assert_compiled_route(campaign) - completed = campaign.run(timeout=2100) + completed = campaign.run() result = campaign.require_success(completed) stage_ids = tuple(node.stage_id for node in campaign.compiled_plan.stages) @@ -60,7 +60,7 @@ def test_tiny_qwen_campaign_uses_current_public_route( assert manifest_paths pruning_paths = _assert_pruning_and_mip_artifacts(campaign) final_checkpoint, post_paths = _assert_post_mip_and_final_checkpoint(campaign) - _assert_selected_checkpoint_runs(final_checkpoint) + _assert_final_checkpoint_runs(final_checkpoint) _assert_final_report(campaign, result) state = CampaignStateStore(campaign.smoke_root) @@ -70,7 +70,7 @@ def test_tiny_qwen_campaign_uses_current_public_route( durable_paths = [*manifest_paths, *pruning_paths, *post_paths] before_resume = _artifact_digests(durable_paths) - resumed = campaign.run(timeout=120) + resumed = campaign.run(timeout=60) resumed_result = campaign.require_success(resumed) _assert_completed_campaign(campaign, resumed_result, stage_ids) @@ -84,17 +84,39 @@ def test_tiny_qwen_campaign_uses_current_public_route( def _assert_compiled_route(campaign: TinyQwenCampaign) -> None: """Verify the saved YAML bundle compiles to the intended one-GPU DAG.""" - for name in ("experiment.yaml", "runner.yaml", "execution.yaml"): + for name in ( + "recipe.yaml", + "experiment.runtime.yaml", + "experiment.resolved.yaml", + "runner.yaml", + "execution.yaml", + "manifest.json", + "provenance.json", + ): assert (campaign.smoke_bundle / name).is_file() - serving = campaign.config["post_mip"]["flows"][campaign.flow_id]["nodes"]["serving"] + flow_nodes = campaign.config["post_mip"]["flows"][campaign.flow_id]["nodes"] + assert campaign.config["sort"]["deferred_axes"] == [ + "kv_groups", + "q_heads_per_group", + "gdn_key_groups", + "gdn_value_heads_per_group", + "gdn_key_head_dim", + "gdn_value_head_dim", + ] + activation_passes = campaign.config["pruning"]["activation_passes"] + assert len(activation_passes) == 1 + assert activation_passes[0]["pruning_mixin"]["layer_descriptor"]["_target_"].endswith( + ".Qwen3P5TextFFNIntermediateLayerDescriptor" + ) + serving = flow_nodes["serving"] assert serving["config"]["allow_aiperf_v011_online_tokenizer_resolution"] is True assert serving["config"]["topology"]["extra_vllm_args"] == [ "-cc.cudagraph_mode=NONE", "--no-enable-flashinfer-autotune", "--max-num-batched-tokens", - "128", + "64", "--max-num-seqs", - "4", + "2", "--gpu-memory-utilization", "0.5", ] @@ -114,19 +136,38 @@ def _assert_compiled_route(campaign: TinyQwenCampaign) -> None: node for node in campaign.compiled_plan.stages if node.stage_id.startswith(prefix) ) assert tuple(node.stage_id.removeprefix(prefix) for node in post_nodes) == expected_nodes + assert flow_nodes["online_eval"]["config"]["eval_samples"] == 1 + assert flow_nodes["final_eval"]["config"]["eval_samples"] == 1 assert post_nodes[0].parents == ("mip",) for parent, node in pairwise(post_nodes): assert node.parents == (parent.stage_id,) + stage_ids = tuple(node.stage_id for node in campaign.compiled_plan.stages) + major_stages = ( + "convert", + "tokenize_data", + "width_importance", + "sort", + "build_library", + "replacement_scoring", + "mip", + ) + assert tuple(stage_id for stage_id in stage_ids if stage_id in major_stages) == major_stages + assert all(node.total_gpus <= 1 for node in campaign.compiled_plan.stages) + replacement_node = next( node for node in campaign.compiled_plan.stages if node.stage_id == "replacement_scoring" ) replacement_work = adapter_for_stage(replacement_node).plan( campaign.compiled_plan, replacement_node ) - assert (replacement_node.instances, replacement_node.gpus_per_instance) == (1, 1) - assert len(replacement_work.items) == 1 - assert replacement_work.items[0].metadata["worker_count"] == 1 + assert campaign.config["embedding_pruning"]["widths"] == [512, 256] + assert ( + replacement_node.strategy.value, + replacement_node.instances, + replacement_node.gpus_per_instance, + len(replacement_work.items), + ) == ("single", 1, 1, 1) def _assert_completed_campaign( @@ -160,18 +201,15 @@ def _assert_pruning_and_mip_artifacts(campaign: TinyQwenCampaign) -> list[Path]: replacement_summary_path = root / "artifacts/replacement_scoring/summary.json" replacement_summary = _json(replacement_summary_path) - assert replacement_summary["widths"] == [256] - assert replacement_summary["scenario_count"] == 1 - replacement_result_paths = list( - root.glob( - "scenarios/width-*/depth-*/distributed_eval/replacement_scoring/results/**/*.json" - ) - ) - replacement_results = [_json(path) for path in replacement_result_paths] - assert replacement_results - provenances = [result["provenance"] for result in replacement_results] - assert {provenance["score_device_type"] for provenance in provenances} == {"cuda"} - assert {provenance["visible_cuda_device_count"] for provenance in provenances} == {1} + expected_widths = campaign.config["embedding_pruning"]["widths"] + assert replacement_summary["widths"] == expected_widths + assert replacement_summary["scenario_count"] == len(expected_widths) + replacement_outputs = [ + Path(output) for child in replacement_summary["children"] for output in child["outputs"] + ] + assert replacement_summary["record_count"] > 0 + assert all(child["record_count"] > 0 for child in replacement_summary["children"]) + assert replacement_outputs and all(path.is_file() for path in replacement_outputs) candidate_library_path = root / "candidate_library.json" candidate_library = _json(candidate_library_path) @@ -193,18 +231,19 @@ def _assert_pruning_and_mip_artifacts(campaign: TinyQwenCampaign) -> list[Path]: if scenario["status"] == "feasible" ] solutions = [solution for path in solution_paths for solution in _json(path)] - assert len(solutions) >= 3 + assert solutions solution_ffn_widths = { int(width) for solution in solutions for width in _nested_values(solution["chosen_block_configs"], "intermediate_size") } - assert solution_ffn_widths.intersection({256, 512}) + configured_ffn_widths = set(campaign.config["pruning"]["intermediate_size_list"]) + assert solution_ffn_widths and solution_ffn_widths <= configured_ffn_widths return [ *pass_manifests, *score_files, replacement_summary_path, - *replacement_result_paths, + *replacement_outputs, candidate_library_path, active_profiles_path, *grid_paths, @@ -215,58 +254,34 @@ def _assert_pruning_and_mip_artifacts(campaign: TinyQwenCampaign) -> list[Path]: def _assert_post_mip_and_final_checkpoint( campaign: TinyQwenCampaign, ) -> tuple[Path, list[Path]]: - """Verify the post-MIP flow selects and trains one usable final checkpoint.""" + """Verify the post-MIP flow evaluates and trains one usable final checkpoint.""" root = campaign.smoke_root ledger = CandidateLedger(root / "artifacts/post_mip") online = ledger.load_candidate_set("online_eval") - assert len(online.revision_ids) >= 3 + assert len(online.revision_ids) == 1 online_observations = ledger.observations["online_eval"] assert all(row.status == "success" for row in online_observations.values()) assert all(Path(row.artifacts["result_path"]).is_file() for row in online_observations.values()) - online_losses = { - revision_id: _finite_metric(ledger, revision_id, "online_eval.lm_loss") - for revision_id in online.revision_ids - } - - best_lm = ledger.load_candidate_set("best_lm") - assert len(best_lm.revision_ids) == 3 - expected_best_lm = tuple( - revision_id - for _loss, revision_id in sorted( - (loss, revision_id) for revision_id, loss in online_losses.items() - )[:3] - ) - assert best_lm.revision_ids == expected_best_lm - assert { - revision_id - for revision_id, row in ledger.observations["best_lm"].items() - if row.status == "selected" - } == set(best_lm.revision_ids) + assert _finite_metric(ledger, online.revision_ids[0], "online_eval.lm_loss") >= 0 materialized = ledger.load_candidate_set("materialized") - assert len(materialized.revision_ids) == 3 + assert len(materialized.revision_ids) == 1 for revision_id in materialized.revision_ids: revision = ledger.revisions[revision_id] - assert revision.parent_revision_id in set(best_lm.revision_ids) checkpoint = Path(revision.artifact["checkpoint"]) assert (checkpoint / "config.json").is_file() assert list(checkpoint.glob("*.safetensors")) - assert { - ledger.revisions[revision_id].parent_revision_id - for revision_id in materialized.revision_ids - } == set(best_lm.revision_ids) - serving = ledger.load_candidate_set("serving") - assert len(serving.revision_ids) >= 2 - assert set(serving.revision_ids) <= set(materialized.revision_ids) - serving_throughputs = { - revision_id: _finite_metric( - ledger, revision_id, "serving.concurrency_1.output_token_throughput" + assert len(serving.revision_ids) == 1 + assert ( + _finite_metric( + ledger, + serving.revision_ids[0], + "serving.concurrency_1.output_token_throughput", ) - for revision_id in serving.revision_ids - } - assert all(throughput > 0 for throughput in serving_throughputs.values()) + > 0 + ) serving_observations = ledger.observations["serving"] aiperf_paths = [ Path(path) @@ -277,29 +292,19 @@ def _assert_post_mip_and_final_checkpoint( ] assert aiperf_paths and all(path.is_file() for path in aiperf_paths) - fastest = ledger.load_candidate_set("fastest") - assert len(fastest.revision_ids) == 2 - expected_fastest = tuple( - revision_id - for _throughput, revision_id in sorted( - (throughput, revision_id) for revision_id, throughput in serving_throughputs.items() - )[-2:][::-1] - ) - assert fastest.revision_ids == expected_fastest short_kd = ledger.load_candidate_set("short_kd") - assert len(short_kd.revision_ids) == 2 + assert len(short_kd.revision_ids) == 1 kd_paths = [] for revision_id in short_kd.revision_ids: revision = ledger.revisions[revision_id] - assert revision.parent_revision_id in set(materialized.revision_ids) checkpoint = Path(revision.artifact["checkpoint"]) summary_path = Path(revision.artifact["summary_path"]) kd_summary = _json(summary_path) records = kd_summary["records"] steps = {int(record.get("step", record.get("global_step"))) for record in records} losses = [float(record.get("loss", record.get("train_loss"))) for record in records] - assert kd_summary["max_steps"] == 2 - assert len(steps) >= 2 + assert kd_summary["max_steps"] == 1 + assert steps assert all(math.isfinite(loss) for loss in losses) assert (checkpoint / "config.json").is_file() assert list(checkpoint.glob("*.safetensors")) @@ -318,22 +323,11 @@ def _assert_post_mip_and_final_checkpoint( *checkpoint.glob("*.safetensors"), ] ) - assert { - ledger.revisions[revision_id].parent_revision_id for revision_id in short_kd.revision_ids - } == set(fastest.revision_ids) - final_eval = ledger.load_candidate_set("final_eval") - assert set(final_eval.revision_ids) == set(short_kd.revision_ids) - final_losses = { - revision_id: _finite_metric(ledger, revision_id, "final_eval.lm_loss") - for revision_id in final_eval.revision_ids - } - best = ledger.load_candidate_set("best") - assert len(best.revision_ids) == 1 - selected_id = best.revision_ids[0] - assert final_losses[selected_id] == min(final_losses.values()) - selected = ledger.revisions[selected_id] - final_checkpoint = Path(selected.artifact["checkpoint"]) + assert len(final_eval.revision_ids) == 1 + assert _finite_metric(ledger, final_eval.revision_ids[0], "final_eval.lm_loss") >= 0 + final_revision = ledger.revisions[final_eval.revision_ids[0]] + final_checkpoint = Path(final_revision.artifact["checkpoint"]) post_paths = [root / "artifacts/post_mip/candidate_registry.json"] post_paths.extend(root.glob("artifacts/post_mip/nodes/*/current.json")) @@ -345,17 +339,17 @@ def _assert_post_mip_and_final_checkpoint( return final_checkpoint, [*post_paths, *aiperf_paths, *kd_paths] -def _assert_selected_checkpoint_runs(final_checkpoint: Path) -> None: - """Reload the selected physical checkpoint and run a finite CUDA forward pass.""" +def _assert_final_checkpoint_runs(final_checkpoint: Path) -> None: + """Reload the final physical checkpoint and run a finite CUDA forward pass.""" resolution = resolve_descriptor_from_pretrained(str(final_checkpoint)) - selected_config = _json(final_checkpoint / "config.json") + final_config = _json(final_checkpoint / "config.json") widths_by_block = [ _nested_values(block_config, "intermediate_size") - for block_config in selected_config["block_configs"] + for block_config in final_config["block_configs"] ] assert all(len(widths) == 1 for widths in widths_by_block) - selected_widths = [int(widths[0]) for widths in widths_by_block] + final_widths = [int(widths[0]) for widths in widths_by_block] model = load_anymodel_for_scoring( str(final_checkpoint), anymodel_descriptor=resolution.name, @@ -364,7 +358,7 @@ def _assert_selected_checkpoint_runs(final_checkpoint: Path) -> None: local_files_only=True, ).cuda() - assert [layer.mlp.down_proj.in_features for layer in model.model.layers] == selected_widths + assert [layer.mlp.down_proj.in_features for layer in model.model.layers] == final_widths with torch.no_grad(): logits = model(torch.tensor([[1, 2, 3, 4]], device="cuda"), use_cache=False).logits assert torch.isfinite(logits).all() diff --git a/tests/gpu/torch/puzzletron/test_qwen3p5_0p8b_smoke.py b/tests/gpu/torch/puzzletron/test_qwen3p5_0p8b_smoke.py index 2507118b8fd..66547efa5be 100644 --- a/tests/gpu/torch/puzzletron/test_qwen3p5_0p8b_smoke.py +++ b/tests/gpu/torch/puzzletron/test_qwen3p5_0p8b_smoke.py @@ -32,6 +32,11 @@ import yaml from datasets import Dataset, DatasetDict +from puzzletron_orchestrator.recipe_config import ( + materialize_resolved_bundle, + resolve_recipe_run, + site_template, +) from tests._test_utils.torch.puzzletron.checkpoint_evaluation import ( assert_pruned_checkpoints_completed_benchmark, ) @@ -69,38 +74,44 @@ def test_qwen3p5_0p8b_orchestrated_full_smoke_completes( project_root_path: Path, tmp_path: Path, ) -> None: - """Run the bounded lifecycle on one H100 with network or populated benchmark caches.""" + """Run a bundle generated through the public configuration contract locally.""" dataset = tmp_path / "dataset" results = tmp_path / "results" cache = tmp_path / "cache" - runner = tmp_path / "runner.yaml" + recipe = tmp_path / "recipe.yaml" + site = tmp_path / "site.yaml" _save_messages_dataset(dataset) - runner.write_text( - yaml.safe_dump( - { - "runner": { - "kind": "slurm", - "slurm": {"account": "local-smoke", "max_nodes": 1}, - "execution_contract": { - "repository": str(project_root_path), - "venv": sys.prefix, - "container": None, - "container_mounts": None, - "prerun_commands": [], - "postrun_commands": [], - }, - } - }, - sort_keys=False, - ) + recipe_payload = yaml.safe_load( + ( + project_root_path / "examples/puzzletron/configs/recipes/qwen3p5_0p8b_text_smoke.yaml" + ).read_text() + ) + recipe_payload.update( + { + "run_root": str(results), + "resource_profile": "selected", + "data": {"path": str(dataset), "revision": "fixture-revision"}, + } ) + recipe.write_text(yaml.safe_dump(recipe_payload, sort_keys=False)) + site_payload = site_template() + site_payload["site"]["environment"].update( + {"repository": str(project_root_path), "venv": sys.prefix} + ) + site_payload["site"]["paths"]["hf_home"] = str(cache / "huggingface") + site_payload["site"]["slurm"].update({"account": "local-smoke", "partition": "local"}) + site_payload["resources"]["selected"] = { + "mode": "per_attempt", + "gpus_per_node": 1, + "max_nodes": 1, + } + site.write_text(yaml.safe_dump(site_payload, sort_keys=False)) + bundle = materialize_resolved_bundle(resolve_recipe_run(recipe, site), activate=True) environment = os.environ.copy() environment.update( { - "PUZZLETRON_RUN_ROOT": str(results), - "PUZZLETRON_DATASET_PATH": str(dataset), "HF_HOME": str(cache / "huggingface"), "HF_DATASETS_CACHE": str(cache / "datasets"), "TORCH_HOME": str(cache / "torch"), @@ -112,17 +123,11 @@ def test_qwen3p5_0p8b_orchestrated_full_smoke_completes( sys.executable, str(project_root_path / "examples/puzzletron/orchestrate.py"), "--experiment", - str( - project_root_path / "examples/puzzletron/configs/families/qwen3_5/" - "qwen3p5_0p8b/runs/full_smoke.yaml" - ), + str(bundle / "experiment.runtime.yaml"), "--runner", - str(runner), + str(bundle / "runner.yaml"), "--execution", - str( - project_root_path - / "examples/puzzletron/configs/orchestration/execution.single_gpu.yaml" - ), + str(bundle / "execution.yaml"), "--stage", "full", "--local", diff --git a/tests/gpu/torch/puzzletron/test_qwen3p5_0p8b_vlm_smoke.py b/tests/gpu/torch/puzzletron/test_qwen3p5_0p8b_vlm_smoke.py index 553af9844dd..41c93f859a6 100644 --- a/tests/gpu/torch/puzzletron/test_qwen3p5_0p8b_vlm_smoke.py +++ b/tests/gpu/torch/puzzletron/test_qwen3p5_0p8b_vlm_smoke.py @@ -31,13 +31,15 @@ from examples.puzzletron.evaluation.vlm import profile as evaluation_profile from examples.puzzletron.evaluation.vlm import suites as evaluation_suites from modelopt.torch.puzzletron.dataset.multimodal import materialize_normalized_conversation_samples +from puzzletron_orchestrator.recipe_config import ( + materialize_resolved_bundle, + resolve_recipe_run, + site_template, +) from tests._test_utils.torch.puzzletron.checkpoint_evaluation import ( assert_pruned_checkpoints_completed_benchmark, ) -RUN_PATH = "examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/full_vlm_smoke.yaml" -EXECUTION_PATH = "examples/puzzletron/configs/orchestration/execution.single_gpu.yaml" - def _materialize_image_conversations(path: Path) -> None: images = [] @@ -83,26 +85,17 @@ def _materialize_image_conversations(path: Path) -> None: assert manifest["sample_count"] == manifest["image_count"] == len(samples) -def _write_local_runner(path: Path, project_root: Path) -> None: - path.write_text( - yaml.safe_dump( - { - "runner": { - "kind": "slurm", - "slurm": {"account": "local-smoke", "max_nodes": 1}, - "execution_contract": { - "repository": str(project_root), - "venv": sys.prefix, - "container": None, - "container_mounts": None, - "prerun_commands": [], - "postrun_commands": [], - }, - } - }, - sort_keys=False, - ) - ) +def _write_local_site(path: Path, project_root: Path, hf_home: Path) -> None: + payload = site_template() + payload["site"]["environment"].update({"repository": str(project_root), "venv": sys.prefix}) + payload["site"]["paths"]["hf_home"] = str(hf_home) + payload["site"]["slurm"].update({"account": "local-smoke", "partition": "local"}) + payload["resources"]["selected"] = { + "mode": "per_attempt", + "gpus_per_node": 1, + "max_nodes": 1, + } + path.write_text(yaml.safe_dump(payload, sort_keys=False)) @pytest.mark.integration @@ -119,7 +112,7 @@ def test_qwen3p5_0p8b_orchestrated_vlm_full_smoke_completes( tmp_path: Path, monkeypatch, ) -> None: - """Run the bounded VLM lifecycle on one H100 with populated offline benchmark caches.""" + """Run a VLM bundle generated through the public configuration contract locally.""" dataset = tmp_path / "dataset" results = tmp_path / "results" @@ -139,10 +132,25 @@ def test_qwen3p5_0p8b_orchestrated_vlm_full_smoke_completes( ) except ValueError as error: pytest.fail(str(error)) - runner = tmp_path / "runner.yaml" + recipe = tmp_path / "recipe.yaml" + site = tmp_path / "site.yaml" _materialize_image_conversations(dataset) - _write_local_runner(runner, project_root_path) - monkeypatch.setenv("PUZZLETRON_RUN_ROOT", str(results)) + recipe_payload = yaml.safe_load( + ( + project_root_path / "examples/puzzletron/configs/recipes/qwen3p5_0p8b_vlm_smoke.yaml" + ).read_text() + ) + recipe_payload.update( + { + "run_root": str(results), + "resource_profile": "selected", + "data": {"path": str(dataset), "revision": "fixture-revision"}, + "advanced": {"experiment": {"prepare_dataset.enabled": False}}, + } + ) + recipe.write_text(yaml.safe_dump(recipe_payload, sort_keys=False)) + _write_local_site(site, project_root_path, benchmark_hf_home_path) + bundle = materialize_resolved_bundle(resolve_recipe_run(recipe, site), activate=True) environment = os.environ.copy() environment.update( { @@ -158,21 +166,13 @@ def test_qwen3p5_0p8b_orchestrated_vlm_full_smoke_completes( sys.executable, str(project_root_path / "examples/puzzletron/orchestrate.py"), "--experiment", - str(project_root_path / RUN_PATH), + str(bundle / "experiment.runtime.yaml"), "--runner", - str(runner), + str(bundle / "runner.yaml"), "--execution", - str(project_root_path / EXECUTION_PATH), + str(bundle / "execution.yaml"), "--stage", "full", - "--override", - "prepare_dataset.enabled=false", - "--override", - f"dataset_path={dataset}", - "--override", - f"data.path={dataset}", - "--override", - "data.revision=fixture-revision", "--local", "--color", "never", @@ -199,7 +199,6 @@ def test_qwen3p5_0p8b_orchestrated_vlm_full_smoke_completes( json.loads(path.read_text(encoding="utf-8")) for path in activation_root.glob("**/args.json") ] - assert activation_markers vision_markers = [ marker for marker in activation_markers @@ -208,7 +207,6 @@ def test_qwen3p5_0p8b_orchestrated_vlm_full_smoke_completes( assert vision_markers for marker in vision_markers: observability = marker["observability"] - assert observability["vision_forward_count"] > 0 assert observability["vision_output_checksums"] assert observability["batch_fingerprints"] diff --git a/tests/unit/torch/puzzletron/test_automodel_config.py b/tests/unit/torch/puzzletron/test_automodel_config.py index e0cb2ba4f2c..28be69407aa 100644 --- a/tests/unit/torch/puzzletron/test_automodel_config.py +++ b/tests/unit/torch/puzzletron/test_automodel_config.py @@ -51,10 +51,36 @@ def _cfg(method="independent", eval_samples=200, micro_batch_size=2): }, }, }, + "scoring": { + "block_size": 32, + "micro_batch_size": 1, + "automodel": { + "force_hf": False, + "parallel": { + "tp": 1, + "cp": 1, + "pp": 1, + "ep": 1, + "dp_shard": 1, + "dp_replicate": 1, + }, + }, + }, } ) +def _disable_descriptor_recipe_injection(monkeypatch): + monkeypatch.setattr( + "modelopt.torch.puzzletron.plugins.automodel.config._inject_descriptor_model_kwargs", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "modelopt.torch.puzzletron.plugins.automodel.config.inject_descriptor_pipeline_config", + lambda *args, **kwargs: None, + ) + + def test_build_recipe_config_generates_recipe_from_stage_parallelism(monkeypatch): cfg = _cfg() cfg.pruning.automodel.parallel = { @@ -67,14 +93,7 @@ def test_build_recipe_config_generates_recipe_from_stage_parallelism(monkeypatch "sequence_parallel": False, "pipeline_schedule": "1f1b", } - monkeypatch.setattr( - "modelopt.torch.puzzletron.plugins.automodel.config._inject_descriptor_model_kwargs", - lambda *args, **kwargs: None, - ) - monkeypatch.setattr( - "modelopt.torch.puzzletron.plugins.automodel.config.inject_descriptor_pipeline_config", - lambda *args, **kwargs: None, - ) + _disable_descriptor_recipe_injection(monkeypatch) recipe = build_recipe_config(cfg) @@ -100,34 +119,33 @@ def test_build_recipe_config_generates_recipe_from_stage_parallelism(monkeypatch def test_authored_automodel_backend_reaches_scoring_and_solution_recipes(monkeypatch): cfg = _cfg() cfg.model = {"automodel_backend": {"attn": "sdpa"}} - cfg.scoring = { - "block_size": 32, - "micro_batch_size": 1, - "automodel": { - "force_hf": False, - "parallel": { - "tp": 1, - "cp": 1, - "pp": 1, - "ep": 1, - "dp_shard": 1, - "dp_replicate": 1, - }, - }, - } - monkeypatch.setattr( - "modelopt.torch.puzzletron.plugins.automodel.config._inject_descriptor_model_kwargs", - lambda *args, **kwargs: None, - ) - monkeypatch.setattr( - "modelopt.torch.puzzletron.plugins.automodel.config.inject_descriptor_pipeline_config", - lambda *args, **kwargs: None, - ) + _disable_descriptor_recipe_injection(monkeypatch) assert build_recipe_config(cfg)["model"]["backend"] == {"attn": "sdpa"} assert build_solution_recipe_config(cfg, "/checkpoint")["model"]["backend"] == {"attn": "sdpa"} +@pytest.mark.parametrize( + "builder", + [ + build_recipe_config, + lambda cfg: build_solution_recipe_config(cfg, "/checkpoint"), + ], + ids=["activation-scoring", "solution-scoring"], +) +def test_recipe_builders_route_checkpoint_bookkeeping_below_puzzle_root(monkeypatch, builder): + cfg = _cfg() + cfg.puzzle_dir = "/selected/run-root" + _disable_descriptor_recipe_injection(monkeypatch) + + recipe = builder(cfg) + + assert recipe["checkpoint"] == { + "enabled": False, + "checkpoint_dir": "/selected/run-root/.runtime/nemo_automodel/checkpoints", + } + + def test_build_recipe_config_rejects_pure_ddp_replication(): cfg = _cfg() cfg.pruning.automodel.parallel.dp_shard = 1 @@ -169,14 +187,7 @@ def test_solution_recipe_uses_inferred_runtime_descriptor(monkeypatch): }, } ) - monkeypatch.setattr( - "modelopt.torch.puzzletron.plugins.automodel.config._inject_descriptor_model_kwargs", - lambda *args, **kwargs: None, - ) - monkeypatch.setattr( - "modelopt.torch.puzzletron.plugins.automodel.config.inject_descriptor_pipeline_config", - lambda *args, **kwargs: None, - ) + _disable_descriptor_recipe_injection(monkeypatch) recipe = build_solution_recipe_config(cfg, "/checkpoint") diff --git a/tests/unit/torch/puzzletron/test_hydra_utils.py b/tests/unit/torch/puzzletron/test_hydra_utils.py index 2aa1af283ae..54738f9b612 100644 --- a/tests/unit/torch/puzzletron/test_hydra_utils.py +++ b/tests/unit/torch/puzzletron/test_hydra_utils.py @@ -68,7 +68,10 @@ def test_static_workload_stats_preserves_resolved_python_objects(monkeypatch): monkeypatch.setattr(calc_subblock_stats, "launch_calc_subblock_stats", launched.append) pipeline_stages._calculate_static_workload_stats( - {"mip": {"workloads": {"interactive": {"isl": 256, "osl": 64, "batch_size": 2}}}}, + { + "embedding_pruning": {"widths": [512, 256]}, + "mip": {"workloads": {"interactive": {"isl": 256, "osl": 64, "batch_size": 2}}}, + }, hydra_cfg, ) @@ -80,6 +83,7 @@ def test_static_workload_stats_preserves_resolved_python_objects(monkeypatch): assert selected.calc_subblock_stats.generation_seq_len == 64 assert selected.calc_subblock_stats.runtime_stats.enabled is False assert selected.calc_subblock_stats.merge_with_existing_stats is True + assert list(selected.calc_subblock_stats.model_hidden_sizes) == [512, 256] def test_warmup_steps_casts_inputs_before_computing(): diff --git a/tests/unit/torch/puzzletron/test_portable_configs.py b/tests/unit/torch/puzzletron/test_portable_configs.py index ab6d20ec754..b41af708b2c 100644 --- a/tests/unit/torch/puzzletron/test_portable_configs.py +++ b/tests/unit/torch/puzzletron/test_portable_configs.py @@ -21,17 +21,9 @@ import pytest import yaml -from puzzletron_orchestrator.compiler import load_execution_config, load_runner_config -from puzzletron_setup import WORKER_REPOSITORY_PLACEHOLDER, WORKER_VENV_PLACEHOLDER -from puzzletron_setup.v2.defaults import load_defaults +from puzzletron_orchestrator.recipe_config import ROUTES, recipe_template, site_template REPOSITORY_ROOT = Path(__file__).resolve().parents[4] -SLURM_RUNNER_CONFIGS = ( - "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", - "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml", - "examples/puzzletron/configs/orchestration/qwen3p5_0p8b/runner.slurm.yaml", -) -NAMED_SLURM_RUNNER_CONFIGS = SLURM_RUNNER_CONFIGS[1:] NEMOTRON3_NANO_30B_MODEL_CONFIG = ( "examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml" ) @@ -42,97 +34,54 @@ ) -def test_slurm_runner_example_is_portable() -> None: - path = REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml" - slurm = load_runner_config(path) - assert slurm.contract.repository == "/opt/puzzletron/src/modelopt" - assert slurm.contract.venv == "/venv" - assert slurm.contract.container.startswith("REPLACE_WITH_") - assert slurm.contract.container_mounts is None - assert not slurm.contract.prerun_commands - assert slurm.slurm is not None - assert slurm.slurm.account.startswith("REPLACE_WITH_") - assert slurm.slurm.partition == ( - "REPLACE_WITH_PRIMARY_SLURM_PARTITION,REPLACE_WITH_ALTERNATE_SLURM_PARTITION" +def test_site_example_is_the_only_portable_environment_contract() -> None: + checked_in = yaml.safe_load( + (REPOSITORY_ROOT / "examples/puzzletron/configs/site.example.yaml").read_text() ) - assert slurm.slurm.log_dir == "logs" - -def test_baremetal_runner_example_is_portable() -> None: - baremetal = load_runner_config( - REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml" - ) - assert baremetal.contract.repository == WORKER_REPOSITORY_PLACEHOLDER - assert baremetal.contract.venv == WORKER_VENV_PLACEHOLDER - assert baremetal.contract.setup_env is None - assert baremetal.baremetal is not None - hostnames = [host.hostname for host in baremetal.baremetal.hosts] - assert hostnames - assert baremetal.baremetal.rendezvous_host in hostnames - assert all(hostname.startswith("REPLACE_WITH_") for hostname in hostnames) - - -def test_qwen_slurm_runner_preserves_portable_environment_contract() -> None: - path = REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml" - runner = load_runner_config(path) - - contract_values = ( - runner.contract.repository, - runner.contract.venv, - runner.contract.container, - runner.contract.container_mounts, + assert checked_in == site_template() + environment = checked_in["site"]["environment"] + assert environment["repository"].startswith("REPLACE_WITH_") + assert environment["venv"].startswith("REPLACE_WITH_") + assert environment["container"] is None + assert environment["container_mounts"] is None + assert not environment["prerun_commands"] + assert checked_in["site"]["slurm"]["account"].startswith("REPLACE_WITH_") + assert checked_in["site"]["slurm"]["partition"].startswith("REPLACE_WITH_") + assert checked_in["resources"]["multinode"] == { + "mode": "per_attempt", + "gpus_per_node": 8, + "max_nodes": 64, + } + + +@pytest.mark.parametrize("route", ROUTES, ids=lambda route: route.route_id) +def test_recipe_template_is_small_and_has_no_inheritance(route) -> None: + recipe = recipe_template( + model=route.model, + workflow=route.workflow, + mode=route.mode, ) - assert contract_values[:2] == ("/opt/puzzletron/src/modelopt", "/venv") - assert all(value and value.startswith("REPLACE_WITH_") for value in contract_values[2:]) - assert not runner.contract.prerun_commands - assert runner.slurm is not None - assert runner.slurm.account.startswith("REPLACE_WITH_") - assert runner.slurm.partition.startswith("REPLACE_WITH_") - - -@pytest.mark.parametrize("relative_path", SLURM_RUNNER_CONFIGS) -def test_checked_in_slurm_runners_only_emit_generic_partitions( - relative_path: str, -) -> None: - payload = yaml.safe_load((REPOSITORY_ROOT / relative_path).read_text()) - partition_keys = {key for key in payload["runner"]["slurm"] if key.startswith("partition")} - - assert "partition" in payload["runner"]["slurm"] - assert partition_keys <= {"partition", "partition_cpu"} - - -@pytest.mark.parametrize("relative_path", NAMED_SLURM_RUNNER_CONFIGS) -def test_named_slurm_runners_keep_logs_below_the_campaign_root( - relative_path: str, -) -> None: - payload = yaml.safe_load((REPOSITORY_ROOT / relative_path).read_text()) - - assert "log_dir" not in payload["runner"]["slurm"] - -def test_execution_example_is_loadable() -> None: - path = REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/execution.example.yaml" - - execution = load_execution_config(path) - - assert set(execution) >= {"defaults", "stages"} - - -def test_setup_defaults_example_is_portable() -> None: - path = REPOSITORY_ROOT / "examples/puzzletron/configs/setup/defaults.example.yaml" - - defaults = load_defaults(path) - - contract = defaults["infrastructure"]["execution_contract"] - assert contract["repository"] == WORKER_REPOSITORY_PLACEHOLDER - assert contract["venv"] == WORKER_VENV_PLACEHOLDER - assert contract["container"] is None - assert contract["container_mounts"] is None - assert not contract["prerun_commands"] - - slurm = defaults["infrastructure"]["runner"]["slurm"] - assert "account" not in slurm - assert slurm["partition"] is None + expected = { + "schema_version", + "name", + "model", + "workflow", + "mode", + "run_root", + "resource_profile", + } + if route.requires_data: + expected.add("data") + assert recipe["data"]["path"].startswith("REPLACE_WITH_") + assert recipe["data"]["revision"].startswith("REPLACE_WITH_") + assert set(recipe) == expected + assert (recipe["model"], recipe["workflow"], recipe["mode"]) == ( + route.model, + route.workflow, + route.mode, + ) def test_model_examples_use_public_hugging_face_identities() -> None: diff --git a/tests/unit/torch/puzzletron/test_qwen3p5_0p8b_full_smoke_plan.py b/tests/unit/torch/puzzletron/test_qwen3p5_0p8b_full_smoke_plan.py index e84ff0184dd..58fb6dc7439 100644 --- a/tests/unit/torch/puzzletron/test_qwen3p5_0p8b_full_smoke_plan.py +++ b/tests/unit/torch/puzzletron/test_qwen3p5_0p8b_full_smoke_plan.py @@ -17,39 +17,54 @@ from pathlib import Path +import yaml + from puzzletron_orchestrator.compiler import ( compile_campaign_plan, load_execution_config, load_runner_config, ) +from puzzletron_orchestrator.recipe_config import ( + materialize_resolved_bundle, + resolve_recipe_run, + site_template, +) REPOSITORY_ROOT = Path(__file__).resolve().parents[4] -FAMILY_ROOT = REPOSITORY_ROOT / "examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b" -RUN_PATH = FAMILY_ROOT / "runs/full_smoke.yaml" -RUNNER_PATH = ( - REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/qwen3p5_0p8b/runner.slurm.yaml" -) -SINGLE_GPU_EXECUTION_PATH = ( - REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/execution.single_gpu.yaml" -) +RECIPE_PATH = REPOSITORY_ROOT / "examples/puzzletron/configs/recipes/qwen3p5_0p8b_text_smoke.yaml" -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")) +def _compile(tmp_path: Path): + recipe = yaml.safe_load(RECIPE_PATH.read_text()) + recipe["run_root"] = str(tmp_path / "run") + recipe["data"] = { + "path": str(tmp_path / "prepared-data"), + "revision": "fixture-revision", + } + recipe_path = tmp_path / "recipe.yaml" + recipe_path.write_text(yaml.safe_dump(recipe, sort_keys=False)) + site = site_template() + site["site"]["environment"].update({"repository": str(REPOSITORY_ROOT), "venv": ".venv"}) + site["site"]["paths"]["hf_home"] = str(tmp_path / "hf") + site["site"]["slurm"].update({"account": "test", "partition": "test"}) + site_path = tmp_path / "site.yaml" + site_path.write_text(yaml.safe_dump(site, sort_keys=False)) + bundle = materialize_resolved_bundle(resolve_recipe_run(recipe_path, site_path), activate=False) return compile_campaign_plan( - experiment_config_path=experiment, - runner=load_runner_config(RUNNER_PATH), - execution=load_execution_config(execution), + experiment_config_path=bundle / "experiment.runtime.yaml", + runner=load_runner_config(bundle / "runner.yaml"), + execution=load_execution_config(bundle / "execution.yaml"), stage_filter="full", ) -def test_full_smoke_compiles_one_complete_bounded_lifecycle(monkeypatch, tmp_path: Path) -> None: - plan = _compile(monkeypatch, tmp_path, RUN_PATH, SINGLE_GPU_EXECUTION_PATH) +def test_full_smoke_compiles_one_complete_bounded_lifecycle(tmp_path: Path) -> None: + plan = _compile(tmp_path) stages = {stage.stage_id: stage for stage in plan.stages} nodes = plan.experiment_config["post_mip"]["flows"]["params-90"]["nodes"] + assert plan.experiment_config["dataset_path"] == str(tmp_path / "prepared-data") + assert plan.experiment_config["data"]["revision"] == "fixture-revision" assert tuple(node for node in stages if node.startswith("post.")) == ( "post.params-90.online_eval", "post.params-90.best_lm", 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 15dbd068062..58896f99dc2 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 @@ -17,47 +17,52 @@ from pathlib import Path +import yaml + from examples.puzzletron.evaluation.vlm import contracts, suites from puzzletron_orchestrator.compiler import ( compile_campaign_plan, load_execution_config, load_runner_config, ) +from puzzletron_orchestrator.recipe_config import ( + materialize_resolved_bundle, + resolve_recipe_run, + site_template, +) from puzzletron_orchestrator.schema import ExecutionMode, ExecutionStrategy REPOSITORY_ROOT = Path(__file__).resolve().parents[4] -FAMILY_ROOT = REPOSITORY_ROOT / "examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b" -SMOKE_PATH = FAMILY_ROOT / "runs/full_vlm_smoke.yaml" -CAMPAIGN_PATH = FAMILY_ROOT / "runs/vlm_campaign.yaml" -RUNNER_PATH = ( - REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/qwen3p5_0p8b/runner.slurm.yaml" -) -SMOKE_EXECUTION_PATH = ( - REPOSITORY_ROOT - / "examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_smoke.yaml" +SMOKE_RECIPE_PATH = ( + REPOSITORY_ROOT / "examples/puzzletron/configs/recipes/qwen3p5_0p8b_vlm_smoke.yaml" ) -CAMPAIGN_EXECUTION_PATH = ( - REPOSITORY_ROOT - / "examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.vlm_campaign.yaml" +CAMPAIGN_RECIPE_PATH = ( + REPOSITORY_ROOT / "examples/puzzletron/configs/recipes/qwen3p5_0p8b_vlm_campaign.yaml" ) -def _compile(monkeypatch, tmp_path: Path, experiment: Path, execution: Path): - monkeypatch.setenv("PUZZLETRON_RUN_ROOT", str(tmp_path / experiment.stem)) - monkeypatch.setenv("PUZZLETRON_DATASET_REVISION", "fixture-revision") - monkeypatch.setenv("HF_HOME", str(tmp_path / "hf-home")) +def _compile(tmp_path: Path, recipe_source: Path): + recipe = yaml.safe_load(recipe_source.read_text()) + recipe["run_root"] = str(tmp_path / recipe_source.stem) + recipe_path = tmp_path / recipe_source.name + recipe_path.write_text(yaml.safe_dump(recipe, sort_keys=False)) + site = site_template() + site["site"]["environment"].update({"repository": str(REPOSITORY_ROOT), "venv": ".venv"}) + site["site"]["paths"]["hf_home"] = str(tmp_path / "hf-home") + site["site"]["slurm"].update({"account": "test", "partition": "test"}) + site_path = tmp_path / "site.yaml" + site_path.write_text(yaml.safe_dump(site, sort_keys=False)) + bundle = materialize_resolved_bundle(resolve_recipe_run(recipe_path, site_path), activate=False) return compile_campaign_plan( - experiment_config_path=experiment, - runner=load_runner_config(RUNNER_PATH), - execution=load_execution_config(execution), + experiment_config_path=bundle / "experiment.runtime.yaml", + runner=load_runner_config(bundle / "runner.yaml"), + execution=load_execution_config(bundle / "execution.yaml"), stage_filter="full", ) -def test_full_vlm_smoke_compiles_one_complete_bounded_lifecycle( - monkeypatch, tmp_path: Path -) -> None: - plan = _compile(monkeypatch, tmp_path, SMOKE_PATH, SMOKE_EXECUTION_PATH) +def test_full_vlm_smoke_compiles_one_complete_bounded_lifecycle(tmp_path: Path) -> None: + plan = _compile(tmp_path, SMOKE_RECIPE_PATH) stages = {stage.stage_id: stage for stage in plan.stages} config = plan.experiment_config nodes = config["post_mip"]["flows"]["params-90"]["nodes"] @@ -87,8 +92,8 @@ def test_full_vlm_smoke_compiles_one_complete_bounded_lifecycle( assert all(stage.total_gpus == 1 for stage in stages.values() if stage.resource != "cpu") -def test_vlm_campaign_compiles_the_multi_axis_flow(monkeypatch, tmp_path: Path) -> None: - plan = _compile(monkeypatch, tmp_path, CAMPAIGN_PATH, CAMPAIGN_EXECUTION_PATH) +def test_vlm_campaign_compiles_the_multi_axis_flow(tmp_path: Path) -> None: + plan = _compile(tmp_path, CAMPAIGN_RECIPE_PATH) stages = {stage.stage_id: stage for stage in plan.stages} config = plan.experiment_config candidates = config["post_mip"]["flows"]["candidates"]["nodes"] diff --git a/tests/unit/torch/puzzletron/test_qwen3p5_4b_vlm_example.py b/tests/unit/torch/puzzletron/test_qwen3p5_4b_vlm_recipes.py similarity index 52% rename from tests/unit/torch/puzzletron/test_qwen3p5_4b_vlm_example.py rename to tests/unit/torch/puzzletron/test_qwen3p5_4b_vlm_recipes.py index 9a9c62ea5ea..d4000d47707 100644 --- a/tests/unit/torch/puzzletron/test_qwen3p5_4b_vlm_example.py +++ b/tests/unit/torch/puzzletron/test_qwen3p5_4b_vlm_recipes.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""CPU contracts for the maintained Qwen 3.5 4B VLM example.""" +"""CPU contracts for the maintained Qwen 3.5 4B VLM recipes.""" from itertools import pairwise from pathlib import Path @@ -25,52 +25,51 @@ load_execution_config, load_runner_config, ) +from puzzletron_orchestrator.recipe_config import ( + materialize_resolved_bundle, + resolve_recipe_run, + site_template, +) REPOSITORY_ROOT = Path(__file__).resolve().parents[4] FAMILY_ROOT = REPOSITORY_ROOT / "examples/puzzletron/configs/families/qwen3_5/qwen3p5_4b" MODEL_PATH = FAMILY_ROOT / "model.yaml" -MIP_RUN_PATH = FAMILY_ROOT / "runs/mip_vlm_smoke.yaml" -LIFECYCLE_RUN_PATH = FAMILY_ROOT / "runs/vlm_lifecycle_smoke.yaml" -LEGACY_FULL_RUN_PATH = FAMILY_ROOT / "runs/full_vlm_smoke.yaml" -CAMPAIGN_RUN_PATH = FAMILY_ROOT / "runs/ffn_width_10to20pct_kd_search.yaml" -LEGACY_CAMPAIGN_RUN_PATH = FAMILY_ROOT / "runs/vlm_campaign.yaml" -RUNNER_PATH = ( - REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/qwen3p5_4b/runner.slurm.yaml" -) -EXECUTION_PATH = ( - REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/execution.single_gpu.yaml" -) -LIFECYCLE_EXECUTION_PATH = ( - REPOSITORY_ROOT - / "examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.vlm_lifecycle_smoke.yaml" -) -LEGACY_FULL_EXECUTION_PATH = ( - REPOSITORY_ROOT - / "examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.full_vlm_smoke.yaml" -) -CAMPAIGN_EXECUTION_PATH = REPOSITORY_ROOT / ( - "examples/puzzletron/configs/orchestration/qwen3p5_4b/" - "execution.ffn_width_10to20pct_kd_search.yaml" +SMOKE_RECIPE_PATH = ( + REPOSITORY_ROOT / "examples/puzzletron/configs/recipes/qwen3p5_4b_vlm_smoke.yaml" ) -LEGACY_CAMPAIGN_EXECUTION_PATH = ( - REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/qwen3p5_4b/execution.campaign.yaml" +CAMPAIGN_RECIPE_PATH = ( + REPOSITORY_ROOT / "examples/puzzletron/configs/recipes/qwen3p5_4b_vlm_campaign.yaml" ) def _compile_plan( - monkeypatch, tmp_path: Path, - run_path: Path, - execution_path=EXECUTION_PATH, - run_root_name: str | None = None, + recipe_source: Path, ): - monkeypatch.setenv("PUZZLETRON_RUN_ROOT", str(tmp_path / (run_root_name or run_path.stem))) - monkeypatch.setenv("PUZZLETRON_DATASET_PATH", str(tmp_path / "dataset")) - monkeypatch.setenv("PUZZLETRON_DATASET_REVISION", "fixture-revision") + run_root = tmp_path / recipe_source.stem + dataset = tmp_path / "dataset" + recipe = yaml.safe_load(recipe_source.read_text()) + recipe["run_root"] = str(run_root) + recipe["resource_profile"] = "selected" + recipe["data"] = {"path": str(dataset), "revision": "fixture-revision"} + recipe_path = tmp_path / f"{run_root.name}.recipe.yaml" + recipe_path.write_text(yaml.safe_dump(recipe, sort_keys=False)) + site = site_template() + site["site"]["environment"].update({"repository": str(REPOSITORY_ROOT), "venv": ".venv"}) + site["site"]["paths"]["hf_home"] = str(tmp_path / "hf") + site["site"]["slurm"].update({"account": "test", "partition": "test"}) + site["resources"]["selected"] = { + "mode": "per_attempt", + "gpus_per_node": 8, + "max_nodes": 1, + } + site_path = tmp_path / "site.yaml" + site_path.write_text(yaml.safe_dump(site, sort_keys=False)) + bundle = materialize_resolved_bundle(resolve_recipe_run(recipe_path, site_path), activate=False) return compile_campaign_plan( - experiment_config_path=run_path, - runner=load_runner_config(RUNNER_PATH), - execution=load_execution_config(execution_path), + experiment_config_path=bundle / "experiment.runtime.yaml", + runner=load_runner_config(bundle / "runner.yaml"), + execution=load_execution_config(bundle / "execution.yaml"), stage_filter="full", ) @@ -89,42 +88,12 @@ def test_qwen3p5_4b_model_pins_the_bounded_ffn_grid() -> None: assert model["pruning"] == {"intermediate_size_list": widths} -def test_qwen3p5_4b_default_compiles_the_complete_ffn_grid_and_stops_at_mip( - monkeypatch, - tmp_path: Path, -) -> None: - plan = _compile_plan(monkeypatch, tmp_path, MIP_RUN_PATH) - config = plan.experiment_config - stage_ids = tuple(stage.stage_id for stage in plan.stages) - - assert stage_ids == ( - "convert", - "width_importance", - "sort", - "sort_sanity", - "width_sanity", - "slicing_sanity", - "build_library", - "replacement_scoring", - "mip", - ) - assert config["vllm_stats"]["enabled"] is False - runs = config["mip"]["runs"] - assert set(runs) == {"params-80", "memory-85"} - assert all(run["solver"]["num_solutions"] == 3 for run in runs.values()) - assert all(run["homogeneous"]["keep"] == 2 for run in runs.values()) - assert config["post_mip"]["flows"] == {} - - -def test_qwen3p5_4b_opt_in_lifecycle_materializes_reloads_and_bounds_kd_and_evaluation( - monkeypatch, +def test_qwen3p5_4b_smoke_materializes_reloads_and_bounds_kd_and_evaluation( tmp_path: Path, ) -> None: plan = _compile_plan( - monkeypatch, tmp_path, - LIFECYCLE_RUN_PATH, - execution_path=LIFECYCLE_EXECUTION_PATH, + SMOKE_RECIPE_PATH, ) config = plan.experiment_config post_stages = tuple(stage for stage in plan.stages if stage.stage_id.startswith("post.")) @@ -155,41 +124,10 @@ def test_qwen3p5_4b_opt_in_lifecycle_materializes_reloads_and_bounds_kd_and_eval ) -def test_qwen3p5_4b_legacy_full_smoke_alias_resolves_to_lifecycle( - monkeypatch, - tmp_path: Path, -) -> None: - lifecycle_plan = _compile_plan( - monkeypatch, - tmp_path, - LIFECYCLE_RUN_PATH, - execution_path=LIFECYCLE_EXECUTION_PATH, - run_root_name="lifecycle-alias-contract", - ) - legacy_plan = _compile_plan( - monkeypatch, - tmp_path, - LEGACY_FULL_RUN_PATH, - execution_path=LIFECYCLE_EXECUTION_PATH, - run_root_name="lifecycle-alias-contract", - ) - - assert legacy_plan.experiment_config["post_mip"] == lifecycle_plan.experiment_config["post_mip"] - assert ( - legacy_plan.experiment_config["global_distillation"] - == (lifecycle_plan.experiment_config["global_distillation"]) - ) - assert load_execution_config(LEGACY_FULL_EXECUTION_PATH) == load_execution_config( - LIFECYCLE_EXECUTION_PATH - ) - - -def test_qwen3p5_4b_campaign_compares_pruning_bands_and_teacher(monkeypatch, tmp_path) -> None: +def test_qwen3p5_4b_campaign_compares_pruning_bands_and_teacher(tmp_path) -> None: plan = _compile_plan( - monkeypatch, tmp_path, - CAMPAIGN_RUN_PATH, - execution_path=CAMPAIGN_EXECUTION_PATH, + CAMPAIGN_RECIPE_PATH, ) config = plan.experiment_config candidates = config["mip"]["runs"]["ffn-candidates"] @@ -227,29 +165,3 @@ def test_qwen3p5_4b_campaign_compares_pruning_bands_and_teacher(monkeypatch, tmp assert stages["post.candidate-evaluation.screening_kd"].total_gpus == 8 assert all(stages[stage_id].gpus_per_node == 8 for stage_id in candidate_stages) assert stages["post.candidate-evaluation.global_kd"].total_gpus == 2 - - -def test_qwen3p5_4b_legacy_campaign_alias_resolves_to_named_search( - monkeypatch, - tmp_path: Path, -) -> None: - campaign_plan = _compile_plan( - monkeypatch, - tmp_path, - CAMPAIGN_RUN_PATH, - execution_path=CAMPAIGN_EXECUTION_PATH, - run_root_name="campaign-alias-contract", - ) - legacy_plan = _compile_plan( - monkeypatch, - tmp_path, - LEGACY_CAMPAIGN_RUN_PATH, - execution_path=CAMPAIGN_EXECUTION_PATH, - run_root_name="campaign-alias-contract", - ) - - assert legacy_plan.experiment_config["mip"] == campaign_plan.experiment_config["mip"] - assert legacy_plan.experiment_config["post_mip"] == campaign_plan.experiment_config["post_mip"] - assert load_execution_config(LEGACY_CAMPAIGN_EXECUTION_PATH) == load_execution_config( - CAMPAIGN_EXECUTION_PATH - ) diff --git a/tests/unit/torch/puzzletron/test_recipe_config.py b/tests/unit/torch/puzzletron/test_recipe_config.py new file mode 100644 index 00000000000..b491f4d14f1 --- /dev/null +++ b/tests/unit/torch/puzzletron/test_recipe_config.py @@ -0,0 +1,430 @@ +# 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. + +"""Behavioral tests for the concise Puzzletron recipe and site contract.""" + +from __future__ import annotations + +import json +import os +import subprocess +from collections import Counter +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from threading import Barrier + +import pytest +import yaml + +from examples.puzzletron import puzzletron as public_cli +from puzzletron_orchestrator import recipe_config +from puzzletron_orchestrator.recipe_config import ( + bundle_for_run_root, + materialize_resolved_bundle, + recipe_template, + resolve_recipe_run, +) + +REPOSITORY_ROOT = Path(__file__).resolve().parents[4] + + +@pytest.fixture(autouse=True) +def _stable_source_identity(monkeypatch): + monkeypatch.setattr( + recipe_config, + "_code_revision", + lambda: {"revision": "a" * 40, "dirty": False}, + ) + + +def _write_yaml(path: Path, payload: dict) -> Path: + path.write_text(yaml.safe_dump(payload, sort_keys=False)) + return path + + +def _site(tmp_path: Path, *, max_nodes: int = 1, mode: str = "reusable_allocation") -> Path: + return _write_yaml( + tmp_path / "site.yaml", + { + "schema_version": 1, + "site": { + "kind": "slurm", + "environment": {"repository": str(REPOSITORY_ROOT), "venv": ".venv"}, + "paths": {"hf_home": str(tmp_path / "hf")}, + "slurm": {"account": "test", "partition": "test"}, + }, + "resources": { + "selected": { + "mode": mode, + "gpus_per_node": 8, + "max_nodes": max_nodes, + } + }, + }, + ) + + +def _recipe(tmp_path: Path, **updates) -> Path: + payload = recipe_template(resource_profile="selected", run_root=str(tmp_path / "run")) + payload.update(updates) + return _write_yaml(tmp_path / "recipe.yaml", payload) + + +def test_resolved_bundle_preserves_integrity_and_source_provenance(tmp_path): + resolved = resolve_recipe_run(_recipe(tmp_path), _site(tmp_path)) + + assert resolved.plan["execution_mode"] == "reusable_allocation" + assert resolved.plan["stages"][-1]["stage_id"] == "post.params-90.best" + assert resolved.provenance["route"]["value"] == "qwen3.5-0.8b/vlm-pruning/smoke" + prerun = resolved.runner["runner"]["execution_contract"]["prerun_commands"] + assert "_assert_worker_source" in prerun[-1] + + bundle = materialize_resolved_bundle(resolved, activate=True) + assert bundle_for_run_root(tmp_path / "run") == bundle + manifest = json.loads((bundle / "manifest.json").read_text()) + assert manifest["code"]["controller"]["revision"] == "a" * 40 + assert manifest["code"]["worker"]["revision"] == "a" * 40 + + (bundle / "execution.yaml").write_text("execution: {}\n") + with pytest.raises(RuntimeError, match="changed after sealing"): + bundle_for_run_root(tmp_path / "run") + + +def test_bundle_identity_includes_authored_input_locations(tmp_path): + recipe = recipe_template(resource_profile="selected", run_root=str(tmp_path / "run")) + site = yaml.safe_load(_site(tmp_path).read_text()) + + first = resolve_recipe_run( + _write_yaml(tmp_path / "first.recipe.yaml", recipe), + _write_yaml(tmp_path / "first.site.yaml", site), + ) + second = resolve_recipe_run( + _write_yaml(tmp_path / "second.recipe.yaml", recipe), + _write_yaml(tmp_path / "second.site.yaml", site), + ) + + assert first.experiment == second.experiment + assert first.bundle_id != second.bundle_id + + +@pytest.mark.parametrize( + ("update", "message"), + [ + ({"resource_profiel": "selected"}, "did you mean"), + ( + {"advanced": {"experiment": {"model.force_hf": False}}}, + "Duplicate no-op advanced value", + ), + ({"data": {"path": "/prepared/data"}}, "requires explicit immutable data values"), + ( + {"advanced": {"experiment": {"depth_importance.eval_samples": 3}}}, + "targets inactive stage", + ), + ], +) +def test_closed_recipe_schema_rejects_unknown_duplicate_or_unused_values(tmp_path, update, message): + payload = recipe_template(resource_profile="selected", run_root=str(tmp_path / "run")) + payload.update(update) + with pytest.raises(ValueError, match=message): + resolve_recipe_run(_write_yaml(tmp_path / "recipe.yaml", payload), _site(tmp_path)) + + +@pytest.mark.parametrize( + ("advanced", "message"), + [ + ({"experiment": {"puzzle_dir": "/tmp/other"}}, "owned by the selected route"), + ({"experiment": {"model": {"force_hf": False}}}, "owned by the selected route"), + ( + {"execution": {"defaults": {"gpus_per_node": 4}}}, + "cannot override site-owned fields", + ), + ], +) +def test_recipe_cannot_override_route_or_site_owned_identity(tmp_path, advanced, message): + recipe = recipe_template(resource_profile="selected", run_root=str(tmp_path / "run")) + recipe["advanced"] = advanced + with pytest.raises(ValueError, match=message): + resolve_recipe_run(_write_yaml(tmp_path / "recipe.yaml", recipe), _site(tmp_path)) + + +def test_duplicate_yaml_keys_report_the_source_line(tmp_path): + recipe = tmp_path / "recipe.yaml" + recipe.write_text( + "schema_version: 1\nname: first\nname: second\nmodel: qwen3.5-0.8b\n" + "workflow: vlm-pruning\nmode: smoke\nrun_root: run\nresource_profile: selected\n" + ) + with pytest.raises(ValueError, match="Duplicate YAML key 'name' at line 3"): + resolve_recipe_run(recipe, _site(tmp_path)) + + +@pytest.mark.parametrize( + ("update", "message"), + [ + ( + lambda site: site["site"]["environment"].update({"setup_env": 1}), + "site.environment.setup_env", + ), + ( + lambda site: site["site"]["environment"].update( + {"prerun_commands": ["export HF_HOME=/other/cache"]} + ), + "assigns HF_HOME", + ), + ( + lambda site: site["site"]["environment"].update({"container_mounts": "/data:/data"}), + "container_mounts is unused", + ), + ( + lambda site: site["site"]["environment"].update({"source_revision": "main"}), + "source_revision must be a full immutable Git commit", + ), + ], +) +def test_closed_site_schema_rejects_invalid_or_ignored_values(tmp_path, update, message): + site = yaml.safe_load(_site(tmp_path).read_text()) + update(site) + with pytest.raises((TypeError, ValueError), match=message): + resolve_recipe_run(_recipe(tmp_path), _write_yaml(tmp_path / "site.yaml", site)) + + +def test_baremetal_site_uses_the_same_recipe_contract(tmp_path): + site = { + "schema_version": 1, + "site": { + "kind": "baremetal", + "environment": {"repository": str(REPOSITORY_ROOT), "venv": ".venv"}, + "paths": {"hf_home": str(tmp_path / "hf")}, + "baremetal": {"hosts": [{"hostname": "worker-a", "gpus": 8}]}, + }, + "resources": {"selected": {"mode": "per_attempt", "gpus_per_node": 8, "max_nodes": 1}}, + } + resolved = resolve_recipe_run(_recipe(tmp_path), _write_yaml(tmp_path / "site.yaml", site)) + assert resolved.plan["runner_kind"] == "baremetal" + assert resolved.runner["runner"]["inventory"]["hosts"] == [{"hostname": "worker-a", "gpus": 8}] + + +def test_environment_defaults_do_not_change_recipe_identity(tmp_path, monkeypatch): + recipe = _recipe(tmp_path) + site = _site(tmp_path) + monkeypatch.setenv("PUZZLETRON_DATASET_REVISION", "first-hidden-value") + first = resolve_recipe_run(recipe, site) + monkeypatch.setenv("PUZZLETRON_DATASET_REVISION", "second-hidden-value") + second = resolve_recipe_run(recipe, site) + assert first.experiment == second.experiment + assert first.bundle_id == second.bundle_id + + +def test_worker_revision_and_dirty_state_are_enforced(tmp_path, monkeypatch): + site = yaml.safe_load(_site(tmp_path).read_text()) + site["site"]["environment"]["repository"] = str(tmp_path / "worker") + site_path = _write_yaml(tmp_path / "site.yaml", site) + monkeypatch.setattr( + recipe_config, + "_repository_revision", + lambda _path: {"revision": "b" * 40, "dirty": False}, + ) + resolved = resolve_recipe_run(_recipe(tmp_path), site_path) + assert resolved.code["worker"]["revision"] == "b" * 40 + assert resolved.experiment["vlm_smoke_evaluation"]["evaluator_revision"] == "b" * 40 + + monkeypatch.setattr( + recipe_config, + "_repository_revision", + lambda _path: { + "revision": "b" * 40, + "dirty": True, + "working_tree_sha256": "c" * 64, + }, + ) + with pytest.raises(RuntimeError, match="source state changed"): + recipe_config._assert_worker_source( + str(tmp_path / "worker"), {"revision": "b" * 40, "dirty": False} + ) + + packaged_repository = tmp_path / "image" / "src" / "modelopt" + packaged_repository.mkdir(parents=True) + (packaged_repository.parents[1] / "modelopt_revision").write_text("d" * 40 + "\n") + assert recipe_config._source_identity.repository_revision(packaged_repository) == { + "revision": "d" * 40, + "dirty": False, + } + + +def test_source_identity_ignores_lfs_materialization_but_detects_code_edits(tmp_path): + repository = tmp_path / "worker" + repository.mkdir() + (repository / ".gitattributes").write_text("report.html filter=lfs diff=lfs -text\n") + (repository / "report.html").write_text( + f"version https://git-lfs.github.com/spec/v1\noid sha256:{'a' * 64}\nsize 12\n" + ) + source = repository / "worker.py" + source.write_text("VALUE = 1\n") + subprocess.run(["git", "init", "-q"], cwd=repository, check=True) + subprocess.run(["git", "add", "."], cwd=repository, check=True) + subprocess.run( + [ + "git", + "-c", + "user.name=PuzzleTron test", + "-c", + "user.email=puzzletron@example.com", + "-c", + "commit.gpgsign=false", + "commit", + "-q", + "-m", + "baseline", + ], + cwd=repository, + check=True, + ) + (repository / "report.html").write_text("materialized report\n") + assert recipe_config._repository_revision(repository)["dirty"] is False + source.write_text("VALUE = 2\n") + detected = recipe_config._repository_revision(repository) + assert detected["dirty"] is True + assert detected["working_tree_sha256"] + + +def test_synthetic_multinode_topology_is_validated_against_site_capacity(tmp_path): + payload = recipe_template( + model="qwen3.5-4b", + workflow="vlm-pruning", + mode="campaign", + resource_profile="selected", + run_root=str(tmp_path / "large-run"), + ) + payload["data"] = {"path": str(tmp_path / "data"), "revision": "fixture-revision"} + payload["advanced"] = { + "execution": { + "stages": { + "post.candidate-evaluation.screening_kd": { + "instances": 2, + "parallel": { + "tp": 8, + "pp": 4, + "cp": 1, + "dp_shard": 2, + "dp_replicate": 2, + "ep": 2, + }, + } + } + } + } + recipe = _write_yaml(tmp_path / "large.yaml", payload) + resolved = resolve_recipe_run(recipe, _site(tmp_path, max_nodes=64, mode="per_attempt")) + kd = next( + stage + for stage in resolved.plan["stages"] + if stage["stage_id"] == "post.candidate-evaluation.screening_kd" + ) + assert (kd["total_gpus"], kd["nodes"]) == (256, 32) + with pytest.raises(ValueError, match=r"provides at most 31 node.*screening_kd"): + resolve_recipe_run(recipe, _site(tmp_path, max_nodes=31, mode="per_attempt")) + + +def test_every_checked_in_recipe_resolves_to_one_catalog_route(tmp_path): + recipe_paths = sorted((REPOSITORY_ROOT / "examples/puzzletron/configs/recipes").glob("*.yaml")) + resolved_routes = [] + for index, source in enumerate(recipe_paths): + payload = yaml.safe_load(source.read_text()) + payload["run_root"] = str(tmp_path / f"run-{index}") + payload["resource_profile"] = "selected" + if "data" in payload: + payload["data"] = { + "path": str(tmp_path / f"dataset-{index}"), + "revision": f"fixture-revision-{index}", + } + resolved = resolve_recipe_run( + _write_yaml(tmp_path / f"recipe-{index}.yaml", payload), + _site(tmp_path, max_nodes=64, mode="per_attempt"), + ) + resolved_routes.append(resolved.route.route_id) + assert Counter(resolved_routes) == Counter(route.route_id for route in recipe_config.ROUTES) + + +def test_cli_validate_explain_dry_run_and_inspect_share_one_contract(tmp_path, monkeypatch, capsys): + recipe = _recipe(tmp_path) + site = _site(tmp_path) + assert public_cli.main(["validate", str(recipe), "--site", str(site)]) == 0 + assert "valid: my-puzzletron-run" in capsys.readouterr().out + assert public_cli.main(["explain", str(recipe), "--site", str(site)]) == 0 + assert "route: qwen3.5-0.8b/vlm-pruning/smoke" in capsys.readouterr().out + assert public_cli.main(["dry-run", str(recipe), "--site", str(site), "--color", "never"]) == 0 + assert "dry-run only; no jobs will be submitted" in capsys.readouterr().err + + resolved = resolve_recipe_run(recipe, site) + materialize_resolved_bundle(resolved, activate=True) + assert public_cli.main(["inspect", str(tmp_path / "run")]) == 0 + assert f"bundle: {resolved.bundle_id}" in capsys.readouterr().out + + +def test_launch_and_resume_delegate_the_same_sealed_inputs(tmp_path, monkeypatch): + calls = [] + monkeypatch.setattr(public_cli.orchestrate, "main", lambda argv: calls.append(argv) or 0) + recipe = _recipe(tmp_path) + site = _site(tmp_path) + assert public_cli.main(["launch", str(recipe), "--site", str(site)]) == 0 + assert public_cli.main(["resume", str(tmp_path / "run")]) == 0 + + def inputs(argv): + return tuple( + argv[argv.index(flag) + 1] for flag in ("--experiment", "--runner", "--execution") + ) + + bundle = bundle_for_run_root(tmp_path / "run") + assert ( + inputs(calls[0]) + == inputs(calls[1]) + == ( + str(bundle / "experiment.runtime.yaml"), + str(bundle / "runner.yaml"), + str(bundle / "execution.yaml"), + ) + ) + assert all(call[call.index("--stage") + 1] == "full" for call in calls) + + +def test_concurrent_activation_binds_exactly_one_bundle(tmp_path, monkeypatch): + site = _site(tmp_path) + first = resolve_recipe_run(_recipe(tmp_path), site) + second = resolve_recipe_run( + _recipe(tmp_path, advanced={"experiment": {"pruning.eval_samples": 3}}), site + ) + materialize_resolved_bundle(first, activate=False) + materialize_resolved_bundle(second, activate=False) + barrier = Barrier(2) + link = os.link + + def synchronized_link(source, destination): + barrier.wait(timeout=5) + link(source, destination) + + monkeypatch.setattr(recipe_config.os, "link", synchronized_link) + + def activate(resolved): + try: + materialize_resolved_bundle(resolved, activate=True) + except ValueError: + return "conflict", resolved.bundle_id + return "bound", resolved.bundle_id + + with ThreadPoolExecutor(max_workers=2) as pool: + outcomes = list(pool.map(activate, (first, second))) + assert Counter(status for status, _ in outcomes) == Counter({"bound": 1, "conflict": 1}) + bound_id = next(bundle_id for status, bundle_id in outcomes if status == "bound") + active = json.loads((tmp_path / "run/orchestration/current_bundle.json").read_text()) + assert active["bundle_id"] == bound_id diff --git a/tests/unit/torch/puzzletron/test_setup_entrypoints.py b/tests/unit/torch/puzzletron/test_setup_entrypoints.py new file mode 100644 index 00000000000..d99b90272a3 --- /dev/null +++ b/tests/unit/torch/puzzletron/test_setup_entrypoints.py @@ -0,0 +1,65 @@ +# 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. + +"""Compatibility checks for setup entry points.""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +import pytest + +from examples.puzzletron import puzzletron as public_cli +from puzzletron_setup.v2 import cli as setup_cli + +REPOSITORY_ROOT = Path(__file__).resolve().parents[4] + + +def test_public_cli_uses_the_existing_setup_wizard(monkeypatch): + calls = [] + monkeypatch.setattr(setup_cli, "main", lambda argv: calls.append(argv) or 12) + + assert public_cli.main(["setup", "--resume", "campaign"]) == 12 + assert calls == [["--resume", "campaign"]] + + +def test_public_help_advertises_one_setup_command(capsys): + with pytest.raises(SystemExit) as error: + public_cli.main(["--help"]) + assert error.value.code == 0 + + output = capsys.readouterr().out + assert "setup" in output + assert "setup custom" not in output + assert "init" not in output + assert "site init" not in output + + +def test_existing_setup_scripts_keep_their_wizard_help(): + for script_name, expected_option in ( + ("puzzletron_setup.py", "--detailed"), + ("puzzletron_setup_v2.py", "--campaign-dir"), + ): + completed = subprocess.run( + [sys.executable, str(REPOSITORY_ROOT / "examples/puzzletron" / script_name), "--help"], + check=False, + capture_output=True, + text=True, + ) + + assert completed.returncode == 0, completed.stderr + assert expected_option in completed.stdout diff --git a/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py b/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py index 8cf142fc5c9..bc5c2f52ba2 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py @@ -530,7 +530,8 @@ def test_generated_readme_separates_plan_inspection_from_launch(tmp_path: Path) resume_command = next(command for command in commands if "--resume" in command) assert resume_command == [ "python", - f"{repository}/examples/puzzletron/puzzletron_setup_v2.py", + f"{repository}/examples/puzzletron/puzzletron.py", + "setup", "--resume", str(campaign_dir), ]