From 8ab8f7a1ae7106f80edceefcdce4e4f0f89660f1 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Wed, 17 Jun 2026 07:52:29 -0700 Subject: [PATCH 01/28] Add a dummy puzzletron skill Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/SKILL.md | 36 ++++++++++++++++++++++++++++++ .claude/skills/puzzletron | 1 + 2 files changed, 37 insertions(+) create mode 100644 .agents/skills/puzzletron/SKILL.md create mode 120000 .claude/skills/puzzletron diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md new file mode 100644 index 00000000000..c0cab4f02c1 --- /dev/null +++ b/.agents/skills/puzzletron/SKILL.md @@ -0,0 +1,36 @@ +--- +name: puzzletron +description: End-to-end workflow for model pruning and MIP-based optimization. Use `all` to run the full workflow or `mip_sweep` to run the MIP sweep. Usage: /puzzletron +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 args do **not exactly match** `all` or `mip_sweep`, 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: +- `all` — Run the full puzzletron workflow +- `mip_sweep` — Run the MIP sweep + +Usage: `/puzzletron ` + +--- + +**STEP 2 — Only if args exactly match a command name, execute it. Never reach this step if args were empty.** + +## Command: all + +Return the following message: hello world: puzzletron all message2 + +## Command: mip_sweep + +Return the following message: hello world: puzzletron mip sweep2 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 From 9a6d71623e1d42e3c25582b2b23b14b51be079f0 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Wed, 17 Jun 2026 08:55:51 -0700 Subject: [PATCH 02/28] Add progress comnand for puzzletron mip_sweep Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/README.md | 58 +++++++++++ .agents/skills/puzzletron/SKILL.md | 148 ++++++++++++++++++++++++++-- examples/puzzletron/README.md | 7 ++ 3 files changed, 203 insertions(+), 10 deletions(-) create mode 100644 .agents/skills/puzzletron/README.md diff --git a/.agents/skills/puzzletron/README.md b/.agents/skills/puzzletron/README.md new file mode 100644 index 00000000000..af43fb496a4 --- /dev/null +++ b/.agents/skills/puzzletron/README.md @@ -0,0 +1,58 @@ +# 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. + +For full environment setup, model configuration, and algorithm details see +[examples/puzzletron/README.md](../../examples/puzzletron/README.md). + +## Using with AI agents + +> **Experimental:** AI agent integration is an experimental feature and may change. + +| Agent | How to invoke | +|---|---| +| **Claude Code** | `/puzzletron ` in the chat | + +## Commands + +### `mip_sweep ` + +Runs the MIP sweep across multiple compression rates. `nproc_per_node` is the number of GPUs per node. + +```text +/puzzletron mip_sweep 4 +``` + +Output is streamed live and also written to `./log.txt`. + +### `mip_sweep progress` + +Parses `./log.txt` and prints a structured progress report: prep steps, per-compression-rate +sub-steps, and a timing summary with elapsed and estimated remaining time. + +```text +/puzzletron mip_sweep progress +``` + +Example output: + +```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 + [RUNNING] compression_rate=0.6 — validating (47/128 batches) 1m 14s + [ ] compression_rate=0.7 pending + [ ] compression_rate=0.8 pending + [ ] compression_rate=0.9 pending + [ ] compression_rate=1.0 pending +────────────────────────────────────────────────────────────── + Started: 08:05:30 + Now: 08:10:56 + Elapsed: 5m 26s + Completed: 1/6 compression rates (avg 3m 52s/rate) + Remaining: ~19m 22s estimated +``` diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md index c0cab4f02c1..d3a951ba5db 100644 --- a/.agents/skills/puzzletron/SKILL.md +++ b/.agents/skills/puzzletron/SKILL.md @@ -1,6 +1,6 @@ --- name: puzzletron -description: End-to-end workflow for model pruning and MIP-based optimization. Use `all` to run the full workflow or `mip_sweep` to run the MIP sweep. Usage: /puzzletron +description: End-to-end workflow for model pruning and MIP-based optimization. Use `mip_sweep` to run the MIP sweep. Usage: /puzzletron license: Apache-2.0 --- @@ -11,26 +11,154 @@ license: Apache-2.0 **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 args do **not exactly match** `all` or `mip_sweep`, 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_sweep`, 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: -- `all` — Run the full puzzletron workflow -- `mip_sweep` — Run the MIP sweep +- `mip_sweep ` — Run the MIP sweep (nproc_per_node: number of GPUs per node) +- `mip_sweep progress` — Show live MIP sweep progress with timing summary -Usage: `/puzzletron ` +Usage: `/puzzletron [args]` --- -**STEP 2 — Only if args exactly match a command name, execute it. Never reach this step if args were empty.** +**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 +## Command: mip_sweep -Return the following message: hello world: puzzletron all message2 +Parse the second word of args. -## Command: mip_sweep +- If no second word is provided, ask the user: "Please provide the number of GPUs per node (nproc_per_node)." and **STOP**. +- If the second word is exactly `progress`, execute the **mip_sweep progress** sub-command below. +- Otherwise treat the second word as `nproc_per_node` and run the sweep. + +### mip_sweep \ + +Run the following Bash command, substituting `` with the parsed value: + +```bash +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_sweep progress + +Run the following Python script verbatim. Do not modify it. Present the output to the user wrapped in a fenced code block (``` ... ```). + +```bash +python3 - << 'PYEOF' +import re, 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_sweep first.") + sys.exit(0) + +def norm(r): return str(float(r)) +def fmt(s): return f"{int(s)//60}m {int(s)%60}s" if s is not None else "—" +def get_ts(line): + 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 + +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 [] + +# Collect start timestamp per rate +rate_start = {} +for line in lines: + if 'sweep.py:258' in line: + 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) + +now = datetime.now().replace(microsecond=0) +sweep_start = rate_start.get(all_rates[0]) if all_rates else None + +# Rate is done when the next rate has started; last rate done when sweep.py:287 appears +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] +sweep_complete_ts = None +for line in lines: + ts = get_ts(line) + if ts and 'sweep.py:292' in line: + sweep_complete_ts = ts + break +if sweep_complete_ts and last in rate_start: + rate_done.add(last) + +# Per-rate elapsed = next rate start - this rate start (or completion ts or now for last) +rate_elapsed = {} +for i, r in enumerate(all_rates): + if r not in rate_start: + continue + if i + 1 < len(all_rates): + end = rate_start[all_rates[i + 1]] + else: + end = sweep_complete_ts if sweep_complete_ts else now + rate_elapsed[r] = int((end - rate_start[r]).total_seconds()) + +# Currently running rate +running_rate = next((r for r in all_rates if r in rate_start and r not in rate_done), None) + +# Sub-step detail for running rate +cur_detail = "" +if running_rate: + batch_matches = re.findall(r'calculate_losses_pipeline[^:]*:\s*(\d+)%.*?(\d+)/(\d+)', text) + cbc_matches = re.findall(r'After (\d+) nodes.*?\(([\d.]+) seconds\)', 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 = sweep_complete_ts if sweep_complete_ts else 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 -Return the following message: hello world: puzzletron mip sweep2 +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] {'Prep (teacher memory + rate list)':<32} {'<1s':>8}") +for r in all_rates: + if r not in rate_start: + print(f" [ ] {f'compression_rate={r}':<32} {'pending':>8}") + elif r == running_rate: + detail = cur_detail + print(f" [RUNNING] {f'compression_rate={r}{detail}':<32} {fmt(rate_elapsed.get(r)):>8}") + else: + print(f" [DONE] {f'compression_rate={r}':<32} {fmt(rate_elapsed.get(r)):>8}") +print(DIV) +print(f" Started: {sweep_start.strftime('%H:%M:%S') if sweep_start else '—'}") +print(f" Finished: {sweep_complete_ts.strftime('%H:%M:%S') if sweep_complete_ts else now.strftime('%H:%M:%S') + ' (in progress)'}") +print(f" Elapsed: {fmt(total_elapsed)}") +print(f" Completed: {done_count}/{len(all_rates)} compression rates", end="") +print(f" (avg {fmt(avg_s)}/rate)" if avg_s else "") +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)}") +PYEOF +``` 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. From 1e794636d8fc2f9b8bde9ef191a593cdb7d09208 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Wed, 17 Jun 2026 09:02:59 -0700 Subject: [PATCH 03/28] update puzzletron readme Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/README.md | 51 +++++++++++++---------------- 1 file changed, 23 insertions(+), 28 deletions(-) diff --git a/.agents/skills/puzzletron/README.md b/.agents/skills/puzzletron/README.md index af43fb496a4..62e94f54fd5 100644 --- a/.agents/skills/puzzletron/README.md +++ b/.agents/skills/puzzletron/README.md @@ -6,53 +6,48 @@ This skill exposes it as a slash command for AI coding agents. For full environment setup, model configuration, and algorithm details see [examples/puzzletron/README.md](../../examples/puzzletron/README.md). -## Using with AI agents - > **Experimental:** AI agent integration is an experimental feature and may change. -| Agent | How to invoke | -|---|---| -| **Claude Code** | `/puzzletron ` in the chat | - -## Commands +Run `/puzzletron` with no arguments to see available commands. -### `mip_sweep ` +## Running the MIP sweep -Runs the MIP sweep across multiple compression rates. `nproc_per_node` is the number of GPUs per node. +Start the sweep by telling the agent how many GPUs per node to use: ```text /puzzletron mip_sweep 4 ``` -Output is streamed live and also written to `./log.txt`. - -### `mip_sweep progress` - -Parses `./log.txt` and prints a structured progress report: prep steps, per-compression-rate -sub-steps, and a timing summary with elapsed and estimated remaining time. +Output is streamed live and also written to `./log.txt`. While it runs (or after it finishes), +check progress with: ```text /puzzletron mip_sweep progress ``` -Example output: +Example output when complete: ```text Overall: Puzzletron step 7/8 — MIP sweep (6 compression rates) ────────────────────────────────────────────────────────────── - Status Phase Elapsed + Status Phase Elapsed ────────────────────────────────────────────────────────────── - [DONE] Prep (teacher memory + rate list) <1s - [DONE] compression_rate=0.5 3m 52s - [RUNNING] compression_rate=0.6 — validating (47/128 batches) 1m 14s - [ ] compression_rate=0.7 pending - [ ] compression_rate=0.8 pending - [ ] compression_rate=0.9 pending - [ ] compression_rate=1.0 pending + [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 - Now: 08:10:56 - Elapsed: 5m 26s - Completed: 1/6 compression rates (avg 3m 52s/rate) - Remaining: ~19m 22s estimated + Finished: 08:30:38 + Elapsed: 25m 8s + Completed: 6/6 compression rates (avg 4m 11s/rate) + 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. From f99c01f837be740769284de76f84f766c50c2df9 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Wed, 17 Jun 2026 09:09:47 -0700 Subject: [PATCH 04/28] fix typo in SKILL.md Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/SKILL.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md index d3a951ba5db..87258b0cf44 100644 --- a/.agents/skills/puzzletron/SKILL.md +++ b/.agents/skills/puzzletron/SKILL.md @@ -1,6 +1,6 @@ --- name: puzzletron -description: End-to-end workflow for model pruning and MIP-based optimization. Use `mip_sweep` to run the MIP sweep. Usage: /puzzletron +description: "End-to-end workflow for model pruning and MIP-based optimization. Use `mip_sweep` to run the MIP sweep. Usage: /puzzletron " license: Apache-2.0 --- From 131296c1cf46fecd564e2f750806bbf19024f0c8 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Thu, 18 Jun 2026 00:01:21 -0700 Subject: [PATCH 05/28] Add PYTHONPATH to mip_sweep skill Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/SKILL.md | 1 + 1 file changed, 1 insertion(+) diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md index 87258b0cf44..afda7a25b8a 100644 --- a/.agents/skills/puzzletron/SKILL.md +++ b/.agents/skills/puzzletron/SKILL.md @@ -40,6 +40,7 @@ Parse the second word of args. Run the following Bash command, substituting `` with the parsed value: ```bash +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" From c5c201561111212929aa4c70f0af1b6f4bf1fb7a Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Thu, 18 Jun 2026 00:52:58 -0700 Subject: [PATCH 06/28] add skill for puzzletron all Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/README.md | 42 +++++- .agents/skills/puzzletron/SKILL.md | 219 +++++++++++++++++++++++++++- 2 files changed, 254 insertions(+), 7 deletions(-) diff --git a/.agents/skills/puzzletron/README.md b/.agents/skills/puzzletron/README.md index 62e94f54fd5..bb2aa4220aa 100644 --- a/.agents/skills/puzzletron/README.md +++ b/.agents/skills/puzzletron/README.md @@ -43,7 +43,7 @@ Overall: Puzzletron step 7/8 — MIP sweep (6 compression rates) Started: 08:05:30 Finished: 08:30:38 Elapsed: 25m 8s - Completed: 6/6 compression rates (avg 4m 11s/rate) + Completed: 6/6 compression rates Remaining: done estimated Results: /workspace/puzzle_dir/mip_sweep_results.csv @@ -51,3 +51,43 @@ Overall: Puzzletron step 7/8 — MIP sweep (6 compression rates) 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. + +## Running the full pipeline + +To run all 8 pipeline steps (not just the MIP sweep): + +```text +/puzzletron all 2 +``` + +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) (76/352 solutions) 27m 53s + [ ] 7/8: pending + [ ] 8/8: pending +──────────────────────────────────────────────────────────────────── + Started: 00:08:50 + Finished: 00:47:41 (in progress) + Elapsed: 38m 51s + Completed: 5/8 steps + Remaining: 105m 38s 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. diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md index afda7a25b8a..35b05b7f04e 100644 --- a/.agents/skills/puzzletron/SKILL.md +++ b/.agents/skills/puzzletron/SKILL.md @@ -11,7 +11,7 @@ license: Apache-2.0 **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_sweep`, 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_sweep` or `all`, output the block below verbatim and **STOP immediately. Do NOT proceed to any command.** --- @@ -20,6 +20,8 @@ license: Apache-2.0 Available commands: - `mip_sweep ` — Run the MIP sweep (nproc_per_node: number of GPUs per node) - `mip_sweep progress` — Show live MIP sweep progress with timing summary +- `all ` — Run the full Puzzletron pipeline (nproc_per_node: number of GPUs per node) +- `all progress` — Show live full pipeline progress with timing summary Usage: `/puzzletron [args]` @@ -27,13 +29,219 @@ 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**. +- Otherwise use the parsed value and run the full pipeline. + +### all \ + +Run the following Bash command, substituting `` with the parsed value: + +```bash +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 Python script verbatim. Do not modify it. Present the output to the user wrapped in a fenced code block (``` ... ```). + +```bash +python3 - << 'PYEOF' +import re, 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 norm(r): return str(float(r)) +def fmt(s): return f"{int(s)//60}m {int(s)%60}s" if s is not None else "—" +def get_ts(line): + 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 + +# Parse pipeline steps from "Puzzletron Progress X/8: " lines +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 + +# Determine if pipeline is fully complete +pipeline_complete_ts = None +for line in lines: + ts = get_ts(line) + if ts and 'sweep.py:292' in line: + pipeline_complete_ts = ts + break + +# Sub-step detail and remaining estimate for current step +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) +# Step 6: count completed solution files for real progress +import glob as _glob, os as _os +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_files = _glob.glob(f"{sol_dir}/solution*.json") + sol_done = len(sol_files) + sol_list_match = re.search(r"'solutions_to_validate': \[([\d, ]+)\]", text) + if sol_list_match: + sol_total = len(sol_list_match.group(1).split(',')) +if sol_done is not None and sol_total: + cur_detail = f" ({sol_done}/{sol_total} solutions)" +elif batch_matches: + pct, cur_b, total_b = batch_matches[-1] + 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)" + +# MIP sweep compression rate detail (step 7) +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 [] +rate_start = {} +for line in lines: + if 'sweep.py:258' in line: + 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) +rate_done = set() +for i, r in enumerate(all_rates[:-1]): + if all_rates[i + 1] in rate_start: + rate_done.add(r) +if pipeline_complete_ts and all_rates and all_rates[-1] in rate_start: + rate_done.add(all_rates[-1]) + +# Overall timing +pipeline_start = step_events[0][3] if step_events else None +end_ts = pipeline_complete_ts if pipeline_complete_ts else now +total_elapsed = int((end_ts - pipeline_start).total_seconds()) if pipeline_start else 0 + +# Estimate remaining time for current running step +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 batch_matches 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)) + elif all_rates and last_step_num == 7: + done_count_r = len(rate_done) + remaining_count_r = len(all_rates) - done_count_r + rate_elapsed = {} + for i, r in enumerate(all_rates): + if r not in rate_start: + continue + if i + 1 < len(all_rates) and all_rates[i+1] in rate_start: + rate_elapsed[r] = int((rate_start[all_rates[i+1]] - rate_start[r]).total_seconds()) + avg_r = sum(rate_elapsed.values()) / len(rate_elapsed) if rate_elapsed else None + if avg_r and remaining_count_r: + step_remaining = avg_r * remaining_count_r + +print(f"\nOverall: Puzzletron full pipeline (steps 1–{total_steps})") +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 if pipeline_complete_ts else 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 + if snum == 7 and all_rates: + detail = f" ({len(rate_done)}/{len(all_rates)} rates done)" + 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}") + +for snum in range(last_step_num + 1, total_steps + 1): + print(f" {'[ ]':<10} {'':<4} {f'{snum}/{total_steps}: pending':<34} {'':>8}") + +print(DIV) +if all_rates and last_step_num >= 7: + print(f" MIP rates: {len(rate_done)}/{len(all_rates)} done", end="") + running_rate = next((r for r in all_rates if r in rate_start and r not in rate_done), None) + if running_rate: + print(f" (running: {running_rate})", end="") + print() +done_steps = len([s for s in seen_steps if s != last_step_num or pipeline_complete_ts]) +avg_step_s = None +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 if pipeline_complete_ts else None) + if next_ts and sts: + step_durations.append(int((next_ts - sts).total_seconds())) +if step_durations: + avg_step_s = sum(step_durations) / len(step_durations) +remaining_steps_after = total_steps - last_step_num # steps not yet started after current +if pipeline_complete_ts: + est_rem = "done" +elif step_remaining is not None: + future_est = (avg_step_s * remaining_steps_after) if avg_step_s else 0 + est_rem = fmt(step_remaining + future_est) +elif avg_step_s: + est_rem = fmt(avg_step_s * (remaining_steps_after + 1)) +else: + est_rem = "calculating..." + +print(f" Started: {pipeline_start.strftime('%H:%M:%S') if pipeline_start else '—'}") +print(f" Finished: {pipeline_complete_ts.strftime('%H:%M:%S') if pipeline_complete_ts else now.strftime('%H:%M:%S') + ' (in progress)'}") +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 results_match: + print(f"\n Results: {results_match.group(1)}") +PYEOF +``` + ## Command: mip_sweep -Parse the second word of args. +Parse `nproc_per_node` from args using either positional or flag syntax: +- Positional: second word is a number, e.g. `mip_sweep 2` +- Flag: `--nproc_per_node ` anywhere in args, e.g. `mip_sweep --nproc_per_node 2` -- If no second word is provided, ask the user: "Please provide the number of GPUs per node (nproc_per_node)." and **STOP**. - If the second word is exactly `progress`, execute the **mip_sweep progress** sub-command below. -- Otherwise treat the second word as `nproc_per_node` and run the sweep. +- 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**. +- Otherwise use the parsed value and run the sweep. ### mip_sweep \ @@ -155,8 +363,7 @@ print(DIV) print(f" Started: {sweep_start.strftime('%H:%M:%S') if sweep_start else '—'}") print(f" Finished: {sweep_complete_ts.strftime('%H:%M:%S') if sweep_complete_ts else now.strftime('%H:%M:%S') + ' (in progress)'}") print(f" Elapsed: {fmt(total_elapsed)}") -print(f" Completed: {done_count}/{len(all_rates)} compression rates", end="") -print(f" (avg {fmt(avg_s)}/rate)" if avg_s else "") +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: From 1a62799925c522642eca2a375737dd8c393dfab8 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Thu, 18 Jun 2026 01:11:11 -0700 Subject: [PATCH 07/28] update progress bar Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/SKILL.md | 37 +++++++++++++++++++++++++----- 1 file changed, 31 insertions(+), 6 deletions(-) diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md index 35b05b7f04e..24a647294d1 100644 --- a/.agents/skills/puzzletron/SKILL.md +++ b/.agents/skills/puzzletron/SKILL.md @@ -211,16 +211,41 @@ for i, (snum, (sdesc, sts)) in enumerate(step_ts_list): step_durations.append(int((next_ts - sts).total_seconds())) if step_durations: avg_step_s = sum(step_durations) / len(step_durations) -remaining_steps_after = total_steps - last_step_num # steps not yet started after current + +# Parse config for sweep settings to estimate step 7 duration +CONFIG_PATH = 'examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml' +sweep_enabled = True +sweep_n_rates = 6 +try: + cfg_text = open(CONFIG_PATH).read() + _en_m = re.search(r'sweep:\s*\n\s+enabled:\s*(true|false)', cfg_text) + if _en_m: + sweep_enabled = _en_m.group(1) == 'true' + _rates_m = re.search(r'memory_compression_rates:\s*\[([^\]]+)\]', cfg_text) + if _rates_m: + sweep_n_rates = len(_rates_m.group(1).split(',')) +except Exception: + pass +# Prefer actual rate count from log if step 7 has already started +effective_n_rates = len(all_rates) if all_rates else sweep_n_rates +RATE_S = 250 # ~4m 10s per compression rate (historical) + +def step_est(snum): + if snum == 7: + return (RATE_S * effective_n_rates) if sweep_enabled else 120 + 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_est = (avg_step_s * remaining_steps_after) if avg_step_s else 0 - est_rem = fmt(step_remaining + future_est) -elif avg_step_s: - est_rem = fmt(avg_step_s * (remaining_steps_after + 1)) + 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: - est_rem = "calculating..." + 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..." print(f" Started: {pipeline_start.strftime('%H:%M:%S') if pipeline_start else '—'}") print(f" Finished: {pipeline_complete_ts.strftime('%H:%M:%S') if pipeline_complete_ts else now.strftime('%H:%M:%S') + ' (in progress)'}") From fa197c2bb9809c908f67feb95c0f78aa702a05cb Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Thu, 18 Jun 2026 01:14:40 -0700 Subject: [PATCH 08/28] Rename mip_sweep to mip command Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/README.md | 8 ++++---- .agents/skills/puzzletron/SKILL.md | 22 +++++++++++----------- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/.agents/skills/puzzletron/README.md b/.agents/skills/puzzletron/README.md index bb2aa4220aa..e4baa46fcb3 100644 --- a/.agents/skills/puzzletron/README.md +++ b/.agents/skills/puzzletron/README.md @@ -10,19 +10,19 @@ For full environment setup, model configuration, and algorithm details see Run `/puzzletron` with no arguments to see available commands. -## Running the MIP sweep +## Running the MIP step -Start the sweep by telling the agent how many GPUs per node to use: +Start the MIP step by telling the agent how many GPUs per node to use: ```text -/puzzletron mip_sweep 4 +/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_sweep progress +/puzzletron mip progress ``` Example output when complete: diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md index 24a647294d1..3cdfb7a2fa7 100644 --- a/.agents/skills/puzzletron/SKILL.md +++ b/.agents/skills/puzzletron/SKILL.md @@ -11,15 +11,15 @@ license: Apache-2.0 **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_sweep` or `all`, 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` or `all`, 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_sweep ` — Run the MIP sweep (nproc_per_node: number of GPUs per node) -- `mip_sweep progress` — Show live MIP sweep progress with timing summary +- `mip ` — Run the MIP step (nproc_per_node: number of GPUs per node) +- `mip progress` — Show live MIP progress with timing summary - `all ` — Run the full Puzzletron pipeline (nproc_per_node: number of GPUs per node) - `all progress` — Show live full pipeline progress with timing summary @@ -258,17 +258,17 @@ if results_match: PYEOF ``` -## Command: mip_sweep +## Command: mip Parse `nproc_per_node` from args using either positional or flag syntax: -- Positional: second word is a number, e.g. `mip_sweep 2` -- Flag: `--nproc_per_node ` anywhere in args, e.g. `mip_sweep --nproc_per_node 2` +- 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_sweep progress** sub-command below. +- If the second word is exactly `progress`, execute the **mip 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**. -- Otherwise use the parsed value and run the sweep. +- Otherwise use the parsed value and run the MIP step. -### mip_sweep \ +### mip \ Run the following Bash command, substituting `` with the parsed value: @@ -281,7 +281,7 @@ torchrun --nproc_per_node examples/puzzletron/main.py \ Stream output to the user as it arrives. When the command finishes, report the exit code. -### mip_sweep progress +### mip progress Run the following Python script verbatim. Do not modify it. Present the output to the user wrapped in a fenced code block (``` ... ```). @@ -295,7 +295,7 @@ try: lines = open(LOG).readlines() text = ''.join(lines) except FileNotFoundError: - print("No log.txt found. Run /puzzletron mip_sweep first.") + print("No log.txt found. Run /puzzletron mip first.") sys.exit(0) def norm(r): return str(float(r)) From 03ceee301528ead4ceeecb54b5b8790ff9eb5093 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Thu, 18 Jun 2026 01:57:30 -0700 Subject: [PATCH 09/28] Move python scripts for puzzletron claude command to py scripts. Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/SKILL.md | 315 +--------------------- .agents/skills/puzzletron/all_progress.py | 237 ++++++++++++++++ .agents/skills/puzzletron/mip_progress.py | 141 ++++++++++ 3 files changed, 383 insertions(+), 310 deletions(-) create mode 100644 .agents/skills/puzzletron/all_progress.py create mode 100644 .agents/skills/puzzletron/mip_progress.py diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md index 3cdfb7a2fa7..a3d51ab799b 100644 --- a/.agents/skills/puzzletron/SKILL.md +++ b/.agents/skills/puzzletron/SKILL.md @@ -1,6 +1,6 @@ --- name: puzzletron -description: "End-to-end workflow for model pruning and MIP-based optimization. Use `mip_sweep` to run the MIP sweep. Usage: /puzzletron " +description: "End-to-end workflow for model pruning and MIP-based optimization. Commands: mip, all. Usage: /puzzletron " license: Apache-2.0 --- @@ -54,208 +54,10 @@ Stream output to the user as it arrives. When the command finishes, report the e ### all progress -Run the following Python script verbatim. Do not modify it. Present the output to the user wrapped in a fenced code block (``` ... ```). +Run the following Bash command. Present the output to the user wrapped in a fenced code block (``` ... ```). ```bash -python3 - << 'PYEOF' -import re, 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 norm(r): return str(float(r)) -def fmt(s): return f"{int(s)//60}m {int(s)%60}s" if s is not None else "—" -def get_ts(line): - 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 - -# Parse pipeline steps from "Puzzletron Progress X/8: " lines -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 - -# Determine if pipeline is fully complete -pipeline_complete_ts = None -for line in lines: - ts = get_ts(line) - if ts and 'sweep.py:292' in line: - pipeline_complete_ts = ts - break - -# Sub-step detail and remaining estimate for current step -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) -# Step 6: count completed solution files for real progress -import glob as _glob, os as _os -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_files = _glob.glob(f"{sol_dir}/solution*.json") - sol_done = len(sol_files) - sol_list_match = re.search(r"'solutions_to_validate': \[([\d, ]+)\]", text) - if sol_list_match: - sol_total = len(sol_list_match.group(1).split(',')) -if sol_done is not None and sol_total: - cur_detail = f" ({sol_done}/{sol_total} solutions)" -elif batch_matches: - pct, cur_b, total_b = batch_matches[-1] - 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)" - -# MIP sweep compression rate detail (step 7) -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 [] -rate_start = {} -for line in lines: - if 'sweep.py:258' in line: - 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) -rate_done = set() -for i, r in enumerate(all_rates[:-1]): - if all_rates[i + 1] in rate_start: - rate_done.add(r) -if pipeline_complete_ts and all_rates and all_rates[-1] in rate_start: - rate_done.add(all_rates[-1]) - -# Overall timing -pipeline_start = step_events[0][3] if step_events else None -end_ts = pipeline_complete_ts if pipeline_complete_ts else now -total_elapsed = int((end_ts - pipeline_start).total_seconds()) if pipeline_start else 0 - -# Estimate remaining time for current running step -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 batch_matches 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)) - elif all_rates and last_step_num == 7: - done_count_r = len(rate_done) - remaining_count_r = len(all_rates) - done_count_r - rate_elapsed = {} - for i, r in enumerate(all_rates): - if r not in rate_start: - continue - if i + 1 < len(all_rates) and all_rates[i+1] in rate_start: - rate_elapsed[r] = int((rate_start[all_rates[i+1]] - rate_start[r]).total_seconds()) - avg_r = sum(rate_elapsed.values()) / len(rate_elapsed) if rate_elapsed else None - if avg_r and remaining_count_r: - step_remaining = avg_r * remaining_count_r - -print(f"\nOverall: Puzzletron full pipeline (steps 1–{total_steps})") -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 if pipeline_complete_ts else 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 - if snum == 7 and all_rates: - detail = f" ({len(rate_done)}/{len(all_rates)} rates done)" - 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}") - -for snum in range(last_step_num + 1, total_steps + 1): - print(f" {'[ ]':<10} {'':<4} {f'{snum}/{total_steps}: pending':<34} {'':>8}") - -print(DIV) -if all_rates and last_step_num >= 7: - print(f" MIP rates: {len(rate_done)}/{len(all_rates)} done", end="") - running_rate = next((r for r in all_rates if r in rate_start and r not in rate_done), None) - if running_rate: - print(f" (running: {running_rate})", end="") - print() -done_steps = len([s for s in seen_steps if s != last_step_num or pipeline_complete_ts]) -avg_step_s = None -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 if pipeline_complete_ts else None) - if next_ts and sts: - step_durations.append(int((next_ts - sts).total_seconds())) -if step_durations: - avg_step_s = sum(step_durations) / len(step_durations) - -# Parse config for sweep settings to estimate step 7 duration -CONFIG_PATH = 'examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml' -sweep_enabled = True -sweep_n_rates = 6 -try: - cfg_text = open(CONFIG_PATH).read() - _en_m = re.search(r'sweep:\s*\n\s+enabled:\s*(true|false)', cfg_text) - if _en_m: - sweep_enabled = _en_m.group(1) == 'true' - _rates_m = re.search(r'memory_compression_rates:\s*\[([^\]]+)\]', cfg_text) - if _rates_m: - sweep_n_rates = len(_rates_m.group(1).split(',')) -except Exception: - pass -# Prefer actual rate count from log if step 7 has already started -effective_n_rates = len(all_rates) if all_rates else sweep_n_rates -RATE_S = 250 # ~4m 10s per compression rate (historical) - -def step_est(snum): - if snum == 7: - return (RATE_S * effective_n_rates) if sweep_enabled else 120 - 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..." - -print(f" Started: {pipeline_start.strftime('%H:%M:%S') if pipeline_start else '—'}") -print(f" Finished: {pipeline_complete_ts.strftime('%H:%M:%S') if pipeline_complete_ts else now.strftime('%H:%M:%S') + ' (in progress)'}") -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 results_match: - print(f"\n Results: {results_match.group(1)}") -PYEOF +python3 .agents/skills/puzzletron/all_progress.py ``` ## Command: mip @@ -283,115 +85,8 @@ Stream output to the user as it arrives. When the command finishes, report the e ### mip progress -Run the following Python script verbatim. Do not modify it. Present the output to the user wrapped in a fenced code block (``` ... ```). +Run the following Bash command. Present the output to the user wrapped in a fenced code block (``` ... ```). ```bash -python3 - << 'PYEOF' -import re, 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): return str(float(r)) -def fmt(s): return f"{int(s)//60}m {int(s)%60}s" if s is not None else "—" -def get_ts(line): - 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 - -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 [] - -# Collect start timestamp per rate -rate_start = {} -for line in lines: - if 'sweep.py:258' in line: - 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) - -now = datetime.now().replace(microsecond=0) -sweep_start = rate_start.get(all_rates[0]) if all_rates else None - -# Rate is done when the next rate has started; last rate done when sweep.py:287 appears -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] -sweep_complete_ts = None -for line in lines: - ts = get_ts(line) - if ts and 'sweep.py:292' in line: - sweep_complete_ts = ts - break -if sweep_complete_ts and last in rate_start: - rate_done.add(last) - -# Per-rate elapsed = next rate start - this rate start (or completion ts or now for last) -rate_elapsed = {} -for i, r in enumerate(all_rates): - if r not in rate_start: - continue - if i + 1 < len(all_rates): - end = rate_start[all_rates[i + 1]] - else: - end = sweep_complete_ts if sweep_complete_ts else now - rate_elapsed[r] = int((end - rate_start[r]).total_seconds()) - -# Currently running rate -running_rate = next((r for r in all_rates if r in rate_start and r not in rate_done), None) - -# Sub-step detail for running rate -cur_detail = "" -if running_rate: - batch_matches = re.findall(r'calculate_losses_pipeline[^:]*:\s*(\d+)%.*?(\d+)/(\d+)', text) - cbc_matches = re.findall(r'After (\d+) nodes.*?\(([\d.]+) seconds\)', 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 = sweep_complete_ts if sweep_complete_ts else 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] {'Prep (teacher memory + rate list)':<32} {'<1s':>8}") -for r in all_rates: - if r not in rate_start: - print(f" [ ] {f'compression_rate={r}':<32} {'pending':>8}") - elif r == running_rate: - detail = cur_detail - print(f" [RUNNING] {f'compression_rate={r}{detail}':<32} {fmt(rate_elapsed.get(r)):>8}") - else: - print(f" [DONE] {f'compression_rate={r}':<32} {fmt(rate_elapsed.get(r)):>8}") -print(DIV) -print(f" Started: {sweep_start.strftime('%H:%M:%S') if sweep_start else '—'}") -print(f" Finished: {sweep_complete_ts.strftime('%H:%M:%S') if sweep_complete_ts else now.strftime('%H:%M:%S') + ' (in progress)'}") -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)}") -PYEOF +python3 .agents/skills/puzzletron/mip_progress.py ``` diff --git a/.agents/skills/puzzletron/all_progress.py b/.agents/skills/puzzletron/all_progress.py new file mode 100644 index 00000000000..9224439df4b --- /dev/null +++ b/.agents/skills/puzzletron/all_progress.py @@ -0,0 +1,237 @@ +# 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 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) +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 +for line in lines: + ts = get_ts(line) + if ts and "sweep.py:292" in line: + pipeline_complete_ts = ts + break + +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(",")) +if sol_done is not None and sol_total: + cur_detail = f" ({sol_done}/{sol_total} solutions)" +elif batch_matches: + pct, cur_b, total_b = batch_matches[-1] + 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)" + +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 [] +rate_start = {} +for line in lines: + if "sweep.py:258" in line: + 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) +rate_done = set() +for i, r in enumerate(all_rates[:-1]): + if all_rates[i + 1] in rate_start: + rate_done.add(r) +if pipeline_complete_ts and all_rates and all_rates[-1] in rate_start: + rate_done.add(all_rates[-1]) + +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 batch_matches 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)) + elif all_rates and last_step_num == 7: + done_count_r = len(rate_done) + remaining_count_r = len(all_rates) - done_count_r + rate_elapsed = {} + for i, r in enumerate(all_rates): + if r not in rate_start: + continue + if i + 1 < len(all_rates) and all_rates[i + 1] in rate_start: + rate_elapsed[r] = int( + (rate_start[all_rates[i + 1]] - rate_start[r]).total_seconds() + ) + avg_r = sum(rate_elapsed.values()) / len(rate_elapsed) if rate_elapsed else None + if avg_r and remaining_count_r: + step_remaining = avg_r * remaining_count_r + +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 + if snum == 7 and all_rates: + detail = f" ({len(rate_done)}/{len(all_rates)} rates done)" + 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}" + ) + +for snum in range(last_step_num + 1, total_steps + 1): + print(f" {'[ ]':<10} {'':<4} {f'{snum}/{total_steps}: pending':<34} {'':>8}") + +print(DIV) +if all_rates and last_step_num >= 7: + print(f" MIP rates: {len(rate_done)}/{len(all_rates)} done", end="") + running_rate = next((r for r in all_rates if r in rate_start and r not in rate_done), None) + if running_rate: + print(f" (running: {running_rate})", end="") + print() + +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 + +CONFIG_PATH = ( + "examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml" +) +sweep_enabled = True +sweep_n_rates = 6 +try: + cfg_text = open(CONFIG_PATH).read() + _en_m = re.search(r"sweep:\s*\n\s+enabled:\s*(true|false)", cfg_text) + if _en_m: + sweep_enabled = _en_m.group(1) == "true" + _rates_m = re.search(r"memory_compression_rates:\s*\[([^\]]+)\]", cfg_text) + if _rates_m: + sweep_n_rates = len(_rates_m.group(1).split(",")) +except Exception: + pass +effective_n_rates = len(all_rates) if all_rates else sweep_n_rates +RATE_S = 250 # ~4m 10s per compression rate (historical) + + +def step_est(snum): + """Estimate duration in seconds for a pending pipeline step.""" + if snum == 7: + return (RATE_S * effective_n_rates) if sweep_enabled else 120 + 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 results_match: + print(f"\n Results: {results_match.group(1)}") diff --git a/.agents/skills/puzzletron/mip_progress.py b/.agents/skills/puzzletron/mip_progress.py new file mode 100644 index 00000000000..90e1dbce814 --- /dev/null +++ b/.agents/skills/puzzletron/mip_progress.py @@ -0,0 +1,141 @@ +# 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 + + +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 [] + +rate_start = {} +for line in lines: + if "sweep.py:258" in line: + 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) + +now = datetime.now().replace(microsecond=0) +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 +sweep_complete_ts = None +for line in lines: + ts = get_ts(line) + if ts and "sweep.py:292" in line: + sweep_complete_ts = ts + break +if sweep_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 + if i + 1 < len(all_rates): + end = rate_start[all_rates[i + 1]] + else: + end = sweep_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) + cbc_matches = re.findall(r"After (\d+) nodes.*?\(([\d.]+) seconds\)", 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 = sweep_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] {'Prep (teacher memory + rate list)':<32} {'<1s':>8}") +for r in all_rates: + if r not in rate_start: + print(f" [ ] {f'compression_rate={r}':<32} {'pending':>8}") + elif r == running_rate: + print( + f" [RUNNING] {f'compression_rate={r}{cur_detail}':<32} {fmt(rate_elapsed.get(r)):>8}" + ) + else: + print(f" [DONE] {f'compression_rate={r}':<32} {fmt(rate_elapsed.get(r)):>8}") +print(DIV) +finished_str = ( + sweep_complete_ts.strftime("%H:%M:%S") + if sweep_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)}") From 02ec52eeb5dd7a403323e1177f405d84ab3d5ef9 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Thu, 18 Jun 2026 02:00:44 -0700 Subject: [PATCH 10/28] update progress bar Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.agents/skills/puzzletron/README.md b/.agents/skills/puzzletron/README.md index e4baa46fcb3..9ee8ed9c6c0 100644 --- a/.agents/skills/puzzletron/README.md +++ b/.agents/skills/puzzletron/README.md @@ -78,15 +78,15 @@ Overall: Puzzletron full pipeline (steps 1–8) [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) (76/352 solutions) 27m 53s + [RUNNING] 6/8: calculating one block scores (multi-gpu) (270/352 solutions) 100m 6s [ ] 7/8: pending [ ] 8/8: pending ──────────────────────────────────────────────────────────────────── Started: 00:08:50 - Finished: 00:47:41 (in progress) - Elapsed: 38m 51s + Finished: 01:59:54 (in progress) + Elapsed: 111m 4s Completed: 5/8 steps - Remaining: 105m 38s estimated + Remaining: 56m 24s estimated ``` Step 6 progress is tracked via completed `solution_N.json` files on disk for an accurate From dd6fbdb60f72c7b19969b26194357dd1624500fd Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Thu, 18 Jun 2026 04:01:44 -0700 Subject: [PATCH 11/28] code clean up Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/all_progress.py | 75 ++-------------------- .agents/skills/puzzletron/mip_progress.py | 76 +++++++++++++++++------ 2 files changed, 64 insertions(+), 87 deletions(-) diff --git a/.agents/skills/puzzletron/all_progress.py b/.agents/skills/puzzletron/all_progress.py index 9224439df4b..f80cfb31054 100644 --- a/.agents/skills/puzzletron/all_progress.py +++ b/.agents/skills/puzzletron/all_progress.py @@ -30,11 +30,6 @@ 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 "—" @@ -64,11 +59,8 @@ def get_ts(line): last_step_num = max(seen_steps.keys()) if seen_steps else 0 pipeline_complete_ts = None -for line in lines: - ts = get_ts(line) - if ts and "sweep.py:292" in line: - pipeline_complete_ts = ts - break +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 @@ -94,23 +86,6 @@ def get_ts(line): nodes, secs = cbc_matches[-1] cur_detail = f" (MIP solver: {int(nodes):,} nodes, {float(secs):.1f}s)" -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 [] -rate_start = {} -for line in lines: - if "sweep.py:258" in line: - 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) -rate_done = set() -for i, r in enumerate(all_rates[:-1]): - if all_rates[i + 1] in rate_start: - rate_done.add(r) -if pipeline_complete_ts and all_rates and all_rates[-1] in rate_start: - rate_done.add(all_rates[-1]) - 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 @@ -125,20 +100,6 @@ def get_ts(line): elif batch_matches 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)) - elif all_rates and last_step_num == 7: - done_count_r = len(rate_done) - remaining_count_r = len(all_rates) - done_count_r - rate_elapsed = {} - for i, r in enumerate(all_rates): - if r not in rate_start: - continue - if i + 1 < len(all_rates) and all_rates[i + 1] in rate_start: - rate_elapsed[r] = int( - (rate_start[all_rates[i + 1]] - rate_start[r]).total_seconds() - ) - avg_r = sum(rate_elapsed.values()) / len(rate_elapsed) if rate_elapsed else None - if avg_r and remaining_count_r: - step_remaining = avg_r * remaining_count_r print(f"\nOverall: Puzzletron full pipeline (steps 1–{total_steps})") # noqa: RUF001 print(DIV) @@ -155,8 +116,6 @@ def get_ts(line): detail = "" if is_last and not is_done: detail = cur_detail - if snum == 7 and all_rates: - detail = f" ({len(rate_done)}/{len(all_rates)} rates done)" label = f"{snum}/{total_steps}: {sdesc}{detail}" status = "[DONE]" if is_done else "[RUNNING]" print( @@ -167,13 +126,6 @@ def get_ts(line): print(f" {'[ ]':<10} {'':<4} {f'{snum}/{total_steps}: pending':<34} {'':>8}") print(DIV) -if all_rates and last_step_num >= 7: - print(f" MIP rates: {len(rate_done)}/{len(all_rates)} done", end="") - running_rate = next((r for r in all_rates if r in rate_start and r not in rate_done), None) - if running_rate: - print(f" (running: {running_rate})", end="") - print() - 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): @@ -184,29 +136,12 @@ def get_ts(line): step_durations.append(int((next_ts - sts).total_seconds())) avg_step_s = sum(step_durations) / len(step_durations) if step_durations else None -CONFIG_PATH = ( - "examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml" -) -sweep_enabled = True -sweep_n_rates = 6 -try: - cfg_text = open(CONFIG_PATH).read() - _en_m = re.search(r"sweep:\s*\n\s+enabled:\s*(true|false)", cfg_text) - if _en_m: - sweep_enabled = _en_m.group(1) == "true" - _rates_m = re.search(r"memory_compression_rates:\s*\[([^\]]+)\]", cfg_text) - if _rates_m: - sweep_n_rates = len(_rates_m.group(1).split(",")) -except Exception: - pass -effective_n_rates = len(all_rates) if all_rates else sweep_n_rates -RATE_S = 250 # ~4m 10s per compression rate (historical) - def step_est(snum): """Estimate duration in seconds for a pending pipeline step.""" if snum == 7: - return (RATE_S * effective_n_rates) if sweep_enabled else 120 + # 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 @@ -233,5 +168,7 @@ def step_est(snum): 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_progress.py b/.agents/skills/puzzletron/mip_progress.py index 90e1dbce814..a52f1120b13 100644 --- a/.agents/skills/puzzletron/mip_progress.py +++ b/.agents/skills/puzzletron/mip_progress.py @@ -45,9 +45,63 @@ def get_ts(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 ("sweep.py:292" 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 results_match: + print(f"\n Results: {results_match.group(1)}") + sys.exit(0) + +# ── Sweep enabled: per-rate progress ───────────────────────────────────────── rate_start = {} for line in lines: if "sweep.py:258" in line: @@ -57,7 +111,6 @@ def get_ts(line): if r in all_rates and r not in rate_start: rate_start[r] = get_ts(line) -now = datetime.now().replace(microsecond=0) sweep_start = rate_start.get(all_rates[0]) if all_rates else None rate_done = set() @@ -65,23 +118,14 @@ def get_ts(line): if all_rates[i + 1] in rate_start: rate_done.add(r) last = all_rates[-1] if all_rates else None -sweep_complete_ts = None -for line in lines: - ts = get_ts(line) - if ts and "sweep.py:292" in line: - sweep_complete_ts = ts - break -if sweep_complete_ts and last and last in rate_start: +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 - if i + 1 < len(all_rates): - end = rate_start[all_rates[i + 1]] - else: - end = sweep_complete_ts or now + end = rate_start[all_rates[i + 1]] if i + 1 < len(all_rates) 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) @@ -89,7 +133,6 @@ def get_ts(line): cur_detail = "" if running_rate: batch_matches = re.findall(r"calculate_losses_pipeline[^:]*:\s*(\d+)%.*?(\d+)/(\d+)", text) - cbc_matches = re.findall(r"After (\d+) nodes.*?\(([\d.]+) seconds\)", text) if batch_matches: pct, cur, total = batch_matches[-1] cur_detail = f" — validating ({cur}/{total} batches)" @@ -97,7 +140,7 @@ def get_ts(line): nodes, secs = cbc_matches[-1] cur_detail = f" — MIP solver ({int(nodes):,} nodes, {float(secs):.1f}s)" -end_ts = sweep_complete_ts or now +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) @@ -110,7 +153,6 @@ def get_ts(line): ) 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}") @@ -127,9 +169,7 @@ def get_ts(line): print(f" [DONE] {f'compression_rate={r}':<32} {fmt(rate_elapsed.get(r)):>8}") print(DIV) finished_str = ( - sweep_complete_ts.strftime("%H:%M:%S") - if sweep_complete_ts - else now.strftime("%H:%M:%S") + " (in progress)" + 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}") From c9f1fe1dccf9a4e69156ab3591f74464bccf2d96 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Thu, 18 Jun 2026 04:05:26 -0700 Subject: [PATCH 12/28] code clean up Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/mip_progress.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/.agents/skills/puzzletron/mip_progress.py b/.agents/skills/puzzletron/mip_progress.py index a52f1120b13..d5d67eab039 100644 --- a/.agents/skills/puzzletron/mip_progress.py +++ b/.agents/skills/puzzletron/mip_progress.py @@ -97,6 +97,8 @@ def get_ts(line): 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) @@ -157,16 +159,16 @@ def get_ts(line): print(DIV) print(f" {'Status':<10} {'Phase':<32} {'Elapsed':>8}") print(DIV) -print(f" [DONE] {'Prep (teacher memory + rate list)':<32} {'<1s':>8}") +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" [ ] {f'compression_rate={r}':<32} {'pending':>8}") + print(f" {'[ ]':<10} {f'compression_rate={r}':<32} {'pending':>8}") elif r == running_rate: print( - f" [RUNNING] {f'compression_rate={r}{cur_detail}':<32} {fmt(rate_elapsed.get(r)):>8}" + f" {'[RUNNING]':<10} {f'compression_rate={r}{cur_detail}':<32} {fmt(rate_elapsed.get(r)):>8}" ) else: - print(f" [DONE] {f'compression_rate={r}':<32} {fmt(rate_elapsed.get(r)):>8}") + 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)" From afb6a712f3f9efd006eb5df1867feba8c36198bf Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Thu, 18 Jun 2026 04:17:07 -0700 Subject: [PATCH 13/28] add changelog for puzzletron clade skill Signed-off-by: Daniel Korzekwa --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d3d0ec160ec..852a352d891 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 ``mip`` and ``all`` commands for running the MIP step or full pipeline, and ``mip progress`` / ``all progress`` sub-commands reporting per-step status, elapsed time, and estimated time remaining. 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``. From 37d7132dff009f234cf1e7e81fae68f6f8e05b0b Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Fri, 19 Jun 2026 00:21:38 -0700 Subject: [PATCH 14/28] NameError: unpack batch_matches before if/elif so cur_b and total_b are always defined Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/all_progress.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.agents/skills/puzzletron/all_progress.py b/.agents/skills/puzzletron/all_progress.py index f80cfb31054..3db4b033e6b 100644 --- a/.agents/skills/puzzletron/all_progress.py +++ b/.agents/skills/puzzletron/all_progress.py @@ -77,10 +77,10 @@ def get_ts(line): 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: - pct, cur_b, total_b = batch_matches[-1] cur_detail = f" ({cur_b}/{total_b} batches)" elif cbc_matches: nodes, secs = cbc_matches[-1] @@ -97,7 +97,7 @@ def get_ts(line): 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 batch_matches and int(cur_b) > 0 and int(cur_b) < int(total_b): + 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)) From 417ae874cb442cd5f59a3a3bb73cb0b88d0a582e Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Fri, 19 Jun 2026 00:23:12 -0700 Subject: [PATCH 15/28] dd nproc_per_node integer validation to prevent shell injection in all and mip commands Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/SKILL.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md index a3d51ab799b..2b5c5bb1e94 100644 --- a/.agents/skills/puzzletron/SKILL.md +++ b/.agents/skills/puzzletron/SKILL.md @@ -37,6 +37,7 @@ Parse `nproc_per_node` from args using either positional or flag syntax: - 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 \ @@ -68,6 +69,7 @@ Parse `nproc_per_node` from args using either positional or flag syntax: - If the second word is exactly `progress`, execute the **mip 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 MIP step. ### mip \ From 708cf7b37341f4249fced362e35b9124f4de87bc Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Fri, 19 Jun 2026 00:25:22 -0700 Subject: [PATCH 16/28] replace hardcoded sweep.py line-number markers with content-based detection in mip_progress.py Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/mip_progress.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/.agents/skills/puzzletron/mip_progress.py b/.agents/skills/puzzletron/mip_progress.py index d5d67eab039..2065cc18ce3 100644 --- a/.agents/skills/puzzletron/mip_progress.py +++ b/.agents/skills/puzzletron/mip_progress.py @@ -54,7 +54,7 @@ def get_ts(line): complete_ts = None for line in lines: ts = get_ts(line) - if ts and ("sweep.py:292" in line or "Puzzletron Progress 8/8" in line): + if ts and ("Results written to:" in line or "Puzzletron Progress 8/8" in line): complete_ts = ts break @@ -106,12 +106,11 @@ def get_ts(line): # ── Sweep enabled: per-rate progress ───────────────────────────────────────── rate_start = {} for line in lines: - if "sweep.py:258" in line: - 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) + 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 From ca9d8b4009d3df0ff06059de46fef64f398cb92d Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Fri, 19 Jun 2026 00:26:40 -0700 Subject: [PATCH 17/28] add set -o pipefail to torchrun pipelines so torchrun failures are not masked by grep exit code Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/SKILL.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md index 2b5c5bb1e94..803cd5bf816 100644 --- a/.agents/skills/puzzletron/SKILL.md +++ b/.agents/skills/puzzletron/SKILL.md @@ -45,7 +45,7 @@ Parse `nproc_per_node` from args using either positional or flag syntax: Run the following Bash command, substituting `` with the parsed value: ```bash -export PYTHONPATH=$PYTHONPATH:/workspace/Model-Optimizer && \ +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" @@ -77,7 +77,7 @@ Parse `nproc_per_node` from args using either positional or flag syntax: Run the following Bash command, substituting `` with the parsed value: ```bash -export PYTHONPATH=$PYTHONPATH:/workspace/Model-Optimizer && \ +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" From a572aa58ec172a3a6b0a0f0cb4cd0739d001e2e0 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Fri, 19 Jun 2026 01:42:12 -0700 Subject: [PATCH 18/28] Add qwen 3.5 descriptor Signed-off-by: Daniel Korzekwa --- .../pruning/attn_pruning.yaml | 15 ++ .../pruning/ffn_pruning.yaml | 20 ++ .../pruning/hidden_dim_pruning.yaml | 15 ++ .../pruning/pruning_defaults.yaml | 33 ++++ .../qwen3_5-2B_pruneffn_memory.yaml | 22 +++ .../qwen3_5-2B_pruneffn_memory/qwen3_5.yaml | 106 +++++++++++ .../validate_model_defaults.yaml | 17 ++ .../validate_solutions_defaults.yaml | 10 + .../puzzletron/anymodel/models/__init__.py | 1 + .../anymodel/models/qwen3_5/__init__.py | 17 ++ .../models/qwen3_5/qwen3_5_converter.py | 44 +++++ .../qwen3_5/qwen3_5_model_descriptor.py | 174 ++++++++++++++++++ 12 files changed, 474 insertions(+) create mode 100644 examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/attn_pruning.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/ffn_pruning.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/hidden_dim_pruning.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/pruning/pruning_defaults.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/qwen3_5-2B_pruneffn_memory.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/qwen3_5.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/validate_model_defaults.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/validate_solutions_defaults.yaml create mode 100644 modelopt/torch/puzzletron/anymodel/models/qwen3_5/__init__.py create mode 100644 modelopt/torch/puzzletron/anymodel/models/qwen3_5/qwen3_5_converter.py create mode 100644 modelopt/torch/puzzletron/anymodel/models/qwen3_5/qwen3_5_model_descriptor.py 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..128eb5dbb0e --- /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 + +# 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..ce1749d9698 --- /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: valid +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..017f04dc300 --- /dev/null +++ b/modelopt/torch/puzzletron/anymodel/models/qwen3_5/qwen3_5_converter.py @@ -0,0 +1,44 @@ +# 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 + +from typing import List + +from ....block_config import AttentionConfig, BlockConfig, FFNConfig +from ...converter import Converter, ConverterFactory + +__all__ = ["Qwen3_5Converter"] + + +@ConverterFactory.register_decorator("qwen3_5") +class Qwen3_5Converter(Converter): + @staticmethod + def create_block_configs_from_main_config(config) -> List[BlockConfig]: + # Qwen3.5 is a VLM; text parameters live in the nested text_config. + text_config = config.text_config if hasattr(config, "text_config") else config + + num_hidden_layers = text_config.num_hidden_layers + + block_configs = [ + 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(num_hidden_layers) + ] + return block_configs 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..eb1ee7a74a3 --- /dev/null +++ b/modelopt/torch/puzzletron/anymodel/models/qwen3_5/qwen3_5_model_descriptor.py @@ -0,0 +1,174 @@ +# 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, List + +from torch import nn +from transformers.models.qwen3_5.modeling_qwen3_5 import ( + Qwen3_5DecoderLayer, + Qwen3_5ForCausalLM, + 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", +] + + +@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): + # model is Qwen3_5ForConditionalGeneration; text model is at model.model.language_model + text_config = Qwen3_5ModelDescriptor.get_language_model_config(model.config) + model.model.language_model.rotary_emb = Qwen3_5TextRotaryEmbedding(config=text_config).to( + device=runtime.device + ) + + @staticmethod + def input_embedding_name(): + return "model.language_model.embed_tokens" + + @staticmethod + def output_embedding_name(): + return "lm_head" + + @staticmethod + def final_norm_name(): + return "model.language_model.norm" + + @staticmethod + def layer_block_name(index: int): + return f"model.language_model.layers.{index}" + + @staticmethod + def layer_name_predicates(num_layers: int) -> Dict[str, re.Pattern]: + layer_name_patterns = { + "embeddings": re.compile(r"^model\.language_model\.embed_tokens\.weight$"), + "lm_head": re.compile(r"^(model\.language_model\.norm\.weight|lm_head\.weight)$"), + "vision_encoding": re.compile(r"^model\.visual\..*"), + } + + def build_ffn_predicates() -> Dict[str, re.Pattern]: + return { + f"block_{layer_idx}_ffn": re.compile( + rf"^model\.language_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\.language_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.language_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.language_model.layers.{layer_idx}.self_attn" + qkvo_weight_names: List[str] = field( + default_factory=lambda: ["q_proj", "k_proj", "v_proj", "o_proj"] + ) From 77c5c96a2a9bc9226ce93ed63167c050f5bea263 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Fri, 19 Jun 2026 02:45:53 -0700 Subject: [PATCH 19/28] Update puzzletron skill tutorial Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/README.md | 90 ++++++++------- .../puzzletron/adding_new_model_tutorial.md | 104 ++++++++++++++++++ 2 files changed, 155 insertions(+), 39 deletions(-) create mode 100644 .agents/skills/puzzletron/adding_new_model_tutorial.md diff --git a/.agents/skills/puzzletron/README.md b/.agents/skills/puzzletron/README.md index 9ee8ed9c6c0..e8042dfbcee 100644 --- a/.agents/skills/puzzletron/README.md +++ b/.agents/skills/puzzletron/README.md @@ -1,7 +1,7 @@ # 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. +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). @@ -10,6 +10,52 @@ For full environment setup, model configuration, and algorithm details see 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: pending + [ ] 8/8: pending +──────────────────────────────────────────────────────────────────── + 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: @@ -52,42 +98,8 @@ Overall: Puzzletron step 7/8 — MIP sweep (6 compression rates) 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. -## Running the full pipeline - -To run all 8 pipeline steps (not just the MIP sweep): - -```text -/puzzletron all 2 -``` - -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: pending - [ ] 8/8: pending -──────────────────────────────────────────────────────────────────── - Started: 00:08:50 - Finished: 01:59:54 (in progress) - Elapsed: 111m 4s - Completed: 5/8 steps - Remaining: 56m 24s estimated -``` +## Adding support for a new model -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. +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/adding_new_model_tutorial.md b/.agents/skills/puzzletron/adding_new_model_tutorial.md new file mode 100644 index 00000000000..abe72d6f5b5 --- /dev/null +++ b/.agents/skills/puzzletron/adding_new_model_tutorial.md @@ -0,0 +1,104 @@ +# Adding a New Model to Puzzletron with Claude + +This tutorial shows the full conversation between a user and Claude when adding +Qwen3.5-2B 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-2B +``` + +> **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-2B") +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-2B + +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-2B_pruneffn_memory/`): + +`intermediate_size_list` chosen by scaling the Qwen3-8B ratios (~21 %, 42 %, 60 %, 81 % +of teacher) to the 2B's `intermediate_size = 6144`: + +```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-2B on 4 GPUs + +Claude constructs the `torchrun` command directly with the Qwen3.5-2B config path and runs the full pipeline. The user monitors progress with: + +```text +/puzzletron all progress +``` From d76a8ebdb1c6c947e40ec2586889bd441d27bec8 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Fri, 19 Jun 2026 03:25:05 -0700 Subject: [PATCH 20/28] fix qwen 3.5 descriptor Signed-off-by: Daniel Korzekwa --- .../qwen3_5-2B_pruneffn_memory.yaml | 2 +- .../models/qwen3_5/qwen3_5_converter.py | 39 +++++++++++++++---- .../qwen3_5/qwen3_5_model_descriptor.py | 33 ++++++++++------ 3 files changed, 55 insertions(+), 19 deletions(-) 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 index 128eb5dbb0e..1dad6b29be8 100644 --- 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 @@ -9,7 +9,7 @@ input_hf_model_path: /workspace/hf_models/Qwen3.5-2B dataset_path: /workspace/datasets/Puzzle-KD-Nemotron-Post-Training-Dataset-v2 # Working directory for compression outputs -puzzle_dir: /workspace/puzzle_dir +puzzle_dir: /workspace/puzzle_dir_qwen3_5-2B # MIP memory constraint (in MiB) mip: 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 index 017f04dc300..b94a4cc770b 100644 --- a/modelopt/torch/puzzletron/anymodel/models/qwen3_5/qwen3_5_converter.py +++ b/modelopt/torch/puzzletron/anymodel/models/qwen3_5/qwen3_5_converter.py @@ -15,30 +15,55 @@ # 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]: - # Qwen3.5 is a VLM; text parameters live in the nested text_config. text_config = config.text_config if hasattr(config, "text_config") else config - - num_hidden_layers = text_config.num_hidden_layers - - block_configs = [ + 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(num_hidden_layers) + for _ in range(text_config.num_hidden_layers) ] - return block_configs + + @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 index eb1ee7a74a3..eb74b261135 100644 --- 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 @@ -17,12 +17,11 @@ import re from dataclasses import dataclass, field -from typing import Dict, List +from typing import Dict, Iterable, List from torch import nn from transformers.models.qwen3_5.modeling_qwen3_5 import ( Qwen3_5DecoderLayer, - Qwen3_5ForCausalLM, Qwen3_5TextRotaryEmbedding, ) @@ -39,6 +38,10 @@ "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): @@ -86,15 +89,14 @@ def mlp_no_op_post_init(decoder_layer: Qwen3_5DecoderLayer): @staticmethod def init_rotary_embedding(model, runtime): - # model is Qwen3_5ForConditionalGeneration; text model is at model.model.language_model - text_config = Qwen3_5ModelDescriptor.get_language_model_config(model.config) - model.model.language_model.rotary_emb = Qwen3_5TextRotaryEmbedding(config=text_config).to( + # 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.language_model.embed_tokens" + return "model.embed_tokens" @staticmethod def output_embedding_name(): @@ -102,18 +104,27 @@ def output_embedding_name(): @staticmethod def final_norm_name(): - return "model.language_model.norm" + return "model.norm" @staticmethod def layer_block_name(index: int): - return f"model.language_model.layers.{index}" + 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.""" + text_names = [n for n in layer_names if not n.startswith(_NON_TEXT_PREFIXES)] + return super().get_weight_groups(text_names, num_hidden_layers) @staticmethod def layer_name_predicates(num_layers: int) -> Dict[str, re.Pattern]: + # Predicates match ORIGINAL checkpoint names (model.language_model.*). + # convert_weight_name remaps them to model.* after grouping. layer_name_patterns = { "embeddings": re.compile(r"^model\.language_model\.embed_tokens\.weight$"), "lm_head": re.compile(r"^(model\.language_model\.norm\.weight|lm_head\.weight)$"), - "vision_encoding": re.compile(r"^model\.visual\..*"), } def build_ffn_predicates() -> Dict[str, re.Pattern]: @@ -159,7 +170,7 @@ def build_attention_predicates() -> Dict[str, re.Pattern]: @dataclass class Qwen3_5FFNIntermediateLayerDescriptor(FFNIntermediateLayerDescriptor): down_proj_name: str = "mlp.down_proj" - ffn_prefix_name: str = "model.language_model.layers.{layer_idx}.mlp" + ffn_prefix_name: str = "model.layers.{layer_idx}.mlp" linear_weight_names: List[str] = field( default_factory=lambda: ["down_proj", "gate_proj", "up_proj"] ) @@ -168,7 +179,7 @@ class Qwen3_5FFNIntermediateLayerDescriptor(FFNIntermediateLayerDescriptor): @dataclass class Qwen3_5KVHeadsLayerDescriptor(KVHeadsLayerDescriptor): o_proj_name: str = "self_attn.o_proj" - attn_prefix_name: str = "model.language_model.layers.{layer_idx}.self_attn" + 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"] ) From 27799f1051be839d8e3c65a72c75250e15245e99 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Fri, 19 Jun 2026 03:36:33 -0700 Subject: [PATCH 21/28] fix qwen 3.5 descriptor Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/README.md | 4 +- .agents/skills/puzzletron/all_progress.py | 14 ++++++- .../qwen3_5/qwen3_5_model_descriptor.py | 41 +++++++++++++++---- 3 files changed, 48 insertions(+), 11 deletions(-) diff --git a/.agents/skills/puzzletron/README.md b/.agents/skills/puzzletron/README.md index e8042dfbcee..899e492d503 100644 --- a/.agents/skills/puzzletron/README.md +++ b/.agents/skills/puzzletron/README.md @@ -43,8 +43,8 @@ Overall: Puzzletron full pipeline (steps 1–8) [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: pending - [ ] 8/8: pending + [ ] 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) diff --git a/.agents/skills/puzzletron/all_progress.py b/.agents/skills/puzzletron/all_progress.py index 3db4b033e6b..5204a0b56e5 100644 --- a/.agents/skills/puzzletron/all_progress.py +++ b/.agents/skills/puzzletron/all_progress.py @@ -122,8 +122,20 @@ def get_ts(line): 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): - print(f" {'[ ]':<10} {'':<4} {f'{snum}/{total_steps}: pending':<34} {'':>8}") + 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]) 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 index eb74b261135..9d31ac1d1a1 100644 --- 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 @@ -114,23 +114,48 @@ def layer_block_name(index: int): def get_weight_groups( cls, layer_names: Iterable[str], num_hidden_layers: int ) -> Dict[str, List[str]]: - """Filter out vision/MTP weights before grouping.""" + """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)] - return super().get_weight_groups(text_names, num_hidden_layers) + + 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 match ORIGINAL checkpoint names (model.language_model.*). - # convert_weight_name remaps them to model.* after grouping. + # 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\.language_model\.embed_tokens\.weight$"), - "lm_head": re.compile(r"^(model\.language_model\.norm\.weight|lm_head\.weight)$"), + "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\.language_model\.layers\.{layer_idx}\.(post_attention_layernorm\.weight" + rf"^model\.layers\.{layer_idx}\.(post_attention_layernorm\.weight" r"|mlp\.up_proj\.weight" r"|mlp\.gate_proj\.weight" r"|mlp\.down_proj\.weight)$" @@ -141,7 +166,7 @@ def build_ffn_predicates() -> Dict[str, re.Pattern]: def build_attention_predicates() -> Dict[str, re.Pattern]: return { f"block_{layer_idx}_attention": re.compile( - rf"^model\.language_model\.layers\.{layer_idx}\.(input_layernorm\.weight" + 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" From cd32126e9cddc8bed49c62f3b73e5b9476537d96 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Fri, 19 Jun 2026 03:41:47 -0700 Subject: [PATCH 22/28] add puzzletron add-model command Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/SKILL.md | 136 ++++++++++++++++++++++++++++- 1 file changed, 134 insertions(+), 2 deletions(-) diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md index 803cd5bf816..fb986e4f328 100644 --- a/.agents/skills/puzzletron/SKILL.md +++ b/.agents/skills/puzzletron/SKILL.md @@ -1,6 +1,6 @@ --- name: puzzletron -description: "End-to-end workflow for model pruning and MIP-based optimization. Commands: mip, all. Usage: /puzzletron " +description: "End-to-end workflow for model pruning and MIP-based optimization. Commands: mip, all, add-model. Usage: /puzzletron " license: Apache-2.0 --- @@ -11,7 +11,7 @@ license: Apache-2.0 **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` or `all`, 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.** --- @@ -22,6 +22,7 @@ Available commands: - `mip progress` — Show live MIP progress with timing summary - `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]` @@ -92,3 +93,134 @@ Run the following Bash command. Present the output to the user wrapped in a fenc ```bash python3 .agents/skills/puzzletron/mip_progress.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 +``` From 93eae9610eafa6accf477c4d7223918ea43071aa Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Fri, 19 Jun 2026 04:08:00 -0700 Subject: [PATCH 23/28] update qwen 3.5 descriptor Signed-off-by: Daniel Korzekwa --- .../qwen3_5-2B_pruneffn_memory/validate_model_defaults.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index ce1749d9698..6b36142a3a8 100644 --- 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 @@ -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 From c05a07814ec600f4e53c9efb80ccfeb995b54448 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Sun, 21 Jun 2026 10:25:29 -0700 Subject: [PATCH 24/28] puzzletron example for qwen 3.5 0.8B Signed-off-by: Daniel Korzekwa --- .../llama-3_1-8B_pruneffn_memory.yaml | 4 +- .../validate_model_defaults.yaml | 2 +- .../pruning/attn_pruning.yaml | 15 +++ .../pruning/ffn_pruning.yaml | 20 ++++ .../pruning/hidden_dim_pruning.yaml | 15 +++ .../pruning/pruning_defaults.yaml | 33 ++++++ .../qwen3_5-0.8b_pruneffn_memory.yaml | 22 ++++ .../qwen3_5-0.8b_pruneffn_memory/qwen3_5.yaml | 106 ++++++++++++++++++ .../validate_model_defaults.yaml | 17 +++ .../validate_solutions_defaults.yaml | 10 ++ 10 files changed, 241 insertions(+), 3 deletions(-) create mode 100644 examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/attn_pruning.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/ffn_pruning.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/hidden_dim_pruning.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/pruning/pruning_defaults.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/qwen3_5-0.8b_pruneffn_memory.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/qwen3_5.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/validate_model_defaults.yaml create mode 100644 examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/validate_solutions_defaults.yaml 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..b031515b39f --- /dev/null +++ b/examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/qwen3_5-0.8b_pruneffn_memory.yaml @@ -0,0 +1,22 @@ +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 + +# 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..25d9d054bac --- /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: 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-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 From b9cd64a8777a04ffea234526326bb3c629c00e44 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Sun, 21 Jun 2026 11:04:26 -0700 Subject: [PATCH 25/28] update adding new model tutorial to use qwen 3.5 0.8 instead of 2B Signed-off-by: Daniel Korzekwa --- .../puzzletron/adding_new_model_tutorial.md | 42 ++++++++++++++----- .../qwen3_5-0.8b_pruneffn_memory/qwen3_5.yaml | 4 +- 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/.agents/skills/puzzletron/adding_new_model_tutorial.md b/.agents/skills/puzzletron/adding_new_model_tutorial.md index abe72d6f5b5..98498eba2f9 100644 --- a/.agents/skills/puzzletron/adding_new_model_tutorial.md +++ b/.agents/skills/puzzletron/adding_new_model_tutorial.md @@ -1,7 +1,7 @@ # Adding a New Model to Puzzletron with Claude This tutorial shows the full conversation between a user and Claude when adding -Qwen3.5-2B support to Puzzletron. Use it as a guide for how to interact with Claude +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. --- @@ -12,12 +12,12 @@ when you want to compress a model that Puzzletron doesn't yet support. The user starts simply — download the model and run the pipeline. -> **User:** can you download to /workspace/hf_models/... +> **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-2B +Successfully downloaded 13 files to /workspace/hf_models/Qwen3.5-0.8B ``` > **User:** run puzzletron all for this model @@ -50,7 +50,7 @@ Claude confirms the model now loads: ```python from transformers import AutoConfig -config = AutoConfig.from_pretrained("/workspace/hf_models/Qwen3.5-2B") +config = AutoConfig.from_pretrained("/workspace/hf_models/Qwen3.5-0.8B") print(config.model_type) # qwen3_5 ✓ ``` @@ -58,7 +58,7 @@ 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-2B +> **User:** ok, implement the descriptor and compression example config for Qwen3.5-0.8B Claude implements everything in one go: @@ -68,10 +68,10 @@ Claude implements everything in one go: **Registration** (`anymodel/models/__init__.py`): gated behind `transformers >= 4.57.0` -**Compression example config** (`examples/puzzletron/configs/qwen3_5-2B_pruneffn_memory/`): +**Compression example config** (`examples/puzzletron/configs/qwen3_5-0.8b_pruneffn_memory/`): -`intermediate_size_list` chosen by scaling the Qwen3-8B ratios (~21 %, 42 %, 60 %, 81 % -of teacher) to the 2B's `intermediate_size = 6144`: +`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] @@ -95,10 +95,32 @@ print('qwen3_5' in ModelDescriptorFactory.CLASS_MAPPING) # True ✓ ### Step 4: Run the pipeline -> **User:** run puzzletron all for Qwen3.5-2B on 4 GPUs +> **User:** run puzzletron all for Qwen3.5-0.8B on 4 GPUs -Claude constructs the `torchrun` command directly with the Qwen3.5-2B config path and runs the full pipeline. The user monitors progress with: +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 +``` + +Steps 1–5 typically complete in under 2 minutes. 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). 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 index 25d9d054bac..2b49a48e420 100644 --- 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 @@ -42,7 +42,7 @@ scoring: teacher_dir: ${to_path:${teacher_dir}} output_dir: ${puzzle_dir}/single_sequence_replacement_solutions--validation - eval_samples: 128 + eval_samples: 8 micro_batch_size: 1 seed: 42 shuffle_seed: 444 @@ -93,7 +93,7 @@ realize_model: solutions_path: skip_validation: false - eval_samples: 128 + eval_samples: 8 micro_batch_size: 1 seed: 42 shuffle_seed: 444 From 6236b3914e0487fe1838bdbc86369879d1af2826 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Sun, 21 Jun 2026 23:59:40 -0700 Subject: [PATCH 26/28] Add puzzletron mip losses command Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/README.md | 17 +++++ .agents/skills/puzzletron/SKILL.md | 12 +++- .../puzzletron/adding_new_model_tutorial.md | 23 +++++- .agents/skills/puzzletron/mip_losses.py | 70 +++++++++++++++++++ 4 files changed, 120 insertions(+), 2 deletions(-) create mode 100644 .agents/skills/puzzletron/mip_losses.py diff --git a/.agents/skills/puzzletron/README.md b/.agents/skills/puzzletron/README.md index 899e492d503..e14f2de0bf1 100644 --- a/.agents/skills/puzzletron/README.md +++ b/.agents/skills/puzzletron/README.md @@ -98,6 +98,23 @@ Overall: Puzzletron step 7/8 — MIP sweep (6 compression rates) 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 + +After the pipeline completes, view the teacher vs. compressed model accuracy for the MIP solution: + +```text +/puzzletron mip losses +``` + +Example output: + +| Metric | Teacher | Compressed (solution_0) | +|---|---|---| +| `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 diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md index fb986e4f328..57e0020b637 100644 --- a/.agents/skills/puzzletron/SKILL.md +++ b/.agents/skills/puzzletron/SKILL.md @@ -1,6 +1,6 @@ --- name: puzzletron -description: "End-to-end workflow for model pruning and MIP-based optimization. Commands: mip, all, add-model. Usage: /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 --- @@ -20,6 +20,7 @@ license: Apache-2.0 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 MIP solution - `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 @@ -69,6 +70,7 @@ Parse `nproc_per_node` from args using either positional or flag syntax: - 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 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. @@ -94,6 +96,14 @@ Run the following Bash command. Present the output to the user wrapped in a fenc 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 +``` + ## 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**. diff --git a/.agents/skills/puzzletron/adding_new_model_tutorial.md b/.agents/skills/puzzletron/adding_new_model_tutorial.md index 98498eba2f9..0c34eea5696 100644 --- a/.agents/skills/puzzletron/adding_new_model_tutorial.md +++ b/.agents/skills/puzzletron/adding_new_model_tutorial.md @@ -123,4 +123,25 @@ Overall: Puzzletron full pipeline (steps 1–8) Elapsed: 21m 26s | Remaining: ~28m estimated ``` -Steps 1–5 typically complete in under 2 minutes. 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 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 + +Once the pipeline completes, check the accuracy of the MIP-selected compressed model against the teacher: + +```text +/puzzletron mip losses +``` + +Example output for Qwen3.5-0.8B: + +| Metric | Teacher | Compressed (solution_0) | +|---|---|---| +| `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 | + +The results are read from `/mip/puzzle_solutions//solutions--validation/solution_0.json` (and `teacher.json` in the same directory). diff --git a/.agents/skills/puzzletron/mip_losses.py b/.agents/skills/puzzletron/mip_losses.py new file mode 100644 index 00000000000..153b7cafd46 --- /dev/null +++ b/.agents/skills/puzzletron/mip_losses.py @@ -0,0 +1,70 @@ +# 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 MIP solution.""" + +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 that contains /mip/puzzle_solutions/ in the log +match = re.search(r"(\S+)/mip/puzzle_solutions/", text) +if not match: + # Fall back: look for the ckpts/teacher path logged early in the run + 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) + +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) + +solutions_dir = solutions_dirs[-1] + +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)) +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}") From 1d37a9a87694e1239d6753acda6c56d30dd9c920 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 22 Jun 2026 00:38:58 -0700 Subject: [PATCH 27/28] Add mip sweep losses to puzzletron skill Signed-off-by: Daniel Korzekwa --- .agents/skills/puzzletron/README.md | 13 ++- .agents/skills/puzzletron/SKILL.md | 12 ++- .../puzzletron/adding_new_model_tutorial.md | 82 ++++++++++++++++--- .agents/skills/puzzletron/mip_losses.py | 51 +++++++++++- .agents/skills/puzzletron/mip_progress.py | 3 +- .agents/skills/puzzletron/mip_sweep.py | 78 ++++++++++++++++++ .../qwen3_5-0.8b_pruneffn_memory.yaml | 5 ++ 7 files changed, 226 insertions(+), 18 deletions(-) create mode 100644 .agents/skills/puzzletron/mip_sweep.py diff --git a/.agents/skills/puzzletron/README.md b/.agents/skills/puzzletron/README.md index e14f2de0bf1..0820f44aee7 100644 --- a/.agents/skills/puzzletron/README.md +++ b/.agents/skills/puzzletron/README.md @@ -100,16 +100,25 @@ or validation batch progress), and an estimated time remaining based on complete ## Checking compressed model accuracy -After the pipeline completes, view the teacher vs. compressed model accuracy for the MIP solution: +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 ``` -Example output: +**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 | diff --git a/.agents/skills/puzzletron/SKILL.md b/.agents/skills/puzzletron/SKILL.md index 57e0020b637..a18398d8fd7 100644 --- a/.agents/skills/puzzletron/SKILL.md +++ b/.agents/skills/puzzletron/SKILL.md @@ -20,7 +20,8 @@ license: Apache-2.0 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 MIP solution +- `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 @@ -71,6 +72,7 @@ Parse `nproc_per_node` from args using either positional or flag syntax: - 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. @@ -104,6 +106,14 @@ Run the following Bash command. Present the output to the user wrapped in a fenc 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**. diff --git a/.agents/skills/puzzletron/adding_new_model_tutorial.md b/.agents/skills/puzzletron/adding_new_model_tutorial.md index 0c34eea5696..275380297a0 100644 --- a/.agents/skills/puzzletron/adding_new_model_tutorial.md +++ b/.agents/skills/puzzletron/adding_new_model_tutorial.md @@ -129,19 +129,81 @@ Step 6 (one-block scoring) is the longest step — it scores all candidate solut ### Step 5: Check the compressed model accuracy -Once the pipeline completes, check the accuracy of the MIP-selected compressed model against the teacher: +> **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 -/puzzletron mip losses +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 ``` -Example output for Qwen3.5-0.8B: +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. + +--- -| Metric | Teacher | Compressed (solution_0) | -|---|---|---| -| `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 | +### 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 +``` -The results are read from `/mip/puzzle_solutions//solutions--validation/solution_0.json` (and `teacher.json` in the same directory). +Use this table to pick the compression rate that best meets your accuracy/memory budget. diff --git a/.agents/skills/puzzletron/mip_losses.py b/.agents/skills/puzzletron/mip_losses.py index 153b7cafd46..00a3e4fd953 100644 --- a/.agents/skills/puzzletron/mip_losses.py +++ b/.agents/skills/puzzletron/mip_losses.py @@ -14,8 +14,12 @@ # limitations under the License. # Generated with Claude Code -"""Teacher vs. compressed model accuracy for the MIP solution.""" +"""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 @@ -31,10 +35,9 @@ print("No log.txt found. Run /puzzletron all first.") sys.exit(1) -# Extract puzzle_dir from any path that contains /mip/puzzle_solutions/ in the log +# 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: - # Fall back: look for the ckpts/teacher path logged early in the run match = re.search(r"(\S+)/ckpts/teacher", text) if not match: print("Could not find puzzle_dir in log.txt. Has the pipeline run?") @@ -42,12 +45,46 @@ 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) -solutions_dir = solutions_dirs[-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"]: @@ -62,9 +99,15 @@ 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 index 2065cc18ce3..4d22cbd89be 100644 --- a/.agents/skills/puzzletron/mip_progress.py +++ b/.agents/skills/puzzletron/mip_progress.py @@ -126,7 +126,8 @@ def get_ts(line): for i, r in enumerate(all_rates): if r not in rate_start: continue - end = rate_start[all_rates[i + 1]] if i + 1 < len(all_rates) else (complete_ts or now) + 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) 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/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 index b031515b39f..9c722bb6b5e 100644 --- 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 @@ -15,6 +15,11 @@ puzzle_dir: /workspace/puzzle_dir_qwen3_5-0.8b 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 From c1f2b9c46db2c687eba260a08a056d8234e77508 Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Mon, 22 Jun 2026 00:48:01 -0700 Subject: [PATCH 28/28] Update changelog Signed-off-by: Daniel Korzekwa --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 852a352d891..5236b54aa64 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,7 +6,7 @@ Changelog **New Features** -- Add **experimental** ``/puzzletron`` Claude Code agent skill (``.agents/skills/puzzletron/``) with ``mip`` and ``all`` commands for running the MIP step or full pipeline, and ``mip progress`` / ``all progress`` sub-commands reporting per-step status, elapsed time, and estimated time remaining. See `.agents/skills/puzzletron/README.md `_. +- 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``.