diff --git a/.agents/skills/puzzletron/README.md b/.agents/skills/puzzletron/README.md new file mode 100644 index 00000000000..0820f44aee7 --- /dev/null +++ b/.agents/skills/puzzletron/README.md @@ -0,0 +1,131 @@ +# Puzzletron Agent Skill + +Puzzletron is an end-to-end workflow for model pruning and MIP-based architecture optimization. +This skill exposes it as a slash command for AI coding agents and via natural language conversation. + +For full environment setup, model configuration, and algorithm details see +[examples/puzzletron/README.md](../../examples/puzzletron/README.md). + +> **Experimental:** AI agent integration is an experimental feature and may change. + +Run `/puzzletron` with no arguments to see available commands. + +## Running the full pipeline + +To run the full 8-step pipeline, use the slash command (where the number is GPUs per node): + +```text +/puzzletron all 2 +``` + +Or in natural language: + +```text +run puzzletron all for Llama-3.1-8B on 2 GPUs +``` + +Check progress with: + +```text +/puzzletron all progress +``` + +Example output while running: + +```text +Overall: Puzzletron full pipeline (steps 1–8) +──────────────────────────────────────────────────────────────────── + Status Step Description Elapsed +──────────────────────────────────────────────────────────────────── + [DONE] 1/8: starting puzzletron pipeline 0m 0s + [DONE] 2/8: converting model to Puzzletron heterogeneous format (single-gpu) 0m 26s + [DONE] 3/8: scoring pruning activations (multi-gpu) 9m 9s + [DONE] 4/8: pruning the model and saving pruned checkpoints (single-gpu) 0m 57s + [DONE] 5/8: building replacement library and subblock statistics (single-gpu) 0m 26s + [RUNNING] 6/8: calculating one block scores (multi-gpu) (270/352 solutions) 100m 6s + [ ] 7/8: running MIP and realizing models (multi-gpu) + [ ] 8/8: puzzletron pipeline completed (multi-gpu) +──────────────────────────────────────────────────────────────────── + Started: 00:08:50 + Finished: 01:59:54 (in progress) + Elapsed: 111m 4s + Completed: 5/8 steps + Remaining: 56m 24s estimated +``` + +Step 6 progress is tracked via completed `solution_N.json` files on disk for an accurate +remaining estimate. Step 7 (MIP sweep) shows per-rate progress once it starts. + +## Running the MIP step + +Start the MIP step by telling the agent how many GPUs per node to use: + +```text +/puzzletron mip 4 +``` + +Output is streamed live and also written to `./log.txt`. While it runs (or after it finishes), +check progress with: + +```text +/puzzletron mip progress +``` + +Example output when complete: + +```text +Overall: Puzzletron step 7/8 — MIP sweep (6 compression rates) +────────────────────────────────────────────────────────────── + Status Phase Elapsed +────────────────────────────────────────────────────────────── + [DONE] Prep (teacher memory + rate list) <1s + [DONE] compression_rate=0.5 3m 52s + [DONE] compression_rate=0.6 4m 41s + [DONE] compression_rate=0.7 4m 46s + [DONE] compression_rate=0.8 3m 55s + [DONE] compression_rate=0.9 3m 55s + [DONE] compression_rate=1.0 3m 59s +────────────────────────────────────────────────────────────── + Started: 08:05:30 + Finished: 08:30:38 + Elapsed: 25m 8s + Completed: 6/6 compression rates + Remaining: done estimated + + Results: /workspace/puzzle_dir/mip_sweep_results.csv +``` + +While running, the report shows which rate is active, sub-step detail (MIP solver node count +or validation batch progress), and an estimated time remaining based on completed rates. + +## Checking compressed model accuracy + +Two commands are available depending on whether you ran a single constrained MIP solve or a sweep: + +**Single constrained run** — teacher vs. solution_0 at the configured target memory: + +```text +/puzzletron mip losses +``` + +**Sweep** — accuracy across all compression rates from the sweep CSV: + +```text +/puzzletron mip sweep losses +``` + +Example `mip losses` output for Qwen3.5-0.8B (target 10,000 MiB): + +| Metric | Teacher | Compressed (solution_0) | +|---|---|---| +| `target_memory` | 20,389 MiB | 10,000 MiB | +| `lm_loss` | 1.1067 | 3.8808 | +| `token_accuracy_top_1` | 0.7365 | 0.2915 | +| `token_accuracy_top_5` | 0.9079 | 0.5500 | +| `token_accuracy_top_10` | 0.9399 | 0.6451 | + +## Adding support for a new model + +See [adding_new_model_tutorial.md](adding_new_model_tutorial.md) for a step-by-step walkthrough +covering: diagnosing why a model isn't supported, upgrading Transformers, writing a model +descriptor and converter, creating YAML configs, and a final checklist. diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md new file mode 100644 index 00000000000..a18398d8fd7 --- /dev/null +++ b/.agents/skills/puzzletron/SKILL.md @@ -0,0 +1,246 @@ +--- +name: puzzletron +description: "End-to-end workflow for model pruning and MIP-based optimization. Commands: mip, all, add-model. Usage: /puzzletron [args]" +license: Apache-2.0 +--- + +# Puzzletron + +## Routing + +**STEP 1 — Check args before doing anything else. This is MANDATORY.** + +- If args are **empty**, output the block below verbatim and **STOP immediately. Do NOT proceed to any command.** +- If the first word of args does **not exactly match** `mip`, `all`, or `add-model`, output the block below verbatim and **STOP immediately. Do NOT proceed to any command.** + +--- + +**Puzzletron** — end-to-end workflow for model pruning and MIP-based optimization. + +Available commands: +- `mip ` — Run the MIP step (nproc_per_node: number of GPUs per node) +- `mip progress` — Show live MIP progress with timing summary +- `mip losses` — Show teacher vs. compressed model accuracy for the single constrained MIP solution +- `mip sweep losses` — Show accuracy across all compression rates from a completed sweep +- `all ` — Run the full Puzzletron pipeline (nproc_per_node: number of GPUs per node) +- `all progress` — Show live full pipeline progress with timing summary +- `add-model ` — Implement descriptor, converter, and configs for an unsupported model + +Usage: `/puzzletron [args]` + +--- + +**STEP 2 — Only if the first word of args exactly matches a command name, execute it. Never reach this step if args were empty.** + +## Command: all + +Parse `nproc_per_node` from args using either positional or flag syntax: +- Positional: second word is a number, e.g. `all 2` +- Flag: `--nproc_per_node ` anywhere in args, e.g. `all --nproc_per_node 2` + +- If the second word is exactly `progress`, execute the **all progress** sub-command below. +- If no `nproc_per_node` value can be found, ask the user: "Please provide the number of GPUs per node (nproc_per_node)." and **STOP**. +- If the value does not match `^[0-9]+$`, ask the user: "nproc_per_node must be a positive integer." and **STOP**. +- Otherwise use the parsed value and run the full pipeline. + +### all \ + +Run the following Bash command, substituting `` with the parsed value: + +```bash +set -o pipefail && export PYTHONPATH=$PYTHONPATH:/workspace/Model-Optimizer && \ +torchrun --nproc_per_node examples/puzzletron/main.py \ + --config examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml \ + 2>&1 | tee ./log.txt | grep "Puzzletron Progress" +``` + +Stream output to the user as it arrives. When the command finishes, report the exit code. + +### all progress + +Run the following Bash command. Present the output to the user wrapped in a fenced code block (``` ... ```). + +```bash +python3 .agents/skills/puzzletron/all_progress.py +``` + +## Command: mip + +Parse `nproc_per_node` from args using either positional or flag syntax: +- Positional: second word is a number, e.g. `mip 2` +- Flag: `--nproc_per_node ` anywhere in args, e.g. `mip --nproc_per_node 2` + +- If the second word is exactly `progress`, execute the **mip progress** sub-command below. +- If the second word is exactly `losses`, execute the **mip losses** sub-command below. +- If the second and third words are exactly `sweep losses`, execute the **mip sweep losses** sub-command below. +- If no `nproc_per_node` value can be found, ask the user: "Please provide the number of GPUs per node (nproc_per_node)." and **STOP**. +- If the value does not match `^[0-9]+$`, ask the user: "nproc_per_node must be a positive integer." and **STOP**. +- Otherwise use the parsed value and run the MIP step. + +### mip \ + +Run the following Bash command, substituting `` with the parsed value: + +```bash +set -o pipefail && export PYTHONPATH=$PYTHONPATH:/workspace/Model-Optimizer && \ +torchrun --nproc_per_node examples/puzzletron/main.py \ + --config examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml \ + --mip-only 2>&1 | tee ./log.txt | grep "Puzzletron Progress" +``` + +Stream output to the user as it arrives. When the command finishes, report the exit code. + +### mip progress + +Run the following Bash command. Present the output to the user wrapped in a fenced code block (``` ... ```). + +```bash +python3 .agents/skills/puzzletron/mip_progress.py +``` + +### mip losses + +Run the following Bash command. Present the output to the user wrapped in a fenced code block (``` ... ```). + +```bash +python3 .agents/skills/puzzletron/mip_losses.py +``` + +### mip sweep losses + +Run the following Bash command. Present the output to the user wrapped in a fenced code block (``` ... ```). + +```bash +python3 .agents/skills/puzzletron/mip_sweep.py +``` + +## Command: add-model + +Parse `hf_model_path` from args (the second word). If missing, ask: "Please provide the HuggingFace model path (local or hub)." and **STOP**. + +Then follow the steps below to implement full Puzzletron support for the model. + +### Step 1 — Check if already supported + +```bash +python3 -c " +import sys; sys.path.insert(0, '.') +from modelopt.torch.puzzletron.anymodel.model_descriptor import ModelDescriptorFactory +from transformers import AutoConfig +cfg = AutoConfig.from_pretrained('', trust_remote_code=True) +supported = cfg.model_type in ModelDescriptorFactory.CLASS_MAPPING +print(f'model_type: {cfg.model_type}') +print(f'already supported: {supported}') +" +``` + +If already supported, tell the user and **STOP**. + +If `AutoConfig` raises an error about an unrecognised model type, the installed Transformers version is too old. Check the version, upgrade with `python3 -m pip install --upgrade transformers`, then re-run. + +### Step 2 — Inspect the architecture + +Run the following to understand what you are implementing: + +```bash +python3 -c " +from transformers import AutoConfig +cfg = AutoConfig.from_pretrained('', trust_remote_code=True) +print(cfg) +# If it has a nested text_config, print that too +if hasattr(cfg, 'text_config'): + print('--- text_config ---') + print(cfg.text_config) +" +``` + +Key things to note: +- **`model_type`** — this becomes the registration key for both descriptor and converter. +- **Nested `text_config`** — VLMs (e.g. Qwen3.5) wrap language model params inside `config.text_config`. Use `config.text_config` wherever you need `num_hidden_layers`, `intermediate_size`, `num_key_value_heads`. The converter must save `text_config` (not the full VLM config) so downstream code can access these fields directly. +- **Hybrid attention** — check `cfg.text_config.layer_type_list` (or similar). If some layers are linear/recurrent and others are full attention, `attn_no_op_post_init` must branch on `decoder_layer.layer_type`. +- **MoE** — if `num_experts` > 1 the model uses a MoE FFN and is not currently supported by the FFN pruning path; skip FFN pruning for such models. +- **Weight name prefixes** — inspect the checkpoint index to understand the layout: + +```bash +python3 -c " +import json, collections +idx = json.load(open('/model.safetensors.index.json')) +prefixes = collections.Counter() +for n in idx['weight_map']: + prefixes['.'.join(n.split('.')[:3])] += 1 +for p, c in sorted(prefixes.items()): + print(f'{c:4d} {p}') +" +``` + +If weight names use a prefix like `model.language_model.*` rather than `model.*`, the converter must implement `convert_weight_name` to remap them, and `get_weight_groups` must handle both the original checkpoint names (used during conversion) and the remapped names (used when saving pruned checkpoints). See the Qwen3_5 descriptor/converter for the reference implementation of this pattern. + +### Step 3 — Create the files + +Create the following files (use an existing descriptor as a reference — `qwen3_5` for VLMs with nested config and weight remapping, `llama` or `qwen2` for standard text-only models): + +**`modelopt/torch/puzzletron/anymodel/models//__init__.py`** + +```python +from ._converter import * +from ._model_descriptor import * +``` + +**`modelopt/torch/puzzletron/anymodel/models//_model_descriptor.py`** + +Must implement (inheriting from `ModelDescriptor`): +- `decoder_layer_cls()` → the HF decoder layer class +- `input_embedding_name()` → e.g. `"model.embed_tokens"` +- `output_embedding_name()` → e.g. `"lm_head"` +- `final_norm_name()` → e.g. `"model.norm"` +- `layer_block_name(index)` → e.g. `f"model.layers.{index}"` +- `block_config_to_layer_overrides(block_config)` → dict with `intermediate_size` and `num_key_value_heads` +- `attn_no_op_post_init(decoder_layer)` → replace attention + input norm with no-ops +- `mlp_no_op_post_init(decoder_layer)` → replace MLP + post-attention norm with no-ops +- `layer_name_predicates(num_layers)` → regex dict grouping weights into `embeddings`, `lm_head`, `block_N_ffn`, `block_N_attention` +- `init_rotary_embedding(model, runtime)` → re-initialise rotary embedding after subblock load + +**Critical:** `layer_name_predicates` patterns must match the **converted** `model.*` names (not the original VLM checkpoint names). If the checkpoint uses a different prefix, override `get_weight_groups` to normalise names before matching and restore originals in the returned groups (so `param_to_file` lookups in `convert_model_weights` still work). See `Qwen3_5ModelDescriptor.get_weight_groups` for the reference pattern. + +**`modelopt/torch/puzzletron/anymodel/models//_converter.py`** + +Must implement (inheriting from `Converter`): +- `create_block_configs_from_main_config(config)` → list of `BlockConfig`, one per layer +- `convert_configs_in_dirs(input_dir, output_dir)` → if the model has a nested `text_config`, save that instead of the full VLM config so `num_hidden_layers` is accessible at the top level +- `convert_weight_name(name)` → remap checkpoint weight names to converted model names (identity if no remapping needed) + +**Register in `modelopt/torch/puzzletron/anymodel/models/__init__.py`** — gate behind the minimum Transformers version that introduced the model: + +```python +if _Version(_transformers_version) >= _Version("X.Y.Z"): + from . import * +``` + +**Compression config** at `examples/puzzletron/configs/-_pruneffn_memory/`: +- Base YAML (`.yaml`): `descriptor: `, MIP constraints +- Main YAML (override): `input_hf_model_path`, `dataset_path`, `puzzle_dir` (use a **model-specific path** to avoid collisions with other models), `pruning.intermediate_size_list` +- Pruning YAML: points `layer_descriptor._target_` at the new `FFNIntermediateLayerDescriptor` subclass + +Choose `intermediate_size_list` by scaling the Llama-3.1-8B ratios (~21%, 42%, 60%, 83% of teacher) to the new model's `intermediate_size`. + +### Step 4 — Verify registration + +```bash +python3 -c " +import sys; sys.path.insert(0, '.') +from modelopt.torch.puzzletron.anymodel.model_descriptor import ModelDescriptorFactory +from modelopt.torch.puzzletron.anymodel.converter import ConverterFactory +print('descriptor:', '' in ModelDescriptorFactory.CLASS_MAPPING) +print('converter: ', '' in ConverterFactory.CLASS_MAPPING) +" +``` + +Both must print `True`. If not, check the `__init__.py` import chain and the `@register_decorator` keys. + +### Step 5 — Tell the user what was created + +List the files created, confirm registration, and suggest running the pipeline: + +```text +run puzzletron all for on GPUs +``` diff --git a/.agents/skills/puzzletron/adding_new_model_tutorial.md b/.agents/skills/puzzletron/adding_new_model_tutorial.md new file mode 100644 index 00000000000..275380297a0 --- /dev/null +++ b/.agents/skills/puzzletron/adding_new_model_tutorial.md @@ -0,0 +1,209 @@ +# Adding a New Model to Puzzletron with Claude + +This tutorial shows the full conversation between a user and Claude when adding +Qwen3.5-0.8B support to Puzzletron. Use it as a guide for how to interact with Claude +when you want to compress a model that Puzzletron doesn't yet support. + +--- + +## The journey + +### Step 1: Download the model and try to run it + +The user starts simply — download the model and run the pipeline. + +> **User:** can you download to /workspace/hf_models/... + +Claude downloads the model using `huggingface-cli`: + +```text +Successfully downloaded 13 files to /workspace/hf_models/Qwen3.5-0.8B +``` + +> **User:** run puzzletron all for this model + +Claude checks whether the model is supported — first by trying to load it with Transformers, then by looking for an existing Puzzletron descriptor and compression config. It finds two blockers: + +1. Transformers doesn't recognise model type `qwen3_5`: + +```text +ValueError: The checkpoint you are trying to load has model type `qwen3_5` +but Transformers does not recognize this architecture. +``` + +2. There is no Puzzletron descriptor for `qwen3_5` in `modelopt/torch/puzzletron/anymodel/models/` and no compression example config under `examples/puzzletron/configs/` + +--- + +### Step 2: Fix the Transformers version + +> **User:** so please install it using python -m pip + +Claude upgrades: + +```bash +python3 -m pip install --upgrade transformers +# → Successfully installed transformers-5.12.1 +``` + +Claude confirms the model now loads: + +```python +from transformers import AutoConfig +config = AutoConfig.from_pretrained("/workspace/hf_models/Qwen3.5-0.8B") +print(config.model_type) # qwen3_5 ✓ +``` + +--- + +### Step 3: Implement the descriptor, converter, and configs + +> **User:** ok, implement the descriptor and compression example config for Qwen3.5-0.8B + +Claude implements everything in one go: + +**Model descriptor** (`modelopt/torch/puzzletron/anymodel/models/qwen3_5/qwen3_5_model_descriptor.py`) + +**Converter** + +**Registration** (`anymodel/models/__init__.py`): gated behind `transformers >= 4.57.0` + +**Compression example config** (`examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/`): + +`intermediate_size_list` chosen by scaling the Llama-3.1-8B ratios (~21 %, 42 %, 60 %, 83 % +of teacher) to the 0.8B's `intermediate_size = 3584`: + +```text +[1280, 2560, 3584, 5120] +``` + +MIP constraints set proportionally to Llama 3.1-8B: + +```yaml +target_memory: 20_000 # 20 GiB +num_params: 1_500_000_000 +``` + +Claude then verifies that both the descriptor and converter register correctly: + +```python +from modelopt.torch.puzzletron.anymodel.model_descriptor import ModelDescriptorFactory +print('qwen3_5' in ModelDescriptorFactory.CLASS_MAPPING) # True ✓ +``` + +--- + +### Step 4: Run the pipeline + +> **User:** run puzzletron all for Qwen3.5-0.8B on 4 GPUs + +Claude constructs the `torchrun` command directly with the Qwen3.5-0.8B config path and runs the full pipeline. The user monitors progress with: + +```text +/puzzletron all progress +``` + +Example output mid-run: + +```text +Overall: Puzzletron full pipeline (steps 1–8) +──────────────────────────────────────────────────────────────────── + Status Step Description Elapsed +──────────────────────────────────────────────────────────────────── + [DONE] 1/8: starting puzzletron pipeline 0m 1s + [DONE] 2/8: converting model to Puzzletron heterogeneous format (single-gpu) 0m 3s + [DONE] 3/8: scoring pruning activations (multi-gpu) 0m 56s + [DONE] 4/8: pruning the model and saving pruned checkpoints (single-gpu) 0m 10s + [DONE] 5/8: building replacement library and subblock statistics (single-gpu) 0m 10s + [RUNNING] 6/8: calculating one block scores (multi-gpu) (127/264 solutions) 20m 6s + [ ] 7/8: running MIP and realizing models (multi-gpu) + [ ] 8/8: puzzletron pipeline completed (multi-gpu) +──────────────────────────────────────────────────────────────────── + Started: 10:35:54 + Elapsed: 21m 26s | Remaining: ~28m estimated +``` + +Step 6 (one-block scoring) is the longest step — it scores all candidate solutions using a proxy metric (cosine embedding loss on hidden states). The number of solutions depends on the model size and `intermediate_size_list`; for Qwen3.5-0.8B with 4 sizes across 28 layers it is 264. Use `eval_samples` in the base YAML to trade off speed vs. score quality (default 128; 8 is useful for quick iteration). + +--- + +### Step 5: Check the compressed model accuracy + +> **User:** show mip losses for Qwen3.5-0.8B + +Claude runs `/puzzletron mip losses` and presents the results as a comparison table: + +Example output for Qwen3.5-0.8B: + +```text +Metric Teacher Compressed (solution_0) +-------------------------------------------------------------------- +target_memory 20,389 MiB 10,000 MiB +-------------------------------------------------------------------- +lm_loss 1.1067 3.8808 +token_accuracy_top_1 0.7365 0.2915 +token_accuracy_top_5 0.9079 0.55 +token_accuracy_top_10 0.9399 0.6451 + +Results from: /workspace/puzzle_dir_qwen3_5-0.8b/mip/puzzle_solutions/target_memory_10000MiB-num_params_1_5G/solutions--validation +Sweep results: use /puzzletron mip sweep losses +``` + +The results are read from `/mip/puzzle_solutions//solutions--validation/solution_0.json` (and `teacher.json` in the same directory). The teacher memory is taken from the sweep CSV if a sweep was also run. + +--- + +### Step 6: Run the MIP sweep and check sweep losses + +If the sweep is enabled in the config YAML (`mip.sweep.enabled: true`), run it after the full pipeline: + +> **User:** run sweep for Qwen3.5-0.8B + +Claude runs the MIP step with the Qwen3.5-0.8B config on the requested number of GPUs. Monitor progress with: + +```text +/puzzletron mip progress +``` + +Example output mid-run: + +```text +Overall: Puzzletron step 7/8 — MIP sweep (6 compression rates) +────────────────────────────────────────────────────────────── + Status Phase Elapsed +────────────────────────────────────────────────────────────── + [DONE] Prep (teacher memory + rate list) <1s + [DONE] compression_rate=0.5 0m 44s + [DONE] compression_rate=0.6 0m 36s + [DONE] compression_rate=0.7 0m 37s + [DONE] compression_rate=0.8 0m 37s + [RUNNING] compression_rate=0.9 — validating (8/8 batches) 0m 28s + [ ] compression_rate=1.0 pending +────────────────────────────────────────────────────────────── + Started: 00:03:28 + Finished: 00:06:30 (in progress) + Elapsed: 3m 2s + Completed: 4/6 compression rates + Remaining: 1m 17s estimated +``` + +Once complete, view accuracy across all compression rates: + +> **User:** show mip sweep losses + +Claude runs `/puzzletron mip sweep losses` and presents the results: + +```text + rate target_mem actual_mem num_params lm_loss top_1 top_5 top_10 +-------------------------------------------------------------------------------------- +0.5000 10194.3640 10143.2768 888,813,280 3.2367 0.3663 0.6384 0.7251 +0.6000 12233.2368 11719.5001 909,901,856 2.6377 0.4434 0.7198 0.7981 +0.7000 14272.1096 14083.8350 941,534,720 1.8532 0.5855 0.8176 0.8735 +0.8000 16310.9824 15660.0582 962,623,296 1.5385 0.6448 0.8576 0.9046 +0.9000 18349.8552 18024.3931 994,256,160 1.2447 0.7064 0.8914 0.9278 +1.0000 20388.7280 20388.7280 1,025,889,024 1.1067 0.7365 0.9079 0.9399 + +Results from: /workspace/puzzle_dir_qwen3_5-0.8b/mip_sweep_results.csv +``` + +Use this table to pick the compression rate that best meets your accuracy/memory budget. diff --git a/.agents/skills/puzzletron/all_progress.py b/.agents/skills/puzzletron/all_progress.py new file mode 100644 index 00000000000..5204a0b56e5 --- /dev/null +++ b/.agents/skills/puzzletron/all_progress.py @@ -0,0 +1,186 @@ +# 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. + +# Generated with Claude Code +"""Progress report for the full Puzzletron pipeline (all 8 steps).""" + +import glob +import re +import sys +from datetime import datetime + +LOG = "./log.txt" +try: + lines = open(LOG).readlines() + text = "".join(lines) +except FileNotFoundError: + print("No log.txt found. Run /puzzletron all first.") + sys.exit(0) + + +def fmt(s): + """Format seconds as 'Xm Ys', or '—' if None.""" + return f"{int(s) // 60}m {int(s) % 60}s" if s is not None else "—" + + +def get_ts(line): + """Extract a datetime from a log line timestamp, or None.""" + m = re.search(r"\[(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2})", line) + return datetime.strptime(m.group(1), "%Y-%m-%d %H:%M:%S") if m else None + + +now = datetime.now().replace(microsecond=0) +DIV = "─" * 68 + +step_events = [] +for line in lines: + m = re.search(r"Puzzletron Progress (\d+)/(\d+): (.+)", line) + if m: + step_num = int(m.group(1)) + total_steps = int(m.group(2)) + desc = m.group(3).strip() + ts = get_ts(line) + step_events.append((step_num, total_steps, desc, ts)) + +total_steps = step_events[-1][1] if step_events else 8 +seen_steps = {e[0]: (e[2], e[3]) for e in step_events} +last_step_num = max(seen_steps.keys()) if seen_steps else 0 + +pipeline_complete_ts = None +if last_step_num == total_steps and total_steps in seen_steps: + pipeline_complete_ts = seen_steps[total_steps][1] + +cur_detail = "" +step_remaining = None +batch_matches = re.findall(r"calculate_losses_pipeline[^:]*:\s*(\d+)%.*?(\d+)/(\d+)", text) +cbc_matches = re.findall(r"After (\d+) nodes.*?\(([\d.]+) seconds\)", text) + +sol_dir_match = re.search( + r"'output_dir': '([^']+single_sequence_replacement_solutions--validation[^']*)'", text +) +sol_done, sol_total = None, None +if sol_dir_match: + sol_dir = sol_dir_match.group(1) + sol_done = len(glob.glob(f"{sol_dir}/solution*.json")) + sol_list_match = re.search(r"'solutions_to_validate': \[([\d, ]+)\]", text) + if sol_list_match: + sol_total = len(sol_list_match.group(1).split(",")) +pct, cur_b, total_b = batch_matches[-1] if batch_matches else (None, None, None) +if sol_done is not None and sol_total: + cur_detail = f" ({sol_done}/{sol_total} solutions)" +elif batch_matches: + cur_detail = f" ({cur_b}/{total_b} batches)" +elif cbc_matches: + nodes, secs = cbc_matches[-1] + cur_detail = f" (MIP solver: {int(nodes):,} nodes, {float(secs):.1f}s)" + +pipeline_start = step_events[0][3] if step_events else None +end_ts = pipeline_complete_ts or now +total_elapsed = int((end_ts - pipeline_start).total_seconds()) if pipeline_start else 0 + +step_ts_list = sorted(seen_steps.items()) +cur_step_start_ts = seen_steps[last_step_num][1] if last_step_num in seen_steps else None +if not pipeline_complete_ts and cur_step_start_ts: + cur_step_elapsed = int((now - cur_step_start_ts).total_seconds()) + if sol_done and sol_total and sol_done > 0: + rate_per_sol = cur_step_elapsed / sol_done + step_remaining = rate_per_sol * (sol_total - sol_done) + elif cur_b is not None and total_b is not None and int(cur_b) > 0 and int(cur_b) < int(total_b): + rate_per_batch = cur_step_elapsed / int(cur_b) + step_remaining = rate_per_batch * (int(total_b) - int(cur_b)) + +print(f"\nOverall: Puzzletron full pipeline (steps 1–{total_steps})") # noqa: RUF001 +print(DIV) +print(f" {'Status':<10} {'Step':<4} {'Description':<34} {'Elapsed':>8}") +print(DIV) + +for i, (snum, (sdesc, sts)) in enumerate(step_ts_list): + next_ts = ( + step_ts_list[i + 1][1][1] if i + 1 < len(step_ts_list) else (pipeline_complete_ts or now) + ) + elapsed = int((next_ts - sts).total_seconds()) if sts and next_ts else None + is_last = snum == last_step_num + is_done = not is_last or pipeline_complete_ts is not None + detail = "" + if is_last and not is_done: + detail = cur_detail + label = f"{snum}/{total_steps}: {sdesc}{detail}" + status = "[DONE]" if is_done else "[RUNNING]" + print( + f" {status:<10} {'':<4} {label:<34} {fmt(elapsed) if elapsed is not None else '—':>8}" + ) + +_STEP_NAMES = { + 1: "starting puzzletron pipeline", + 2: "converting model to Puzzletron heterogeneous format (single-gpu)", + 3: "scoring pruning activations (multi-gpu)", + 4: "pruning the model and saving pruned checkpoints (single-gpu)", + 5: "building replacement library and subblock statistics (single-gpu)", + 6: "calculating one block scores (multi-gpu)", + 7: "running MIP and realizing models (multi-gpu)", + 8: "puzzletron pipeline completed (multi-gpu)", +} + +for snum in range(last_step_num + 1, total_steps + 1): + desc = _STEP_NAMES.get(snum, "pending") + print(f" {'[ ]':<10} {'':<4} {f'{snum}/{total_steps}: {desc}':<34} {'':>8}") + +print(DIV) +done_steps = len([s for s in seen_steps if s != last_step_num or pipeline_complete_ts]) +step_durations = [] +for i, (snum, (sdesc, sts)) in enumerate(step_ts_list): + next_ts = ( + step_ts_list[i + 1][1][1] if i + 1 < len(step_ts_list) else (pipeline_complete_ts or None) + ) + if next_ts and sts: + step_durations.append(int((next_ts - sts).total_seconds())) +avg_step_s = sum(step_durations) / len(step_durations) if step_durations else None + + +def step_est(snum): + """Estimate duration in seconds for a pending pipeline step.""" + if snum == 7: + # Step 7 in the full pipeline is a single MIP solve (~5m), not a sweep + return 296 + elif snum == 8: + return 60 + return avg_step_s or 0 + + +if pipeline_complete_ts: + est_rem = "done" +elif step_remaining is not None: + future_s = sum(step_est(s) for s in range(last_step_num + 1, total_steps + 1)) + est_rem = fmt(step_remaining + future_s) +else: + cur_s = step_est(last_step_num) + future_s = cur_s + sum(step_est(s) for s in range(last_step_num + 1, total_steps + 1)) + est_rem = fmt(future_s) if (cur_s or future_s) else "calculating..." + +finished_str = ( + pipeline_complete_ts.strftime("%H:%M:%S") + if pipeline_complete_ts + else now.strftime("%H:%M:%S") + " (in progress)" +) +print(f" Started: {pipeline_start.strftime('%H:%M:%S') if pipeline_start else '—'}") +print(f" Finished: {finished_str}") +print(f" Elapsed: {fmt(total_elapsed)}") +print(f" Completed: {done_steps}/{total_steps} steps") +print(f" Remaining: {est_rem} estimated") +results_match = re.search(r"Results written to: (\S+)", text) +if not results_match: + results_match = re.search(r"\[run_puzzle\.py:335\]\s+(\S+)", text) +if results_match: + print(f"\n Results: {results_match.group(1)}") diff --git a/.agents/skills/puzzletron/mip_losses.py b/.agents/skills/puzzletron/mip_losses.py new file mode 100644 index 00000000000..00a3e4fd953 --- /dev/null +++ b/.agents/skills/puzzletron/mip_losses.py @@ -0,0 +1,113 @@ +# 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. + +# Generated with Claude Code +"""Teacher vs. compressed model accuracy for the single constrained MIP solution. + +For sweep results across multiple compression rates use mip_sweep.py instead. +""" + +import csv +import glob +import json +import os +import re +import sys + +KEYS = ["lm_loss", "token_accuracy_top_1", "token_accuracy_top_5", "token_accuracy_top_10"] + +LOG = "./log.txt" +try: + text = open(LOG).read() +except FileNotFoundError: + print("No log.txt found. Run /puzzletron all first.") + sys.exit(1) + +# Extract puzzle_dir from any path containing /mip/puzzle_solutions/ in the log +match = re.search(r"(\S+)/mip/puzzle_solutions/", text) +if not match: + match = re.search(r"(\S+)/ckpts/teacher", text) +if not match: + print("Could not find puzzle_dir in log.txt. Has the pipeline run?") + sys.exit(1) + +puzzle_dir = match.group(1) + +# Extract the configured target_memory from the log (logged in the args dict) +target_mem_match = re.search(r"'target_memory':\s*([\d.]+)", text) +configured_target = float(target_mem_match.group(1)) if target_mem_match else None + +# Find all validation directories +solutions_dirs = sorted(glob.glob(f"{puzzle_dir}/mip/puzzle_solutions/*/solutions--validation")) +if not solutions_dirs: + print(f"No MIP validation results found under {puzzle_dir}. Has the pipeline completed?") + sys.exit(1) + +# Prefer the directory whose name matches the configured target_memory (the constrained run), +# not the sweep directories which have the teacher memory as target. +chosen_dir = None +if configured_target is not None: + target_str = str(int(configured_target)) + for d in solutions_dirs: + if f"target_memory_{target_str}" in d and "num_params" in d: + chosen_dir = d + break + +if chosen_dir is None: + # Fall back to the first (smallest target = most compressed) + chosen_dir = solutions_dirs[0] + +solutions_dir = chosen_dir + +# Get teacher memory from sweep CSV if available, else from dir name +teacher_memory_mib = None +sweep_csv = os.path.join(puzzle_dir, "mip_sweep_results.csv") +if os.path.exists(sweep_csv): + with open(sweep_csv) as f: + reader = csv.DictReader(f) + for row in reader: + teacher_memory_mib = float(row["teacher_memory_mib"]) + break + +# Get target_memory from the chosen dir name +dir_name = os.path.basename(os.path.dirname(solutions_dir)) +mem_match = re.search(r"target_memory_([\d_]+)MiB", dir_name) +target_memory_mib = float(mem_match.group(1).replace("_", "")) if mem_match else configured_target + +results = {} +for name in ["teacher", "solution_0"]: + path = os.path.join(solutions_dir, f"{name}.json") + if not os.path.exists(path): + print(f"{name}.json not found at {path}") + sys.exit(1) + with open(path) as f: + data = json.load(f) + results[name] = {k: round(data[k]["avg"], 4) for k in KEYS if k in data} + +col_w = 30 +print(f"\n{'Metric':<{col_w}} {'Teacher':>10} {'Compressed (solution_0)':>24}") +print("-" * (col_w + 38)) +mem_teacher = f"{teacher_memory_mib:,.0f} MiB" if teacher_memory_mib else "n/a" +mem_solution = f"{target_memory_mib:,.0f} MiB" if target_memory_mib else "n/a" +print(f"{'target_memory':<{col_w}} {mem_teacher:>10} {mem_solution:>24}") +print("-" * (col_w + 38)) +for k in KEYS: + teacher_val = results["teacher"].get(k, "n/a") + student_val = results["solution_0"].get(k, "n/a") + print(f"{k:<{col_w}} {teacher_val!s:>10} {student_val!s:>24}") +print() +print(f"Results from: {solutions_dir}") +if os.path.exists(sweep_csv): + print("Sweep results: use /puzzletron mip sweep losses") diff --git a/.agents/skills/puzzletron/mip_progress.py b/.agents/skills/puzzletron/mip_progress.py new file mode 100644 index 00000000000..4d22cbd89be --- /dev/null +++ b/.agents/skills/puzzletron/mip_progress.py @@ -0,0 +1,183 @@ +# 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. + +# Generated with Claude Code +"""Progress report for the Puzzletron MIP step.""" + +import re +import sys +from datetime import datetime + +LOG = "./log.txt" +try: + lines = open(LOG).readlines() + text = "".join(lines) +except FileNotFoundError: + print("No log.txt found. Run /puzzletron mip first.") + sys.exit(0) + + +def norm(r): + """Normalize a compression rate to a canonical float string.""" + return str(float(r)) + + +def fmt(s): + """Format seconds as 'Xm Ys', or '—' if None.""" + return f"{int(s) // 60}m {int(s) % 60}s" if s is not None else "—" + + +def get_ts(line): + """Extract a datetime from a log line timestamp, or None.""" + m = re.search(r"\[(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2})", line) + return datetime.strptime(m.group(1), "%Y-%m-%d %H:%M:%S") if m else None + + +now = datetime.now().replace(microsecond=0) + +rates_match = re.search(r"Compression rates: \[(.*?)\]", text) +all_rates = [norm(r.strip()) for r in rates_match.group(1).split(",")] if rates_match else [] + +# Detect completion via step 8 marker or sweep.py:292 +complete_ts = None +for line in lines: + ts = get_ts(line) + if ts and ("Results written to:" in line or "Puzzletron Progress 8/8" in line): + complete_ts = ts + break + +cbc_matches = re.findall(r"After (\d+) nodes.*?\(([\d.]+) seconds\)", text) + +# ── Sweep disabled: single MIP solve ───────────────────────────────────────── +if not all_rates: + step7_ts = None + for line in lines: + ts = get_ts(line) + if ts and "Puzzletron Progress 7/8" in line: + step7_ts = ts + break + + end_ts = complete_ts or now + total_elapsed = int((end_ts - step7_ts).total_seconds()) if step7_ts else 0 + + cbc_detail = "" + if cbc_matches: + nodes, secs = cbc_matches[-1] + cbc_detail = f" ({int(nodes):,} nodes, {float(secs):.1f}s)" + + DIV = "─" * 62 + print("\nOverall: Puzzletron step 7/8 — MIP solve (sweep disabled)") + print(DIV) + print(f" {'Status':<10} {'Phase':<32} {'Elapsed':>8}") + print(DIV) + print(f" {'[DONE]':<10} {'Prep (loading model + scores)':<32} {'<1s':>8}") + status = "[DONE]" if complete_ts else "[RUNNING]" + label = f"MIP solve{cbc_detail}" + print(f" {status:<10} {label:<32} {fmt(total_elapsed):>8}") + print(DIV) + finished_str = ( + complete_ts.strftime("%H:%M:%S") + if complete_ts + else now.strftime("%H:%M:%S") + " (in progress)" + ) + print(f" Started: {step7_ts.strftime('%H:%M:%S') if step7_ts else '—'}") + print(f" Finished: {finished_str}") + print(f" Elapsed: {fmt(total_elapsed)}") + print(f" Remaining: {'done' if complete_ts else 'calculating...'}") + results_match = re.search(r"Results written to: (\S+)", text) + if not results_match: + results_match = re.search(r"\[run_puzzle\.py:335\]\s+(\S+)", text) + if results_match: + print(f"\n Results: {results_match.group(1)}") + sys.exit(0) + +# ── Sweep enabled: per-rate progress ───────────────────────────────────────── +rate_start = {} +for line in lines: + m = re.search(r"compression_rate=([\d.]+)", line) + if m: + r = norm(m.group(1)) + if r in all_rates and r not in rate_start: + rate_start[r] = get_ts(line) + +sweep_start = rate_start.get(all_rates[0]) if all_rates else None + +rate_done = set() +for i, r in enumerate(all_rates[:-1]): + if all_rates[i + 1] in rate_start: + rate_done.add(r) +last = all_rates[-1] if all_rates else None +if complete_ts and last and last in rate_start: + rate_done.add(last) + +rate_elapsed = {} +for i, r in enumerate(all_rates): + if r not in rate_start: + continue + next_rate = all_rates[i + 1] if i + 1 < len(all_rates) else None + end = rate_start[next_rate] if next_rate and next_rate in rate_start else (complete_ts or now) + rate_elapsed[r] = int((end - rate_start[r]).total_seconds()) + +running_rate = next((r for r in all_rates if r in rate_start and r not in rate_done), None) + +cur_detail = "" +if running_rate: + batch_matches = re.findall(r"calculate_losses_pipeline[^:]*:\s*(\d+)%.*?(\d+)/(\d+)", text) + if batch_matches: + pct, cur, total = batch_matches[-1] + cur_detail = f" — validating ({cur}/{total} batches)" + elif cbc_matches: + nodes, secs = cbc_matches[-1] + cur_detail = f" — MIP solver ({int(nodes):,} nodes, {float(secs):.1f}s)" + +end_ts = complete_ts or now +total_elapsed = int((end_ts - sweep_start).total_seconds()) if sweep_start else 0 + +done_count = len(rate_done) +remaining_count = len(all_rates) - done_count +avg_s = sum(rate_elapsed[r] for r in rate_done) / done_count if done_count else None +est_rem = ( + fmt(avg_s * remaining_count) + if avg_s and remaining_count + else ("done" if not remaining_count else "calculating...") +) + +DIV = "─" * 62 +print(f"\nOverall: Puzzletron step 7/8 — MIP sweep ({len(all_rates)} compression rates)") +print(DIV) +print(f" {'Status':<10} {'Phase':<32} {'Elapsed':>8}") +print(DIV) +print(f" {'[DONE]':<10} {'Prep (teacher memory + rate list)':<32} {'<1s':>8}") +for r in all_rates: + if r not in rate_start: + print(f" {'[ ]':<10} {f'compression_rate={r}':<32} {'pending':>8}") + elif r == running_rate: + print( + f" {'[RUNNING]':<10} {f'compression_rate={r}{cur_detail}':<32} {fmt(rate_elapsed.get(r)):>8}" + ) + else: + print(f" {'[DONE]':<10} {f'compression_rate={r}':<32} {fmt(rate_elapsed.get(r)):>8}") +print(DIV) +finished_str = ( + complete_ts.strftime("%H:%M:%S") if complete_ts else now.strftime("%H:%M:%S") + " (in progress)" +) +print(f" Started: {sweep_start.strftime('%H:%M:%S') if sweep_start else '—'}") +print(f" Finished: {finished_str}") +print(f" Elapsed: {fmt(total_elapsed)}") +print(f" Completed: {done_count}/{len(all_rates)} compression rates") +print(f" Remaining: {est_rem} estimated") +results_match = re.search(r"Results written to: (\S+)", text) +if results_match: + print(f"\n Results: {results_match.group(1)}") diff --git a/.agents/skills/puzzletron/mip_sweep.py b/.agents/skills/puzzletron/mip_sweep.py new file mode 100644 index 00000000000..b3df36721ab --- /dev/null +++ b/.agents/skills/puzzletron/mip_sweep.py @@ -0,0 +1,78 @@ +# 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. + +# Generated with Claude Code +"""Display MIP sweep results across all compression rates from the sweep CSV.""" + +import contextlib +import csv +import os +import re +import sys + +LOG = "./log.txt" +try: + text = open(LOG).read() +except FileNotFoundError: + print("No log.txt found. Run /puzzletron all first.") + sys.exit(1) + +match = re.search(r"(\S+)/mip/puzzle_solutions/", text) +if not match: + match = re.search(r"(\S+)/ckpts/teacher", text) +if not match: + print("Could not find puzzle_dir in log.txt.") + sys.exit(1) + +puzzle_dir = match.group(1) +sweep_csv = os.path.join(puzzle_dir, "mip_sweep_results.csv") + +if not os.path.exists(sweep_csv): + print(f"No sweep results found at {sweep_csv}.") + print("Enable sweep in the config YAML and re-run /puzzletron mip .") + sys.exit(1) + +with open(sweep_csv) as f: + rows = list(csv.DictReader(f)) + +if not rows: + print("Sweep CSV is empty.") + sys.exit(1) + +COLS = [ + ("compression_rate", "rate", 6), + ("target_memory_mib", "target_mem", 12), + ("actual_memory_mib", "actual_mem", 12), + ("num_params", "num_params", 12), + ("lm_loss", "lm_loss", 8), + ("token_accuracy_top_1", "top_1", 7), + ("token_accuracy_top_5", "top_5", 7), + ("token_accuracy_top_10", "top_10", 8), +] + +header = " ".join(f"{label:>{w}}" for _, label, w in COLS) +divider = "-" * len(header) +print(f"\n{header}") +print(divider) +for row in rows: + line_parts = [] + for key, _, w in COLS: + val = row.get(key, "n/a") + with contextlib.suppress(ValueError, TypeError): + val = f"{float(val):.4f}" if "." in val else f"{int(val):,}" + line_parts.append(f"{val!s:>{w}}") + print(" ".join(line_parts)) +print() +print(f"Results from: {sweep_csv}") diff --git a/.claude/skills/puzzletron b/.claude/skills/puzzletron new file mode 120000 index 00000000000..ef76b5489dd --- /dev/null +++ b/.claude/skills/puzzletron @@ -0,0 +1 @@ +../../.agents/skills/puzzletron \ No newline at end of file diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d3d0ec160ec..5236b54aa64 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,7 @@ Changelog **New Features** +- Add **experimental** ``/puzzletron`` Claude Code agent skill (``.agents/skills/puzzletron/``) with commands to run the MIP step or full pipeline, monitor progress, inspect teacher-vs-compressed accuracy, and add support for new models — including a step-by-step tutorial (``.agents/skills/puzzletron/adding_new_model_tutorial.md``). See `.agents/skills/puzzletron/README.md `_. - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred. - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``. - Add a fused Triton fast path for ``local_hessian`` NVFP4 weight-scale search (the Hessian-weighted FP8-E4M3 scale sweep). For each NVFP4 block it minimizes ``dwᵀ H dw`` over the 126 candidate scales using the per-cin-block local Hessian on tensor cores, replacing the per-weight Python reference sweep — roughly **34x** faster on a single 8192x4096 weight and bit-exact with the reference for fp32/fp16 weights. Used automatically during ``local_hessian`` calibration for both dense and fused-MoE expert weights; falls back to the reference sweep on CPU, when Triton is unavailable, or via ``MODELOPT_NVFP4_TRITON_SWEEP=0``. diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 48954a2b773..d5ce1e4535c 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -388,3 +388,10 @@ Due to non-linear extension of the runtime stats of single subblocks to the tota ## Advanced Usage Modify `llama-3_1-8B_pruneffn_memory.yaml` file for advanced compression scenarios. + +## Using with AI agents + +> **Experimental:** AI agent integration is an experimental feature and may change. + +Puzzletron ships a skill for AI coding agents (Claude Code, Cursor, Codex). +See [`.agents/skills/puzzletron/README.md`](../../.agents/skills/puzzletron/README.md) for setup, commands, and example output. diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml index bfac4ef6944..cf00853c2ae 100644 --- a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml @@ -9,7 +9,7 @@ input_hf_model_path: /workspace/hf_models/meta-llama/Llama-3.1-8B-Instruct dataset_path: /workspace/datasets/Puzzle-KD-Nemotron-Post-Training-Dataset-v2 # Working directory for puzzletron outputs -puzzle_dir: /workspace/puzzle_dir +puzzle_dir: /workspace/puzzle_dir_llama-3_1-8B # MIP memory constraint (in MiB) mip: @@ -18,7 +18,7 @@ mip: # Memory sweep configuration (optional) sweep: enabled: false - memory_compression_rates: [0.5, 0.6, 0.7, 0.8, 0.9] + memory_compression_rates: [0.5, 0.6, 0.7, 0.8, 0.9, 1.0] output_csv: ${puzzle_dir}/mip_sweep_results.csv # FFN intermediate sizes to search over (heterogeneous architecture) diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml index ce1749d9698..6b36142a3a8 100644 --- a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml @@ -3,7 +3,7 @@ autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model block_size: 8192 bos_rate: 0.5 data_column: messages -val_dataset_name: valid +val_dataset_name: validation shuffle_seed: 81436 seed: 42 fim_rate: 0 diff --git a/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/attn_pruning.yaml b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/attn_pruning.yaml new file mode 100644 index 00000000000..a660d8dc01e --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/attn_pruning.yaml @@ -0,0 +1,15 @@ +defaults: + - pruning_defaults + +activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/attn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id} + +activation_hooks_kwargs: + method: independent_kv_head_contribution + optimize_for: memory + target_layer: "self_attn.o_proj" + layer_input_descriptors_path: + +# Qwen3.5-2B has 2 KV heads in full_attention layers; only 1 grouping is possible. +# KV-head pruning is not the primary compression method for this model. +n_heads_in_group_list: [2] +gqa_init_mode: "PruneKVHeads" diff --git a/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/ffn_pruning.yaml b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/ffn_pruning.yaml new file mode 100644 index 00000000000..aedb6cd0c10 --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/ffn_pruning.yaml @@ -0,0 +1,20 @@ +defaults: + - pruning_defaults + +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.Qwen3_5FFNIntermediateLayerDescriptor + +hook_class: ${get_object:modelopt.torch.prune.importance_hooks.base_hooks.IterativeChannelContributionHook} + +activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/ffn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id} + +activation_hooks_kwargs: + method: iterative + target_layer: "mlp.down_proj" + layer_input_descriptors_path: + +# teacher_intermediate_size is 6144 +intermediate_size_list: [1280, 2560, 3584, 5120] +mlp_init_mode: "PruneByActivationsLog" diff --git a/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/hidden_dim_pruning.yaml b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/hidden_dim_pruning.yaml new file mode 100644 index 00000000000..982e35436fd --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/hidden_dim_pruning.yaml @@ -0,0 +1,15 @@ +defaults: + - pruning_defaults + +activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/hidden_dim_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id} + +activation_hooks_kwargs: + method: layer_norm_contribution + target_layer: "layernorm" + +# Qwen3.5-2B hidden_size is 2048 +hidden_size_list: [1024, 1536] +hidden_size_init_mode: "PruneByChannelRanking" +mlp_init_mode: "Truncate" +gqa_init_mode: "AverageKV" +linear_init_mode: "FromTeacher" diff --git a/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/pruning_defaults.yaml b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/pruning_defaults.yaml new file mode 100644 index 00000000000..857332fdbd7 --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/pruning_defaults.yaml @@ -0,0 +1,33 @@ +defaults: + - /validate_model_defaults + +descriptor: ${descriptor} +model_name_or_path: ${teacher_dir} +experiment_id: ${pruning.eval_samples}samples_diverse_mini +activations_log_dir: ??? +activation_hooks_kwargs: ??? + +# Data: +eval_samples: 1000 +micro_batch_size: 4 +dataset_path: ${dataset_path} +val_dataset_name: train + +# Prune ckpts +pruned_ckpts_output_dir: ${puzzle_dir}/pruning/${pruning.experiment_id} + +## FFN pruning +ffn_list: +mlp_init_mode: "Truncate" + +## KV-heads pruning +n_heads_in_group_list: +gqa_init_mode: "AverageKV" + +## Hidden dimension pruning +hidden_size_list: +hidden_size_init_mode: "PruneByChannelRanking" +linear_init_mode: "FromTeacher" + +mlp_init_config_yaml: + activations_log_dir: ${pruning.activations_log_dir} diff --git a/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/qwen3_5-0.8b_pruneffn_memory.yaml b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/qwen3_5-0.8b_pruneffn_memory.yaml new file mode 100644 index 00000000000..9c722bb6b5e --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/qwen3_5-0.8b_pruneffn_memory.yaml @@ -0,0 +1,27 @@ +defaults: + - qwen3_5 + - _self_ + +# Input Hugging Face model to compress +input_hf_model_path: /workspace/hf_models/Qwen/Qwen3.5-0.8B + +# Dataset path for pruning and NAS scoring +dataset_path: /workspace/datasets/Puzzle-KD-Nemotron-Post-Training-Dataset-v2 + +# Working directory for compression outputs +puzzle_dir: /workspace/puzzle_dir_qwen3_5-0.8b + +# MIP memory constraint (in MiB) +mip: + human_constraints: + target_memory: 10_000 # 10 GiB + # Memory sweep configuration (optional) + sweep: + enabled: false + memory_compression_rates: [0.5, 0.6, 0.7, 0.8, 0.9, 1.0] + output_csv: ${puzzle_dir}/mip_sweep_results.csv + +# FFN intermediate sizes to search over (heterogeneous architecture) +# teacher_intermediate_size is 3584 +pruning: + intermediate_size_list: [768, 1536, 2048, 3072] diff --git a/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/qwen3_5.yaml b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/qwen3_5.yaml new file mode 100644 index 00000000000..2b49a48e420 --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/qwen3_5.yaml @@ -0,0 +1,106 @@ +defaults: + - pruning: ffn_pruning + - scoring: ../validate_solutions_defaults + - realize_model: ../validate_solutions_defaults + - bypass: + - override hydra/hydra_logging: disabled + - _self_ + +puzzle_dir: ??? +descriptor: qwen3_5 +teacher_dir: ${puzzle_dir}/ckpts/teacher/ +replacement_library_path: ${puzzle_dir}/replacement_library.json +dataset_path: ??? # path to Nemotron-Post-Training-Dataset-v2 + +skip_realize_model: false + +build_replacement_library: + add_ffn_no_ops: true + add_attention_no_ops: true + +calc_subblock_stats: + batch_sizes: [64, 96, 128] + prefill_seq_len: 4096 + generation_seq_len: 4096 + num_active_tokens_override: + prefill_queue_size: 0 + allocate_prefill_query: false + runtime_stats: + backend: trt_torch + benchmark_iterations: + merge_with_existing_stats: false + subblock_stats_filename: "subblock_stats.json" + moe_stats_filename: "moe_stats.json" + +scoring: + descriptor: ${descriptor} + solutions_to_validate: + skip_existing_solutions: true + + replacement_library_path: ${replacement_library_path} + solutions_path: ${to_path:${puzzle_dir}/single_sequence_replacement_solutions.json} + teacher_dir: ${to_path:${teacher_dir}} + output_dir: ${puzzle_dir}/single_sequence_replacement_solutions--validation + + eval_samples: 8 + micro_batch_size: 1 + seed: 42 + shuffle_seed: 444 + dataset_path: ${dataset_path} + +mip: + single_block_replacement_validation_dir: ${to_path:${scoring.output_dir}} + subblock_stats_path: ${to_path:${puzzle_dir}/${calc_subblock_stats.subblock_stats_filename}} + output_path: ${to_path:${puzzle_dir}/mip/puzzle_solutions} + gathered_metrics_path: + puzzle_profile: + + objective: metrics.cosine_embedding_loss_hidden_states + bigger_is_better: false + + subblock_stats_args: + - batch_size: 96 + weights_dtype: torch.bfloat16 + activations_dtype: torch.bfloat16 + kv_cache_dtype: torch.bfloat16 + + report_additional_costs: + - stats.memory_mib + - stats.num_params + - stats.num_kv_heads + - stats.has_attention + - stats.has_ffn + - stats.kv_cache_memory_mib + - stats.attention_memory_mib + - stats.ffn_memory_mib + - stats.ffn_num_params + - stats.attention_num_params + + human_constraints: + target_memory: 20_000 + num_params: 1_500_000_000 + + mip_constraints: + metric_overrides: + max_seconds_per_solution: 60 + +realize_model: + descriptor: ${descriptor} + teacher_dir: ${to_path:${teacher_dir}} + tokenizer_name: ${to_path:${teacher_dir}} + replacement_library_path: ${replacement_library_path} + save_models: true + solutions_path: + + skip_validation: false + eval_samples: 8 + micro_batch_size: 1 + seed: 42 + shuffle_seed: 444 + dataset_path: ${dataset_path} + +nccl_timeout_minutes: ${timedelta_minutes:10} + +hydra: + run: + dir: ${puzzle_dir}/hydra_logs/${now:%Y-%m-%d}/${now:%H-%M-%S} diff --git a/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/validate_model_defaults.yaml b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/validate_model_defaults.yaml new file mode 100644 index 00000000000..6b36142a3a8 --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/validate_model_defaults.yaml @@ -0,0 +1,17 @@ +model_dtype: torch.bfloat16 # dtype to cast the model for validate_model +autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model +block_size: 8192 +bos_rate: 0.5 +data_column: messages +val_dataset_name: validation +shuffle_seed: 81436 +seed: 42 +fim_rate: 0 +fim_spm_rate: 0 +source_datasets_to_discard: +varlen: false +write_results: false +calc_losses_on_cpu: false +activations_log_dir: +model_name_or_path: +load_dataset_fn: ${get_object:modelopt.torch.puzzletron.utils.data.dataloaders.load_from_disk_fn} diff --git a/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/validate_solutions_defaults.yaml b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/validate_solutions_defaults.yaml new file mode 100644 index 00000000000..ec139023794 --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/validate_solutions_defaults.yaml @@ -0,0 +1,10 @@ +defaults: + - /validate_model_defaults + - _self_ + +solutions_to_validate: +skip_validation: false +save_models: false +bigger_is_better: false +sort_solutions_by: +calculate_full_score_ablations: false diff --git a/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/attn_pruning.yaml b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/attn_pruning.yaml new file mode 100644 index 00000000000..a660d8dc01e --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/attn_pruning.yaml @@ -0,0 +1,15 @@ +defaults: + - pruning_defaults + +activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/attn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id} + +activation_hooks_kwargs: + method: independent_kv_head_contribution + optimize_for: memory + target_layer: "self_attn.o_proj" + layer_input_descriptors_path: + +# Qwen3.5-2B has 2 KV heads in full_attention layers; only 1 grouping is possible. +# KV-head pruning is not the primary compression method for this model. +n_heads_in_group_list: [2] +gqa_init_mode: "PruneKVHeads" diff --git a/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/ffn_pruning.yaml b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/ffn_pruning.yaml new file mode 100644 index 00000000000..aedb6cd0c10 --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/ffn_pruning.yaml @@ -0,0 +1,20 @@ +defaults: + - pruning_defaults + +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.Qwen3_5FFNIntermediateLayerDescriptor + +hook_class: ${get_object:modelopt.torch.prune.importance_hooks.base_hooks.IterativeChannelContributionHook} + +activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/ffn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id} + +activation_hooks_kwargs: + method: iterative + target_layer: "mlp.down_proj" + layer_input_descriptors_path: + +# teacher_intermediate_size is 6144 +intermediate_size_list: [1280, 2560, 3584, 5120] +mlp_init_mode: "PruneByActivationsLog" diff --git a/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/hidden_dim_pruning.yaml b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/hidden_dim_pruning.yaml new file mode 100644 index 00000000000..982e35436fd --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/hidden_dim_pruning.yaml @@ -0,0 +1,15 @@ +defaults: + - pruning_defaults + +activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/hidden_dim_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id} + +activation_hooks_kwargs: + method: layer_norm_contribution + target_layer: "layernorm" + +# Qwen3.5-2B hidden_size is 2048 +hidden_size_list: [1024, 1536] +hidden_size_init_mode: "PruneByChannelRanking" +mlp_init_mode: "Truncate" +gqa_init_mode: "AverageKV" +linear_init_mode: "FromTeacher" diff --git a/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/pruning_defaults.yaml b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/pruning_defaults.yaml new file mode 100644 index 00000000000..857332fdbd7 --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/pruning_defaults.yaml @@ -0,0 +1,33 @@ +defaults: + - /validate_model_defaults + +descriptor: ${descriptor} +model_name_or_path: ${teacher_dir} +experiment_id: ${pruning.eval_samples}samples_diverse_mini +activations_log_dir: ??? +activation_hooks_kwargs: ??? + +# Data: +eval_samples: 1000 +micro_batch_size: 4 +dataset_path: ${dataset_path} +val_dataset_name: train + +# Prune ckpts +pruned_ckpts_output_dir: ${puzzle_dir}/pruning/${pruning.experiment_id} + +## FFN pruning +ffn_list: +mlp_init_mode: "Truncate" + +## KV-heads pruning +n_heads_in_group_list: +gqa_init_mode: "AverageKV" + +## Hidden dimension pruning +hidden_size_list: +hidden_size_init_mode: "PruneByChannelRanking" +linear_init_mode: "FromTeacher" + +mlp_init_config_yaml: + activations_log_dir: ${pruning.activations_log_dir} diff --git a/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/qwen3_5-2B_pruneffn_memory.yaml b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/qwen3_5-2B_pruneffn_memory.yaml new file mode 100644 index 00000000000..1dad6b29be8 --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/qwen3_5-2B_pruneffn_memory.yaml @@ -0,0 +1,22 @@ +defaults: + - qwen3_5 + - _self_ + +# Input Hugging Face model to compress +input_hf_model_path: /workspace/hf_models/Qwen3.5-2B + +# Dataset path for pruning and NAS scoring +dataset_path: /workspace/datasets/Puzzle-KD-Nemotron-Post-Training-Dataset-v2 + +# Working directory for compression outputs +puzzle_dir: /workspace/puzzle_dir_qwen3_5-2B + +# MIP memory constraint (in MiB) +mip: + human_constraints: + target_memory: 20_000 # 20 GiB + +# FFN intermediate sizes to search over (heterogeneous architecture) +# teacher_intermediate_size is 6144 +pruning: + intermediate_size_list: [1280, 2560, 3584, 5120] diff --git a/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/qwen3_5.yaml b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/qwen3_5.yaml new file mode 100644 index 00000000000..25d9d054bac --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/qwen3_5.yaml @@ -0,0 +1,106 @@ +defaults: + - pruning: ffn_pruning + - scoring: ../validate_solutions_defaults + - realize_model: ../validate_solutions_defaults + - bypass: + - override hydra/hydra_logging: disabled + - _self_ + +puzzle_dir: ??? +descriptor: qwen3_5 +teacher_dir: ${puzzle_dir}/ckpts/teacher/ +replacement_library_path: ${puzzle_dir}/replacement_library.json +dataset_path: ??? # path to Nemotron-Post-Training-Dataset-v2 + +skip_realize_model: false + +build_replacement_library: + add_ffn_no_ops: true + add_attention_no_ops: true + +calc_subblock_stats: + batch_sizes: [64, 96, 128] + prefill_seq_len: 4096 + generation_seq_len: 4096 + num_active_tokens_override: + prefill_queue_size: 0 + allocate_prefill_query: false + runtime_stats: + backend: trt_torch + benchmark_iterations: + merge_with_existing_stats: false + subblock_stats_filename: "subblock_stats.json" + moe_stats_filename: "moe_stats.json" + +scoring: + descriptor: ${descriptor} + solutions_to_validate: + skip_existing_solutions: true + + replacement_library_path: ${replacement_library_path} + solutions_path: ${to_path:${puzzle_dir}/single_sequence_replacement_solutions.json} + teacher_dir: ${to_path:${teacher_dir}} + output_dir: ${puzzle_dir}/single_sequence_replacement_solutions--validation + + eval_samples: 128 + micro_batch_size: 1 + seed: 42 + shuffle_seed: 444 + dataset_path: ${dataset_path} + +mip: + single_block_replacement_validation_dir: ${to_path:${scoring.output_dir}} + subblock_stats_path: ${to_path:${puzzle_dir}/${calc_subblock_stats.subblock_stats_filename}} + output_path: ${to_path:${puzzle_dir}/mip/puzzle_solutions} + gathered_metrics_path: + puzzle_profile: + + objective: metrics.cosine_embedding_loss_hidden_states + bigger_is_better: false + + subblock_stats_args: + - batch_size: 96 + weights_dtype: torch.bfloat16 + activations_dtype: torch.bfloat16 + kv_cache_dtype: torch.bfloat16 + + report_additional_costs: + - stats.memory_mib + - stats.num_params + - stats.num_kv_heads + - stats.has_attention + - stats.has_ffn + - stats.kv_cache_memory_mib + - stats.attention_memory_mib + - stats.ffn_memory_mib + - stats.ffn_num_params + - stats.attention_num_params + + human_constraints: + target_memory: 20_000 + num_params: 1_500_000_000 + + mip_constraints: + metric_overrides: + max_seconds_per_solution: 60 + +realize_model: + descriptor: ${descriptor} + teacher_dir: ${to_path:${teacher_dir}} + tokenizer_name: ${to_path:${teacher_dir}} + replacement_library_path: ${replacement_library_path} + save_models: true + solutions_path: + + skip_validation: false + eval_samples: 128 + micro_batch_size: 1 + seed: 42 + shuffle_seed: 444 + dataset_path: ${dataset_path} + +nccl_timeout_minutes: ${timedelta_minutes:10} + +hydra: + run: + dir: ${puzzle_dir}/hydra_logs/${now:%Y-%m-%d}/${now:%H-%M-%S} diff --git a/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/validate_model_defaults.yaml b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/validate_model_defaults.yaml new file mode 100644 index 00000000000..6b36142a3a8 --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/validate_model_defaults.yaml @@ -0,0 +1,17 @@ +model_dtype: torch.bfloat16 # dtype to cast the model for validate_model +autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model +block_size: 8192 +bos_rate: 0.5 +data_column: messages +val_dataset_name: validation +shuffle_seed: 81436 +seed: 42 +fim_rate: 0 +fim_spm_rate: 0 +source_datasets_to_discard: +varlen: false +write_results: false +calc_losses_on_cpu: false +activations_log_dir: +model_name_or_path: +load_dataset_fn: ${get_object:modelopt.torch.puzzletron.utils.data.dataloaders.load_from_disk_fn} diff --git a/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/validate_solutions_defaults.yaml b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/validate_solutions_defaults.yaml new file mode 100644 index 00000000000..ec139023794 --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/validate_solutions_defaults.yaml @@ -0,0 +1,10 @@ +defaults: + - /validate_model_defaults + - _self_ + +solutions_to_validate: +skip_validation: false +save_models: false +bigger_is_better: false +sort_solutions_by: +calculate_full_score_ablations: false diff --git a/modelopt/torch/puzzletron/anymodel/models/__init__.py b/modelopt/torch/puzzletron/anymodel/models/__init__.py index c126d61b887..20911533833 100644 --- a/modelopt/torch/puzzletron/anymodel/models/__init__.py +++ b/modelopt/torch/puzzletron/anymodel/models/__init__.py @@ -26,4 +26,5 @@ from .qwen3 import * if _Version(_transformers_version) >= _Version("4.57.0"): + from .qwen3_5 import * from .qwen3_vl import * diff --git a/modelopt/torch/puzzletron/anymodel/models/qwen3_5/__init__.py b/modelopt/torch/puzzletron/anymodel/models/qwen3_5/__init__.py new file mode 100644 index 00000000000..4415aa08ca8 --- /dev/null +++ b/modelopt/torch/puzzletron/anymodel/models/qwen3_5/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +from .qwen3_5_converter import * +from .qwen3_5_model_descriptor import * diff --git a/modelopt/torch/puzzletron/anymodel/models/qwen3_5/qwen3_5_converter.py b/modelopt/torch/puzzletron/anymodel/models/qwen3_5/qwen3_5_converter.py new file mode 100644 index 00000000000..b94a4cc770b --- /dev/null +++ b/modelopt/torch/puzzletron/anymodel/models/qwen3_5/qwen3_5_converter.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +# mypy: ignore-errors + +import copy +from pathlib import Path +from typing import List + +from ....block_config import AttentionConfig, BlockConfig, FFNConfig +from ....tools.checkpoint_utils_hf import load_model_config, save_model_config +from ...converter import Converter, ConverterFactory + +__all__ = ["Qwen3_5Converter"] + +_LANGUAGE_MODEL_PREFIX = "model.language_model." + + +@ConverterFactory.register_decorator("qwen3_5") +class Qwen3_5Converter(Converter): + @staticmethod + def create_block_configs_from_main_config(config) -> List[BlockConfig]: + text_config = config.text_config if hasattr(config, "text_config") else config + return [ + BlockConfig( + attention=AttentionConfig( + no_op=False, num_key_value_heads=text_config.num_key_value_heads + ), + ffn=FFNConfig(no_op=False, intermediate_size=text_config.intermediate_size), + ).to_dict() + for _ in range(text_config.num_hidden_layers) + ] + + @classmethod + def convert_configs_in_dirs( + cls, input_dir: Path, output_dir: Path, trust_remote_code: bool = False + ): + """Save text_config (not the full VLM config) so downstream code can access + num_hidden_layers and other text-model fields directly.""" + config = load_model_config(input_dir, trust_remote_code=trust_remote_code) + text_config = config.text_config if hasattr(config, "text_config") else config + block_configs = cls.create_block_configs_from_main_config(config) + out_config = copy.deepcopy(text_config) + out_config.block_configs = block_configs + save_model_config(out_config, output_dir) + return out_config + + @staticmethod + def convert_weight_name(name: str) -> str: + """Remap VLM weight names to text-model paths. + + model.language_model.X → model.X + All other names are unchanged (lm_head, etc.). + """ + if name.startswith(_LANGUAGE_MODEL_PREFIX): + return "model." + name[len(_LANGUAGE_MODEL_PREFIX) :] + return name diff --git a/modelopt/torch/puzzletron/anymodel/models/qwen3_5/qwen3_5_model_descriptor.py b/modelopt/torch/puzzletron/anymodel/models/qwen3_5/qwen3_5_model_descriptor.py new file mode 100644 index 00000000000..9d31ac1d1a1 --- /dev/null +++ b/modelopt/torch/puzzletron/anymodel/models/qwen3_5/qwen3_5_model_descriptor.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +# mypy: ignore-errors + +import re +from dataclasses import dataclass, field +from typing import Dict, Iterable, List + +from torch import nn +from transformers.models.qwen3_5.modeling_qwen3_5 import ( + Qwen3_5DecoderLayer, + Qwen3_5TextRotaryEmbedding, +) + +from ....block_config import BlockConfig +from ....pruning.ffn_intermediate_pruning_mixin import FFNIntermediateLayerDescriptor +from ....pruning.kv_heads_pruning_mixin import KVHeadsLayerDescriptor +from ....utils.dummy_modules import DummyBlock +from ...model_descriptor import ModelDescriptor, ModelDescriptorFactory +from ...puzzformer.no_op import MatchingZeros, Same, return_tuple_of_size + +__all__ = [ + "Qwen3_5ModelDescriptor", + "Qwen3_5FFNIntermediateLayerDescriptor", + "Qwen3_5KVHeadsLayerDescriptor", +] + +# Weight prefixes that belong to the vision encoder and MTP head — not part of the +# text model and skipped during subblock conversion. +_NON_TEXT_PREFIXES = ("model.visual.", "mtp.") + + +@ModelDescriptorFactory.register_decorator("qwen3_5") +class Qwen3_5ModelDescriptor(ModelDescriptor): + @staticmethod + def get_language_model_config(config): + """Qwen3.5 is a VLM; language model parameters live in the nested text_config.""" + return config.text_config if hasattr(config, "text_config") else config + + @staticmethod + def decoder_layer_cls(): + return Qwen3_5DecoderLayer + + @classmethod + def create_dummy_block(cls, original_layer: nn.Module, block_index: int) -> nn.Module: + """Preserve layer_type so the model forward can select the right attention path.""" + dummy = DummyBlock(block_index=block_index) + if hasattr(original_layer, "layer_type"): + dummy.layer_type = original_layer.layer_type + return dummy + + @staticmethod + def block_config_to_layer_overrides(block_config: BlockConfig): + return { + "intermediate_size": block_config.ffn.intermediate_size, + "num_key_value_heads": block_config.attention.num_key_value_heads, + } + + @staticmethod + def attn_no_op_post_init(decoder_layer: Qwen3_5DecoderLayer): + """Zero out the attention sub-block, branching on the hybrid layer type. + + full_attention layers return a (hidden_states, attn_weights) tuple; + linear_attention (GatedDeltaNet) layers return hidden_states directly. + """ + decoder_layer.input_layernorm = Same() + if decoder_layer.layer_type == "full_attention": + decoder_layer.self_attn = return_tuple_of_size(MatchingZeros, size=2)() + else: + decoder_layer.linear_attn = MatchingZeros() + + @staticmethod + def mlp_no_op_post_init(decoder_layer: Qwen3_5DecoderLayer): + decoder_layer.post_attention_layernorm = Same() + decoder_layer.mlp = MatchingZeros() + + @staticmethod + def init_rotary_embedding(model, runtime): + # After conversion the model is Qwen3_5ForCausalLM; text model is at model.model + model.model.rotary_emb = Qwen3_5TextRotaryEmbedding(config=model.config).to( + device=runtime.device + ) + + @staticmethod + def input_embedding_name(): + return "model.embed_tokens" + + @staticmethod + def output_embedding_name(): + return "lm_head" + + @staticmethod + def final_norm_name(): + return "model.norm" + + @staticmethod + def layer_block_name(index: int): + return f"model.layers.{index}" + + @classmethod + def get_weight_groups( + cls, layer_names: Iterable[str], num_hidden_layers: int + ) -> Dict[str, List[str]]: + """Filter out vision/MTP weights before grouping. + + get_weight_groups is called from two places with different name formats: + - convert_model_weights: original VLM checkpoint names (model.language_model.*) + - _save_checkpoint: already-converted state dict names (model.*) + + Predicates use model.* format. When original names are detected we remap + internally for matching but restore originals in the returned groups so + that the param_to_file lookup in convert_model_weights still works. + """ + _lm_prefix = "model.language_model." + text_names = [n for n in layer_names if not n.startswith(_NON_TEXT_PREFIXES)] + + if not any(n.startswith(_lm_prefix) for n in text_names): + # Already-converted names — pass through directly. + return super().get_weight_groups(text_names, num_hidden_layers) + + # Original checkpoint names: remap to model.* for predicate matching, + # then un-remap so returned groups contain the original names. + name_map: Dict[str, str] = {} # remapped → original + remapped = [] + for n in text_names: + r = "model." + n[len(_lm_prefix) :] if n.startswith(_lm_prefix) else n + name_map[r] = n + remapped.append(r) + + groups_remapped = super().get_weight_groups(remapped, num_hidden_layers) + return {group: [name_map[r] for r in names] for group, names in groups_remapped.items()} + + @staticmethod + def layer_name_predicates(num_layers: int) -> Dict[str, re.Pattern]: + # Predicates use converted model.* names (matching Qwen3_5ForCausalLM). + # get_weight_groups normalises original checkpoint names before matching. + layer_name_patterns = { + "embeddings": re.compile(r"^model\.embed_tokens\.weight$"), + "lm_head": re.compile(r"^(model\.norm\.weight|lm_head\.weight)$"), + } + + def build_ffn_predicates() -> Dict[str, re.Pattern]: + return { + f"block_{layer_idx}_ffn": re.compile( + rf"^model\.layers\.{layer_idx}\.(post_attention_layernorm\.weight" + r"|mlp\.up_proj\.weight" + r"|mlp\.gate_proj\.weight" + r"|mlp\.down_proj\.weight)$" + ) + for layer_idx in range(num_layers) + } + + def build_attention_predicates() -> Dict[str, re.Pattern]: + return { + f"block_{layer_idx}_attention": re.compile( + rf"^model\.layers\.{layer_idx}\.(input_layernorm\.weight" + # full_attention (Qwen3_5Attention) weights + r"|self_attn\.q_proj\.weight" + r"|self_attn\.k_proj\.weight" + r"|self_attn\.v_proj\.weight" + r"|self_attn\.o_proj\.weight" + r"|self_attn\.q_norm\.weight" + r"|self_attn\.k_norm\.weight" + # linear_attention (GatedDeltaNet) weights + r"|linear_attn\.in_proj_qkv\.weight" + r"|linear_attn\.in_proj_z\.weight" + r"|linear_attn\.in_proj_b\.weight" + r"|linear_attn\.in_proj_a\.weight" + r"|linear_attn\.out_proj\.weight" + r"|linear_attn\.conv1d\.weight" + r"|linear_attn\.norm\.weight" + r"|linear_attn\.dt_bias" + r"|linear_attn\.A_log)$" + ) + for layer_idx in range(num_layers) + } + + layer_name_patterns.update(**build_ffn_predicates(), **build_attention_predicates()) + return layer_name_patterns + + +@dataclass +class Qwen3_5FFNIntermediateLayerDescriptor(FFNIntermediateLayerDescriptor): + down_proj_name: str = "mlp.down_proj" + ffn_prefix_name: str = "model.layers.{layer_idx}.mlp" + linear_weight_names: List[str] = field( + default_factory=lambda: ["down_proj", "gate_proj", "up_proj"] + ) + + +@dataclass +class Qwen3_5KVHeadsLayerDescriptor(KVHeadsLayerDescriptor): + o_proj_name: str = "self_attn.o_proj" + attn_prefix_name: str = "model.layers.{layer_idx}.self_attn" + qkvo_weight_names: List[str] = field( + default_factory=lambda: ["q_proj", "k_proj", "v_proj", "o_proj"] + )