NVIDIA · danielkorzekwa · Jun 17, 2026 · Jun 17, 2026 · Jun 17, 2026 · Jun 17, 2026
@@ -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.
@@ -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 <command> [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 <nproc_per_node>` — 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 <nproc_per_node>` — 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 <hf_model_path>` — Implement descriptor, converter, and configs for an unsupported model
+
+Usage: `/puzzletron <command> [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 <value>` 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 \<nproc_per_node\>
+
+Run the following Bash command, substituting `<nproc_per_node>` with the parsed value:
+
+```bash
+set -o pipefail && export PYTHONPATH=$PYTHONPATH:/workspace/Model-Optimizer && \
+torchrun --nproc_per_node <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 <value>` 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 \<nproc_per_node\>
+
+Run the following Bash command, substituting `<nproc_per_node>` with the parsed value:
+
+```bash
+set -o pipefail && export PYTHONPATH=$PYTHONPATH:/workspace/Model-Optimizer && \
+torchrun --nproc_per_node <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('<hf_model_path>', 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('<hf_model_path>', 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('<hf_model_path>/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/<model_type>/__init__.py`**
+
+```python
+from .<model_type>_converter import *
+from .<model_type>_model_descriptor import *
+```
+
+**`modelopt/torch/puzzletron/anymodel/models/<model_type>/<model_type>_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/<model_type>/<model_type>_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 .<model_type> import *
+```
+
+**Compression config** at `examples/puzzletron/configs/<model_type>-<size>_pruneffn_memory/`:
+- Base YAML (`<model_type>.yaml`): `descriptor: <model_type>`, 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:', '<model_type>' in ModelDescriptorFactory.CLASS_MAPPING)
+print('converter: ', '<model_type>' 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 <model_name> on <N> GPUs
+```