reorg: remove superseded path llama2-13b/

llama2-13b/README.md

@@ -1,43 +0,0 @@
----
-license: llama2
-base_model: NousResearch/Llama-2-13b-chat-hf
-tags:
-  - tensor-decomposition
-  - tensor-ring
-  - model-compression
-  - llama2
----
-
-# TR-3: LLaMA-2-13B-chat compressed with a Tensor Ring
-
-Every attention and MLP projection replaced by a **3-site Tensor Ring** (cyclic — the ring
-closes with a non-trivial boundary bond, so it is not a train/TT) plus a 5% sparse residual,
-then healed with pure knowledge distillation.
-
-- **7.740B parameters**, down from 13.016B — **40.53% reduction**
-- **60M healing tokens**, pure forward-KL from the unmodified 13B teacher (alpha=1, no gold labels)
-- Held-out token-mean KL: **0.3612 → 0.1333**
-- MMLU (5-shot, 40/subtask probe): **39.17% → 49.65%**, teacher **54.43%** — **91.2% retained**
-
-The highest retention of any model in this study, consistent with larger models tolerating a
-fixed compression rate better: at ~40% reduction the pre-heal KL is 3.5810 for Qwen2.5-3B,
-0.6915 for LLaMA-2-7B and 0.3612 for this model.
-
-## ⚠️ Full benchmark table not yet measured
-
-The five-task evaluation on the full benchmarks has **not been run on this checkpoint**. The
-numbers above are healing-time probes (held-out KL, MMLU subsampled to 40 questions per
-subtask) — internally consistent and comparable across this study, but **not** full-benchmark
-figures. The dense teacher scores MMLU 53.53 / HellaSwag 82.10 / BoolQ 81.93 / TriviaQA 65.18
-/ GSM8K 35.18 (average 63.58) on the full benchmarks under `eval_protocol_13b.json`.
-
-## ⚠️ This is NOT a standard `transformers` checkpoint
-
-`state_dict_fp16_healed.pt` holds **ring cores**, not dense `nn.Linear` weights. Rebuild the
-architecture with `TRLinear`/`SparseAugTRLinear` from `tensor_formats.py` using `ranks.json`;
-`AutoModelForCausalLM.from_pretrained()` will not work. LLaMA-2 has no attention biases and
-untied embeddings, so the loader is the plain case — see `../qwen2.5-3b/README.md` for the
-recipe, substituting this directory's base model and files.
-
-Healing used `micro=16, grad_accum=1, seq_len=1024` — `tokens_per_step = 16384`, the same
-global batch as the other TR-3 arms.

llama2-13b/config.json

@@ -1,26 +0,0 @@
-{
-  "_name_or_path": null,
-  "architectures": [
-    "LlamaForCausalLM"
-  ],
-  "bos_token_id": 1,
-  "eos_token_id": 2,
-  "hidden_act": "silu",
-  "hidden_size": 5120,
-  "initializer_range": 0.02,
-  "intermediate_size": 13824,
-  "max_position_embeddings": 4096,
-  "model_type": "llama",
-  "num_attention_heads": 40,
-  "num_hidden_layers": 40,
-  "num_key_value_heads": 40,
-  "pad_token_id": 0,
-  "pretraining_tp": 1,
-  "rms_norm_eps": 1e-05,
-  "rope_scaling": null,
-  "tie_word_embeddings": false,
-  "torch_dtype": "float16",
-  "transformers_version": "4.31.0.dev0",
-  "use_cache": true,
-  "vocab_size": 32000
-}

llama2-13b/eval_protocol_13b.json

@@ -1,55 +0,0 @@
-{
-  "model_id": "NousResearch/Llama-2-13b-chat-hf",
-  "model_revision": "d73f5fa9c4bc135502e04c27b39660747172d76b",
-  "lm_eval_revision": "0.4.12",
-  "tasks": {
-    "mmlu": "mmlu",
-    "hellaswag": "hellaswag",
-    "boolq": "boolq",
-    "triviaqa": "triviaqa",
-    "gsm8k": "gsm8k"
-  },
-  "num_fewshot": {
-    "mmlu": 5,
-    "hellaswag": 10,
-    "boolq": 0,
-    "triviaqa": 5,
-    "gsm8k": 5
-  },
-  "apply_chat_template": false,
-  "batch_size": "auto",
-  "dtype": "float16",
-  "seed": 1234,
-  "score_scale": "fraction",
-  "result_parser": {
-    "mmlu": "acc,none",
-    "hellaswag": "acc_norm,none",
-    "boolq": "acc,none",
-    "triviaqa": "exact_match,remove_whitespace",
-    "gsm8k": "exact_match,strict-match"
-  },
-  "hardware": {
-    "gpus_max": 4,
-    "gpu_model": "NVIDIA GB200 (Grace-Blackwell NVL4, sm_100)",
-    "note": "RIKYU compute node, 4x GB200 with 184GiB each. NOT the hardware the 7B results were produced on (that was 4x RTX A6000, sm_86). Numerically irrelevant for fp16 eval but recorded because the software stack differs too -- see torch_deviation below."
-  },
-  "torch_deviation": {
-    "locked": "torch==2.4.1+cu121",
-    "actual": "torch==2.7.1+cu128",
-    "reason": "requirements.compactifai.lock.txt was exported from an x86_64 box. torch 2.4.1+cu121 ships no aarch64 wheel, and the whole nvidia-*-cu12==12.1.x stack it depends on is x86_64/win_amd64 only. Independently, CUDA 12.1 predates Blackwell: GB200 is sm_100 and needs CUDA 12.8+. Every other package in the lock is installed at the locked version. This deviation is FORCED, not chosen, and applies to the 13B numbers but NOT to the already-recorded 7B numbers.",
-    "implication": "13B results in this protocol are not same-stack comparable to the 7B results in eval_protocol.json. To claim a controlled 7B-vs-13B comparison, re-run the 7B under this stack."
-  },
-  "added_packages": {
-    "sentencepiece": "0.2.2",
-    "protobuf": "7.36.1",
-    "reason": "ADDITIVE, not a version change: neither package appears in requirements.compactifai.lock.txt at any version. The 7B repo ships a prebuilt fast tokenizer (tokenizer.json), so transformers never needed to convert anything; the 13B repo ships only tokenizer.model (the raw sentencepiece model), and converting it requires both libraries. Without them AutoTokenizer.from_pretrained raises ImportError before any compute starts.",
-    "equivalence_check": "The converted 13B tokenizer was compared against the 7B's prebuilt tokenizer.json on four probes (English prose, Python source, an MMLU-shaped few-shot prompt, and Japanese text): identical input_ids on all four, identical vocab_size 32000, identical bos/eos/pad ids 1/2/0. Both repos carry the same LLaMA-2 sentencepiece model, so the conversion reproduces the tokenizer rather than approximating it."
-  },
-  "assumptions": [
-    "Model source: mirrors the 7B decision. We use the non-gated NousResearch/Llama-2-13b-chat-hf pinned at commit d73f5fa9c4bc135502e04c27b39660747172d76b, because meta-llama/Llama-2-13b-chat-hf is manually gated. All three safetensors shards match the gated repo's byte-for-byte in SIZE (9,948,693,272 / 9,904,129,368 / 6,178,962,272); shard sizes agreeing is consistent with a straight redistribution but is NOT a content hash, so this is evidence, not proof. Reproduction level for weights: close (unofficial host, contents not independently hashed against Meta's).",
-    "Every protocol field except model_id/model_revision/hardware is copied UNCHANGED from eval_protocol.json. Shot counts, metric keys, dtype, seed, batch_size and score_scale are identical, so the 7B and 13B rows differ only in the model -- which is the whole point of the scaling claim. Do not tune any of these per-scale.",
-    "Architecture: hidden 5120, intermediate 13824, 40 layers, vocab 32000 -> 13,015,864,320 dense parameters. These are read from the model config at runtime by scripts/whitened_lowrank.py:_arch_dims(), not hardcoded here, so this file cannot drift from the actual weights.",
-    "Compression budget is set to match the 7B's ACHIEVED ratio, not its absolute size: 4,023,119,424 / 6,738,415,616 = 0.597042 (40.30% reduction). Applied to 13B that is 7,771,022,084 total, of which 634,388,480 is the 5% sparse residual, giving --target-params-b 7.1366 for the ring cores + embeddings + norms.",
-    "TriviaQA remains the least certain task, exactly as documented in the 7B protocol. The same 5-shot exact_match (remove_whitespace) convention is locked here so the 7B-vs-13B delta stays meaningful even where the absolute deviates from the paper."
-  ]
-}

llama2-13b/ranks.json

@@ -1,3362 +0,0 @@
-{
- "model.layers.0.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.0.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.0.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.0.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.0.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.0.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.0.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.1.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.1.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.1.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.1.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.1.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.1.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.1.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.2.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.2.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.2.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.2.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.2.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.2.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.2.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.3.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.3.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.3.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.3.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.3.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.3.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.3.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.4.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.4.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.4.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.4.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.4.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.4.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.4.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.5.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.5.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.5.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.5.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.5.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.5.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.5.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.6.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.6.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.6.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.6.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.6.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.6.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.6.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.7.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.7.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.7.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.7.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.7.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.7.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.7.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.8.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.8.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.8.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.8.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.8.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.8.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.8.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.9.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.9.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.9.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.9.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.9.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.9.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.9.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.10.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.10.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.10.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.10.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.10.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.10.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.10.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.11.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.11.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.11.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.11.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.11.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.11.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.11.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.12.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.12.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.12.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.12.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.12.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.12.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.12.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.13.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.13.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.13.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.13.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.13.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.13.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.13.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.14.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.14.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.14.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.14.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.14.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.14.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.14.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.15.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.15.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.15.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.15.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.15.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.15.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.15.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.16.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.16.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.16.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.16.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.16.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.16.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.16.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.17.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.17.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.17.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.17.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.17.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.17.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.17.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.18.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.18.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.18.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.18.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.18.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.18.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.18.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.19.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.19.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.19.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.19.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.19.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.19.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.19.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.20.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.20.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.20.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.20.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.20.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.20.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.20.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.21.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.21.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.21.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.21.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.21.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.21.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.21.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.22.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.22.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.22.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.22.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.22.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.22.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.22.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.23.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.23.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.23.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.23.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.23.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.23.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.23.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.24.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.24.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.24.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.24.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.24.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.24.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.24.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.25.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.25.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.25.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.25.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.25.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.25.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.25.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.26.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.26.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.26.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.26.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.26.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.26.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.26.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.27.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.27.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.27.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.27.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.27.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.27.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.27.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.28.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.28.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.28.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.28.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.28.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.28.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.28.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.29.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.29.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.29.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.29.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.29.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.29.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.29.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.30.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.30.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.30.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.30.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.30.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.30.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.30.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.31.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.31.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.31.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.31.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.31.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.31.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.31.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.32.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.32.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.32.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.32.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.32.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.32.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.32.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.33.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.33.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.33.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.33.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.33.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.33.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.33.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.34.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.34.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.34.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.34.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.34.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.34.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.34.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.35.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.35.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.35.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.35.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.35.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.35.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.35.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.36.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.36.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.36.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.36.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.36.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.36.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.36.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.37.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.37.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.37.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.37.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.37.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.37.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.37.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.38.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.38.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.38.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.38.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.38.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.38.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.38.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.39.self_attn.q_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.39.self_attn.k_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.39.self_attn.v_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.39.self_attn.o_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 124,
-  "r1": 124,
-  "r2": 124,
-  "nnz": 1310720
- },
- "model.layers.39.mlp.gate_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.39.mlp.up_proj": {
-  "o1": 24,
-  "o2": 24,
-  "o3": 24,
-  "i1": 16,
-  "i2": 16,
-  "i3": 20,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- },
- "model.layers.39.mlp.down_proj": {
-  "o1": 16,
-  "o2": 16,
-  "o3": 20,
-  "i1": 24,
-  "i2": 24,
-  "i3": 24,
-  "r0": 174,
-  "r1": 174,
-  "r2": 174,
-  "nnz": 3538944
- }
-}

llama2-13b/state_dict_fp16_healed.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1f344018d52f0c8001b3721e6dcec601e82bb404ed755cecfbffcaf48057b52f
-size 20556180503

llama2-13b/tensor_formats.py

@@ -1,768 +0,0 @@
-from __future__ import annotations
-
-"""General N-core tensor formats for a linear layer: tensor train (MPO) and tensor ring.
-
-WHY THIS EXISTS. The 2-core and 3-core results say the tensorisation step, not the healing
-step, is where the leverage is: at equal parameters the 3-core beat the 2-core AFTER healing
-at both budgets tried (4.72B: 38.94 vs ~36.5 full MMLU at 60M tokens; 5.74B: 44.91 vs 41.54
-probe), and it got there on roughly a quarter of the tokens. That invites two orthogonal
-questions this module exists to answer: how many cores, and which topology.
-
-TENSOR TRAIN vs TENSOR RING. A TT/MPO pins the two boundary bonds to 1 -- the chain has open
-ends. A tensor ring closes the chain: core k has shape (r_{k-1}, d_k, r_k) with r_N == r_0,
-and the weight is the TRACE of the cyclic product. TT is exactly the TR special case r_0 = 1,
-so at equal ranks TR is strictly more expressive. It is also strictly harder to fit: TT-SVD
-is a sequence of Eckart-Young-optimal truncations, while a ring has no such closed form, so
-TR is fitted here by TT initialisation followed by alternating least squares. That asymmetry
-matters when reading results -- a TR loss could mean "the format is worse" OR "the fit did
-not converge", so `fit_tr` reports its per-sweep error curve and refuses to return a fit
-worse than its own TT initialisation.
-
-INDEX CONVENTION (shared with mpo2.py / mpo3.py). Output mode k travels with input mode k:
-
-    R(W)[d_1, ..., d_N] = W[(o_1..o_N), (i_1..i_N)],   d_k = (o_k, i_k) flattened
-
-Every mode must be > 1. A mode of 1 collapses a site and silently reduces the core count,
-which would make an "N-core" result actually be an (N-1)-core one.
-
-WHITENING is diagonal and factorised across the N input modes, as in mpo3.py: the whitening
-is then a mode-local rescaling that survives the fit exactly and un-whitens by division.
-"""
-
-import torch
-
-
-# --------------------------------------------------------------------------- rearrangement
-def rearrange_n(W: torch.Tensor, osplit, isplit) -> torch.Tensor:
-    """W[(o_1..o_N),(i_1..i_N)] -> R[d_1,...,d_N] with d_k = o_k*i_k."""
-    n = len(osplit)
-    if len(isplit) != n:
-        raise ValueError("output and input must be split into the same number of modes")
-    perm = [j for k in range(n) for j in (k, n + k)]
-    return (W.reshape(*osplit, *isplit).permute(*perm)
-             .reshape(*[o * i for o, i in zip(osplit, isplit)]))
-
-
-def inv_rearrange_n(R: torch.Tensor, osplit, isplit) -> torch.Tensor:
-    """Inverse of `rearrange_n`."""
-    n = len(osplit)
-    inter = [x for k in range(n) for x in (osplit[k], isplit[k])]
-    perm = list(range(0, 2 * n, 2)) + list(range(1, 2 * n, 2))
-    o = 1
-    i = 1
-    for k in range(n):
-        o *= osplit[k]
-        i *= isplit[k]
-    return R.reshape(*inter).permute(*perm).reshape(o, i)
-
-
-def _apply_scales(R: torch.Tensor, osplit, isplit, scales, invert: bool):
-    """Multiply (or divide) R by sqrt(d_k[i_k]) along each input mode, in place of a dense L."""
-    n = len(osplit)
-    shape = [x for k in range(n) for x in (osplit[k], isplit[k])]
-    T = R.reshape(*shape)
-    for k in range(n):
-        s = scales[k].clamp_min(1e-12).sqrt()
-        view = [1] * (2 * n)
-        view[2 * k + 1] = isplit[k]
-        T = T / s.view(*view) if invert else T * s.view(*view)
-    return T.reshape(R.shape)
-
-
-# --------------------------------------------------------------------------- contraction
-def contract_tt(cores) -> torch.Tensor:
-    """Contract TT cores [(1,d1,r1), (r1,d2,r2), ..., (r_{N-1},dN,1)] -> R[d_1,...,d_N]."""
-    out = cores[0]
-    for g in cores[1:]:
-        a, b = out.shape[-1], g.shape[0]
-        if a != b:
-            raise ValueError(f"TT bond mismatch {a} vs {b}")
-        out = torch.tensordot(out, g, dims=([out.dim() - 1], [0]))
-    if out.shape[0] != 1 or out.shape[-1] != 1:
-        raise ValueError("TT boundary bonds must be 1")
-    return out.reshape(*out.shape[1:-1])
-
-
-def contract_tr(cores, chunk: int = 0) -> torch.Tensor:
-    """Contract TR cores [(r0,d1,r1), ..., (r_{N-1},dN,r0)] -> R[d_1,...,d_N].
-
-    The obvious route -- chain every tensordot, then take the diagonal over the two ring
-    indices -- materialises (r_0, d_1..d_N, r_0), i.e. r_0^2 times the size of the ANSWER.
-    On one LLaMA attention layer at ring rank 118 that is 934 GB, so it does not merely run
-    slowly, it cannot run at all. Folding the last contraction and the trace into a single
-    einsum keeps the peak at r_0 * prod(d_1..d_{N-1}) * r_{N-1}, and chunking the first mode
-    bounds it further for the wide MLP layers.
-    """
-    n = len(cores)
-    if cores[0].shape[0] != cores[-1].shape[-1]:
-        raise ValueError(f"ring does not close: {cores[0].shape[0]} vs {cores[-1].shape[-1]}")
-    d0 = cores[0].shape[1]
-    if chunk and d0 > chunk:
-        parts = []
-        for lo in range(0, d0, chunk):
-            sub = [cores[0][:, lo:min(lo + chunk, d0), :]] + list(cores[1:])
-            parts.append(contract_tr(sub, chunk=0))
-        return torch.cat(parts, dim=0)
-    out = cores[0]
-    for g in cores[1:-1]:
-        out = torch.tensordot(out, g, dims=([out.dim() - 1], [0]))
-    return torch.einsum("i...k,kdi->...d", out, cores[-1])
-
-
-# --------------------------------------------------------------------------- fitting
-def fit_tt(R: torch.Tensor, ranks):
-    """TT-SVD with the given internal ranks. Singular values split symmetrically at each bond,
-    matching fit_mpo2/fit_mpo3's gauge so cores stay on comparable scales (fp16-safe and
-    kinder to a blockwise quantizer)."""
-    dims = list(R.shape)
-    n = len(dims)
-    if len(ranks) != n - 1:
-        raise ValueError(f"{n} sites need {n-1} internal ranks, got {len(ranks)}")
-    cores, rest, rprev = [], R.reshape(dims[0], -1), 1
-    for k in range(n - 1):
-        M = rest.reshape(rprev * dims[k], -1)
-        U, S, Vh = torch.linalg.svd(M, full_matrices=False)
-        r = min(int(ranks[k]), S.shape[0])
-        q = S[:r].clamp_min(0).sqrt()
-        cores.append((U[:, :r] * q).reshape(rprev, dims[k], r))
-        rest = q[:, None] * Vh[:r]
-        rprev = r
-    cores.append(rest.reshape(rprev, dims[-1], 1))
-    return cores
-
-
-def tt_to_tr(cores):
-    """Reinterpret TT cores as TR cores with r_0 = 1 (an exact, lossless re-labelling)."""
-    return [c.clone() for c in cores]
-
-
-def _subchain(cores, k):
-    """Contract every core except k, in cyclic order, -> E[r_k, D, r_{k-1}] with D = prod_{j!=k} d_j."""
-    n = len(cores)
-    order = [(k + 1 + t) % n for t in range(n - 1)]
-    out = cores[order[0]]
-    for j in order[1:]:
-        out = torch.tensordot(out, cores[j], dims=([out.dim() - 1], [0]))
-    return out
-
-
-def _subchain_chunk(cores, k, lo, hi):
-    """`_subchain` with the FIRST following mode restricted to [lo, hi). The restriction rides
-    along the contraction untouched, so concatenating the chunks reproduces `_subchain`."""
-    n = len(cores)
-    order = [(k + 1 + t) % n for t in range(n - 1)]
-    out = cores[order[0]][:, lo:hi, :]
-    for j in order[1:]:
-        out = torch.tensordot(out, cores[j], dims=([out.dim() - 1], [0]))
-    return out
-
-
-def rank_of(cores, k):
-    """r_{k-1} * r_k for core k -- the width of that core's least-squares system."""
-    n = len(cores)
-    return cores[k].shape[0] * cores[k].shape[2]
-
-
-def _solve_maybe_cpu(AtA, AtW):
-    """Solve AtA X = AtW, falling back to the CPU when the GPU cannot hold the factorisation.
-
-    The normal-equation matrix is (r_{k-1} r_k) square, i.e. it grows as r^4. On down_proj at
-    ring rank 219 that is 47961^2 = 9.2 GB, and the factorisation needs a second copy -- no
-    amount of chunking the data helps, because this matrix does not depend on the data size.
-    Cholesky is tried first (the Gram is symmetric positive definite after the ridge) and
-    needs only one triangular factor; if even that does not fit, the solve moves to the host.
-    Slow, but a bounded slowdown beats an unbounded failure on the layers that matter most.
-    """
-    try:
-        L = torch.linalg.cholesky(AtA)
-        return torch.cholesky_solve(AtW, L)
-    except torch.OutOfMemoryError:
-        torch.cuda.empty_cache()
-    except Exception:
-        pass
-    try:
-        return torch.linalg.solve(AtA, AtW)
-    except torch.OutOfMemoryError:
-        torch.cuda.empty_cache()
-    dev = AtA.device
-    X = torch.linalg.solve(AtA.cpu().double(), AtW.cpu().double())
-    return X.to(device=dev, dtype=AtW.dtype)
-
-
-def tr_rel_err(cores, R: torch.Tensor, chunk_elems: int = 250_000_000) -> float:
-    """||contract_tr(cores) - R||_F / ||R||_F, accumulated in chunks.
-
-    Materialising the reconstruction to subtract it needs the reconstruction AND the einsum's
-    intermediate at once; on down_proj at ring rank 155 that intermediate is (155, 64, 688,
-    256) = 7 GB even with the first mode already chunked. Accumulating the squared error one
-    slice at a time keeps only one slice live, and the chunk width is derived from an element
-    budget rather than a hard-coded constant -- a constant is exactly what silently failed to
-    engage when a mode happened to equal it.
-    """
-    n = len(cores)
-    d0 = cores[0].shape[1]
-    mid = 1
-    for k in range(1, n - 1):
-        mid *= cores[k].shape[1]
-    per_unit = max(1, mid * cores[0].shape[0] * cores[-1].shape[0])
-    step = max(1, min(d0, int(chunk_elems // per_unit)))
-    acc = 0.0
-    for lo in range(0, d0, step):
-        hi = min(lo + step, d0)
-        sub = [cores[0][:, lo:hi, :]] + list(cores[1:])
-        diff = contract_tr(sub) - R[lo:hi]
-        acc += float(diff.pow(2).sum())
-        del sub, diff
-    return (acc ** 0.5) / float(R.norm())
-
-
-def fit_tr(R: torch.Tensor, ranks, sweeps: int = 30, tol: float = 1e-8, verbose: bool = False,
-           seed: int = 1234, restarts: int = 2, solver: str = "lstsq",
-           chunk: int = 0, chunk_elems: int = 250_000_000):
-    """Fit a tensor ring by TT initialisation + alternating least squares.
-
-    ranks = [r_0, r_1, ..., r_{N-1}] with r_N == r_0. The TT init uses r_0 = 1 (a legal ring),
-    then each ALS sweep re-solves every core in closed form given the others; the ring's extra
-    boundary rank is opened up by padding the initialisation. Because ALS on a ring is not
-    guaranteed to improve monotonically from an arbitrary start, the best iterate seen is kept
-    and the routine REFUSES to return anything worse than its own TT initialisation -- without
-    that guard a fitting failure would be indistinguishable from the format being worse, which
-    is precisely the comparison this module exists to make.
-    """
-    n = len(R.shape)
-    if len(ranks) != n:
-        raise ValueError(f"a ring over {n} sites needs {n} ranks (r_0..r_{{N-1}})")
-    dims = list(R.shape)
-    rnorm = float(R.norm())
-    rnorm2 = float(R.pow(2).sum())
-    last_sq = None
-    tt = fit_tt(R, ranks[1:])
-    base_err = float((contract_tt(tt) - R).norm() / R.norm())
-    r0 = int(ranks[0])
-    gen = torch.Generator(device="cpu").manual_seed(seed)
-
-    def padded_init():
-        """TT solution embedded in the ring, with the OPENED bond seeded on BOTH ends.
-
-        Seeding only one end is not enough and is a silent failure: core 0's least-squares
-        environment is built from cores 1..N-1, whose ring index comes from core N-1's
-        trailing bond. If that stays zero, the environment vanishes on exactly the components
-        the ring adds, every ALS solve returns zero there, and the fit never leaves the TT
-        solution -- it looks like "the ring buys nothing" rather than "the ring was never
-        initialised".
-        """
-        out = []
-        for k, c in enumerate(tt):
-            a, d, b = c.shape
-            na = r0 if k == 0 else a
-            nb = r0 if k == n - 1 else b
-            g = torch.zeros(na, d, nb, dtype=R.dtype, device=R.device)
-            g[:a, :, :b] = c
-            eps = 0.05 * float(c.norm()) / max(1, c.numel()) ** 0.5
-            if k == 0 and na > a:
-                g[a:, :, :] = eps * torch.randn(na - a, d, nb, generator=gen,
-                                                dtype=R.dtype).to(R.device)
-            if k == n - 1 and nb > b:
-                g[:, :, b:] = eps * torch.randn(na, d, nb - b, generator=gen,
-                                                dtype=R.dtype).to(R.device)
-            out.append(g)
-        return out
-
-    def random_init(g=None):
-        """i.i.d. cores whose ring contraction has the right variance.
-
-        The scale is derived, not guessed. For cores with i.i.d. N(0, s^2) entries the trace
-        of the cyclic product is a sum of prod_k r_k independent products of N Gaussians, so
-        Var(W) ~= (prod_k r_k) * s^(2N). Matching that to the target mean square gives
-        s = (E[R^2] / prod_k r_k)^(1/(2N)). Getting this wrong by a constant factor per core
-        compounds to that factor^N in the product, which for N=4,5 puts the initial iterate
-        many orders of magnitude away from the data and leaves ALS crawling.
-        """
-        rr = [int(ranks[k]) for k in range(n)]
-        prod = 1.0
-        for x in rr:
-            prod *= x
-        ms = float(R.pow(2).mean())
-        s = (ms / prod) ** (1.0 / (2 * n))
-        gg = g if g is not None else gen
-        return [s * torch.randn(rr[k], dims[k], rr[(k + 1) % n], generator=gg,
-                                dtype=R.dtype).to(R.device) for k in range(n)]
-
-    # ALS on a ring has genuine local minima -- a single start that stalls would be reported
-    # as "the ring is not better", which is exactly the confound this module must avoid. So
-    # several starts are run and the best kept; restarts are cheap next to the SVDs anyway.
-    inits = [("tt_padded", padded_init)]
-    for j in range(restarts):
-        g = torch.Generator(device="cpu").manual_seed(seed + 977 * (j + 1))
-        inits.append((f"random{j}", lambda g=g: random_init(g)))
-    best, best_err, hists, winner = tt_to_tr(tt), base_err, {}, "tt_init"
-    for init_name, init_fn in inits:
-        cores, ihist = init_fn(), []
-        for it in range(sweeps):
-            for k in range(n):
-                # The subchain leaves every d index separate and in CYCLIC order starting at
-                # k+1. R's remaining modes must be flattened in that SAME order -- movedim
-                # gives natural order, which for interior k is a different permutation and
-                # would solve the least-squares problem against mismatched columns. It still
-                # returns a finite "solution", just to the wrong system.
-                # The chunk decision is made from SHAPES ONLY and BEFORE anything large is
-                # built. An earlier version computed the full subchain first and then decided
-                # whether to chunk, so the unchunked allocation always happened anyway and the
-                # chunked branch saved nothing -- the memory-bounded path existed but could
-                # never take effect.
-                # Auto-size the chunk from a memory budget instead of comparing against a
-                # fixed number. A constant threshold silently does nothing when the mode
-                # happens to equal it -- which is exactly what let a 4-site fit try to
-                # materialise a 134 GB subchain (mode 64, chunk 64, "64 > 64" is false).
-                m_ = (k + 1) % n
-                other = 1
-                for j in range(n):
-                    if j != k:
-                        other *= dims[j]
-                per_unit = max(1, (other // max(1, dims[m_])) * rank_of(cores, k))
-                auto = max(1, min(dims[m_], int(chunk_elems // per_unit)))
-                use_chunk = auto < dims[m_]
-                if solver == "normal" and use_chunk:
-                    # Memory-bounded normal equations. Materialising the subchain costs
-                    # r_k * (prod_{j!=k} d_j) * r_{k-1} floats -- on down_proj at these ranks
-                    # that is (155, 688, 256, 155) = 17 GB, so the unchunked path cannot run
-                    # on the MLP layers at all. The Gram and right-hand side are sums over the
-                    # flattened "others" index, so they accumulate exactly over slices of the
-                    # first subchain mode; this is the same arithmetic, not an approximation.
-                    m = m_
-                    d_next = dims[m]
-                    rk_, rkm1_ = None, None
-                    AtA = AtW = None
-                    Rp = R.permute(*([k] + [(k + 1 + t) % n for t in range(n - 1)]))
-                    for lo in range(0, d_next, auto):
-                        hi = min(lo + auto, d_next)
-                        Ec = _subchain_chunk(cores, k, lo, hi)
-                        rk_, rkm1_ = Ec.shape[0], Ec.shape[-1]
-                        Ec = Ec.reshape(rk_, -1, rkm1_)
-                        Ac = Ec.permute(1, 2, 0).reshape(Ec.shape[1], rkm1_ * rk_)
-                        Wc = Rp[:, lo:hi].reshape(dims[k], -1)
-                        g = Ac.transpose(0, 1) @ Ac
-                        h = Ac.transpose(0, 1) @ Wc.transpose(0, 1)
-                        AtA = g if AtA is None else AtA + g
-                        AtW = h if AtW is None else AtW + h
-                        del Ec, Ac, Wc, g, h
-                    lam = 1e-8 * float(AtA.diagonal().abs().mean().clamp_min(1e-30))
-                    AtA_r = AtA + lam * torch.eye(AtA.shape[0], dtype=AtA.dtype,
-                                                  device=AtA.device)
-                    G = _solve_maybe_cpu(AtA_r, AtW)
-                    del AtA_r
-                    # ||R - A G||^2 = ||R||^2 - 2<AtW, G> + <G, AtA G>, all from quantities
-                    # already formed. Computing it by contracting the whole ring instead would
-                    # cost one full O(prod d_k) tensor per sweep per core.
-                    sq = rnorm2 - 2.0 * float((AtW * G).sum()) + float((G * (AtA @ G)).sum())
-                    last_sq = max(sq, 0.0)
-                    rkm1, rk = rkm1_, rk_
-                else:
-                    # ONE branch builds A/Wk and ONE consumes them. Splitting the decision
-                    # across two separate if/else chains left the non-chunked solve reachable
-                    # after the chunked path had run, where `A` is whatever the PREVIOUS core
-                    # left behind -- a silent wrong answer rather than an error.
-                    E = _subchain(cores, k)               # (r_k, d_{k+1},..,d_{k-1}, r_{k-1})
-                    rk, rkm1 = E.shape[0], E.shape[-1]
-                    E = E.reshape(rk, -1, rkm1)
-                    D = E.shape[1]
-                    A = E.permute(1, 2, 0).reshape(D, rkm1 * rk)
-                    perm = [k] + [(k + 1 + t) % n for t in range(n - 1)]
-                    # R's remaining modes must be flattened in the subchain's CYCLIC order --
-                    # movedim gives natural order, which for interior k solves the
-                    # least-squares problem against mismatched columns and still returns a
-                    # finite "solution", just to the wrong system.
-                    Wk = R.permute(*perm).reshape(dims[k], D)
-                    del E
-                    if solver == "normal":
-                        AtA = A.transpose(0, 1) @ A
-                        lam = 1e-8 * float(AtA.diagonal().abs().mean().clamp_min(1e-30))
-                        AtA += lam * torch.eye(AtA.shape[0], dtype=A.dtype, device=A.device)
-                        G = _solve_maybe_cpu(AtA, A.transpose(0, 1) @ Wk.transpose(0, 1))
-                        del AtA
-                    else:
-                        G = torch.linalg.lstsq(A, Wk.transpose(0, 1)).solution
-                    last_sq = None
-                    del A, Wk
-                cores[k] = (G.transpose(0, 1).reshape(dims[k], rkm1, rk)
-                             .permute(1, 0, 2).contiguous())
-            err = (float(last_sq) ** 0.5 / rnorm if last_sq is not None
-                   else float((contract_tr(cores, chunk=chunk) - R).norm() / rnorm))
-            ihist.append(round(err, 8))
-            if verbose:
-                print(f"    TR-ALS[{init_name}] sweep {it+1}: rel err {err:.6f}", flush=True)
-            if err < best_err - 1e-12:
-                best, best_err, winner = [c.clone() for c in cores], err, init_name
-            if it and abs(ihist[-2] - err) < tol * max(err, 1e-12):
-                break
-        hists[init_name] = ihist
-    # The algebraic residual is a difference of nearly equal large numbers; at a relative
-    # error of ~1e-2 it has already lost most of its fp32 significance. It is fine for
-    # steering ALS but not for the number that gets reported, so the winner is re-measured
-    # once by explicit contraction.
-    best_err = tr_rel_err(best, R)
-    if best_err > base_err:
-        return tt_to_tr(tt), base_err, {"fell_back_to_tt": True, "tt_err": base_err,
-                                        "winner": "tt_init", "als_hist": hists}
-    return best, best_err, {"fell_back_to_tt": winner == "tt_init", "tt_err": base_err,
-                            "winner": winner, "als_hist": hists}
-
-
-# --------------------------------------------------------------------------- bookkeeping
-def tt_params(dims, ranks) -> int:
-    r = [1] + [int(x) for x in ranks] + [1]
-    return sum(r[k] * dims[k] * r[k + 1] for k in range(len(dims)))
-
-
-def tr_params(dims, ranks) -> int:
-    r = [int(x) for x in ranks]
-    n = len(dims)
-    return sum(r[k] * dims[k] * r[(k + 1) % n] for k in range(n))
-
-
-def tt_rank_ceilings(dims):
-    """r_k <= min(prod_{j<=k} d_j, prod_{j>k} d_j): above that the truncation is vacuous."""
-    n, out = len(dims), []
-    for k in range(n - 1):
-        left = 1
-        for j in range(k + 1):
-            left *= dims[j]
-        right = 1
-        for j in range(k + 1, n):
-            right *= dims[j]
-        out.append(min(left, right))
-    return out
-
-
-def ranks_for_budget(dims, budget: int, ring: bool, alpha: float = 1.0):
-    """Largest ranks fitting `budget`, shaped by r_k proportional to ceiling_k ** alpha.
-
-    Bisection rather than a closed form: the count is linear in each rank but bilinear in
-    adjacent pairs, and cyclic for a ring, so there is no single quadratic to solve.
-
-    ALPHA EXISTS BECAUSE EQUAL FILL FRACTION IS WRONG WHEN THE CEILINGS ARE LOPSIDED. For a
-    4-site split of 4096 the ceilings are (64, 4096, 64): alpha=1 spends the budget on the
-    huge middle bond and returns ranks like (36, 2327, 36), strangling the two bonds that
-    actually gate information flow through the chain. alpha=0 gives every bond the same rank
-    instead. Neither is right a priori, so the caller is expected to MEASURE across alpha --
-    the reconstruction error, not the heuristic, decides.
-    """
-    ceil_tt = tt_rank_ceilings(dims)
-    ceils = ([min(ceil_tt)] + ceil_tt) if ring else ceil_tt
-    fn = tr_params if ring else tt_params
-    mx = max(ceils)
-    shape = [(c / mx) ** alpha for c in ceils]
-
-    def build(t):
-        return [max(1, min(int(round(t * mx * s)), c)) for s, c in zip(shape, ceils)]
-
-    lo, hi = 0.0, 1.0
-    if fn(dims, build(1.0)) <= budget:
-        return build(1.0)
-    for _ in range(60):
-        mid = (lo + hi) / 2
-        if fn(dims, build(mid)) > budget:
-            hi = mid
-        else:
-            lo = mid
-    r = build(lo)
-    while fn(dims, r) > budget and max(r) > 1:
-        r[r.index(max(r))] -= 1
-    return r
-
-
-def balanced_split(n: int, k: int):
-    """Most balanced k-way factorisation of n with every factor >= 2 (spread-minimising)."""
-    best = None
-    def rec(rem, parts):
-        nonlocal best
-        if len(parts) == k - 1:
-            if rem >= 2:
-                cand = parts + [rem]
-                key = max(cand) / min(cand)
-                if best is None or key < best[0]:
-                    best = (key, cand)
-            return
-        d = 2
-        while d * d <= rem or d <= rem:
-            if d > rem:
-                break
-            if rem % d == 0:
-                rec(rem // d, parts + [d])
-            d += 1
-    rec(n, [])
-    if best is None:
-        raise ValueError(f"no {k}-way factorisation of {n} with all factors >= 2")
-    return sorted(best[1])
-
-
-def self_test(device: str = "cpu") -> None:
-    torch.manual_seed(0)
-    dt = torch.float64
-    for osp, isp in [((4, 3), (2, 5)), ((4, 3, 2), (2, 5, 3)), ((2, 3, 2, 2), (2, 2, 3, 2))]:
-        n = len(osp)
-        out = 1
-        inn = 1
-        for x in osp:
-            out *= x
-        for x in isp:
-            inn *= x
-        W = torch.randn(out, inn, dtype=dt, device=device)
-        R = rearrange_n(W, osp, isp)
-        dims = [o * i for o, i in zip(osp, isp)]
-        assert list(R.shape) == dims, (R.shape, dims)
-        assert torch.equal(inv_rearrange_n(R, osp, isp), W), f"{n}-site rearrange not invertible"
-        # index identity, not just the round trip: a wrong-but-self-inverse permutation would
-        # pass the round trip and scramble every weight silently
-        for _ in range(40):
-            oi = [torch.randint(0, x, (1,)).item() for x in osp]
-            ii = [torch.randint(0, x, (1,)).item() for x in isp]
-            ridx = tuple(oi[k] * isp[k] + ii[k] for k in range(n))
-            orow = 0
-            for k in range(n):
-                orow = orow * osp[k] + oi[k]
-            ocol = 0
-            for k in range(n):
-                ocol = ocol * isp[k] + ii[k]
-            assert R[ridx] == W[orow, ocol], f"{n}-site rearrange index map wrong"
-
-        ceils = tt_rank_ceilings(dims)
-        cores = fit_tt(R, ceils)
-        e = float((contract_tt(cores) - R).norm() / R.norm())
-        assert e < 1e-10, f"{n}-site full-rank TT is not exact ({e:.3e})"
-        assert tt_params(dims, ceils) == sum(c.numel() for c in cores), "tt_params disagrees"
-
-        if n >= 3:
-            trunc = [max(1, c // 2) for c in ceils]
-            cores = fit_tt(R, trunc)
-            assert tt_params(dims, trunc) == sum(c.numel() for c in cores)
-            for a in (0.0, 0.5, 1.0):
-                rr = ranks_for_budget(dims, tt_params(dims, trunc), ring=False, alpha=a)
-                assert tt_params(dims, rr) <= tt_params(dims, trunc), \
-                    f"TT budget solver overran at alpha={a}"
-            # a ring at the SAME parameter budget must not do worse than the TT it starts from
-            budget = tt_params(dims, trunc)
-            rring = ranks_for_budget(dims, budget, ring=True)
-            assert tr_params(dims, rring) <= budget, "TR budget solver overran"
-            _c, err_tr, info = fit_tr(R, rring, sweeps=6)
-            assert err_tr <= info["tt_err"] + 1e-12, "fit_tr returned worse than its own TT init"
-            assert abs(float((contract_tr(_c) - R).norm() / R.norm()) - err_tr) < 1e-8, \
-                "reported TR error does not match the returned cores"
-
-    ring_chk = [torch.randn(3, 5, 4, dtype=dt), torch.randn(4, 6, 2, dtype=dt),
-                torch.randn(2, 7, 3, dtype=dt)]
-    full = contract_tr(ring_chk)
-    assert torch.allclose(full, contract_tr(ring_chk, chunk=2), atol=1e-12), \
-        "chunked contract_tr disagrees with the unchunked one"
-    ref = torch.einsum("iaj,jbk,kci->abc", *ring_chk)
-    assert torch.allclose(full, ref, atol=1e-12), "contract_tr is not the cyclic trace"
-
-    # a ring with r_0 > 1 must be able to represent things a TT of the same internal ranks cannot
-    dims = [6, 6, 6]
-    torch.manual_seed(1)
-    ring = [torch.randn(3, 6, 3, dtype=dt) for _ in range(3)]
-    Rr = contract_tr(ring)
-    tt_same = fit_tt(Rr, [3, 3])
-    e_tt = float((contract_tt(tt_same) - Rr).norm() / Rr.norm())
-    assert e_tt > 1e-3, ("a rank-(3,3,3) ring was reproduced by a rank-(3,3) TT; the ring "
-                         f"contraction is probably not closing the loop (err {e_tt:.3e})")
-    for slv in ("lstsq", "normal"):
-        got, err, info = fit_tr(Rr, [3, 3, 3], sweeps=200, restarts=3, solver=slv)
-        assert err < 1e-4, (f"TR-ALS[{slv}] failed to recover an exact rank-3 ring "
-                            f"(err {err:.3e}); winner={info['winner']} "
-                            f"tt_init={info['tt_err']:.3e}")
-    got, err, info = fit_tr(Rr, [3, 3, 3], sweeps=200, restarts=3)
-    # chunking changes only how the Gram is accumulated, so it must reach the same optimum
-    _g, err_ch, _i = fit_tr(Rr, [3, 3, 3], sweeps=200, restarts=3, solver="normal", chunk=2)
-    assert err_ch < 1e-4, f"chunked TR-ALS failed to recover an exact ring (err {err_ch:.3e})"
-    assert err < 1e-4, (f"TR-ALS failed to recover an exact rank-3 ring (err {err:.3e}); "
-                        f"winner={info['winner']} tt_init={info['tt_err']:.3e}")
-
-    assert balanced_split(4096, 3) == [16, 16, 16], balanced_split(4096, 3)
-    assert balanced_split(4096, 2) == [64, 64]
-    b = balanced_split(11008, 3)
-    assert 43 in b and b[0] * b[1] * b[2] == 11008, b
-    print("tensor_formats self-test PASSED (N-site rearrange exact & index-checked, full-rank "
-          "TT exact, budget solvers within budget, TR >= its TT init, TR strictly beats TT on "
-          "a true ring, balanced splits correct)")
-
-
-if __name__ == "__main__":
-    self_test()
-
-
-# --------------------------------------------------------------------------- layers
-class TRLinear(torch.nn.Module):
-    """Tensor-ring linear layer: W = Tr[G_0 G_1 ... G_{N-1}] after the shared rearrangement.
-
-    NOT a LowRankLinear subclass, unlike MPO2Linear/MPO3Linear. Those two could pretend to be
-    a pair of factors because a train has open ends; a ring genuinely has N 3-D cores and no
-    A/B roles to borrow, so anything keyed on `.A`/`.B` (the blockwise quantizer, the bond
-    rebalancer) has to be taught about `.cores` explicitly rather than silently doing the
-    wrong thing to a subclass it thinks it understands.
-
-    RECONSTRUCTION IS CHUNKED BY DEFAULT. Merging any two cores of a 3-site ring costs
-    r*d_i*d_j*r entries -- on down_proj at ring rank 142 that is 7.1 GB in bf16, per layer,
-    per forward. Chunking the first mode bounds it; the cost is more, smaller matmuls.
-    """
-
-    def __init__(self, cores, osplit, isplit, chunk_elems: int = 500_000_000,
-                 bias=None):
-        super().__init__()
-        n = len(cores)
-        if len(osplit) != n or len(isplit) != n:
-            raise ValueError(f"{n} cores need {n} output and {n} input modes")
-        if min(list(osplit) + list(isplit)) < 2:
-            raise ValueError(f"degenerate split o={tuple(osplit)} i={tuple(isplit)}: a mode of "
-                             f"1 collapses a site, making this an (N-1)-site network")
-        for k in range(n):
-            if cores[k].shape[1] != osplit[k] * isplit[k]:
-                raise ValueError(f"core {k} has mode {cores[k].shape[1]}, expected "
-                                 f"{osplit[k]}*{isplit[k]}")
-            if cores[k].shape[2] != cores[(k + 1) % n].shape[0]:
-                raise ValueError(f"ring does not close between cores {k} and {(k+1)%n}")
-        self.cores = torch.nn.ParameterList([torch.nn.Parameter(c) for c in cores])
-        self.osplit, self.isplit = tuple(int(x) for x in osplit), tuple(int(x) for x in isplit)
-        o = i = 1
-        for k in range(n):
-            o *= self.osplit[k]
-            i *= self.isplit[k]
-        self.out_features, self.in_features = o, i
-        self.ranks = tuple(int(c.shape[0]) for c in cores)
-        # Elements per unit of the FIRST mode in contract_tr's peak intermediate, which is
-        # (r_0, c, d_1, ..., d_{n-2}, r_{n-1}) -- the LAST core's mode is contracted away by
-        # the closing einsum and must not be counted. Including it over-estimated the cost by
-        # d_{n-1} (256x on an MLP layer), which drove the chunk to 1 and would have made every
-        # MLP forward 256 sequential contractions instead of 8.
-        mid = 1
-        for k in range(1, n - 1):
-            mid *= cores[k].shape[1]
-        per_unit = max(1, mid * self.ranks[0] * self.ranks[-1])
-        self._chunk = max(1, min(cores[0].shape[1], int(chunk_elems // per_unit)))
-        # non-persistent: never wants to show up in a checkpoint, and materialize_() is an
-        # inference-only opt-in (see below), so a freshly-loaded layer must start uncached.
-        self.register_buffer("_cached_weight", None, persistent=False)
-        # LLaMA projections have no bias, but Qwen2's q/k/v do. Dropping it on
-        # replacement is SILENT -- the shapes still match, the parameter-count gate
-        # still balances (achieved excludes bias, and so would the rebuilt model), and
-        # only the accuracy moves. So the bias is carried, not discarded.
-        self.bias = None if bias is None else torch.nn.Parameter(bias)
-
-    def weight(self) -> torch.Tensor:
-        if self._cached_weight is not None:
-            return self._cached_weight
-        cs = list(self.cores)
-        R = contract_tr(cs, chunk=self._chunk if self._chunk < cs[0].shape[1] else 0)
-        return inv_rearrange_n(R, self.osplit, self.isplit)
-
-    @torch.no_grad()
-    def materialize_(self):
-        """Cache the dense weight once. INFERENCE ONLY -- gradients cannot flow through the
-        cached tensor back to the cores, so this must never be called before/during training.
-
-        forward() calls weight() on every call, which redoes the full chunked ring
-        contraction (two tensor contractions + a trace per core) from scratch. That is fine
-        for a single forward pass, but autoregressive generation calls forward() once per
-        layer per GENERATED TOKEN with the cores completely unchanged between calls. Measured
-        on tr3_ring4b04_h1's GSM8K eval before this existed: 209s/item, ETA 73h for 1319
-        items -- every one of ~200 generated tokens re-contracted all 224 ring layers.
-        A train layer (LowRankLinear/MPO2Linear) never had this problem: F.linear(F.linear(x,
-        B), A) applies the two factors straight to the activation and never materializes W at
-        all, so there was nothing to cache there.
-        """
-        self._cached_weight = self.weight().detach().clone()
-
-    def unmaterialize_(self):
-        self._cached_weight = None
-
-    def forward(self, x):
-        b = None if self.bias is None else self.bias.to(x.dtype)
-        return torch.nn.functional.linear(x, self.weight().to(x.dtype), b)
-
-    def extra_repr(self) -> str:
-        return (f"o={self.osplit}, i={self.isplit}, ranks={self.ranks}, "
-                f"chunk={self._chunk}, materialized={self._cached_weight is not None}")
-
-
-class SparseAugTRLinear(TRLinear):
-    """TRLinear plus a sparse additive residual. Same int64-index contract as the MPO layers:
-    `sp_idx` must NEVER be cast to a floating dtype -- a `.half()` on a 45M-scale index
-    silently rounds it to a different, valid-looking position."""
-
-    def __init__(self, cores, osplit, isplit, sp_idx, sp_val,
-                 chunk_elems: int = 500_000_000, bias=None):
-        super().__init__(cores, osplit, isplit, chunk_elems, bias)
-        if sp_idx.numel() != sp_val.numel():
-            raise ValueError(f"{sp_idx.numel()} indices for {sp_val.numel()} values")
-        n = self.out_features * self.in_features
-        if sp_idx.numel() and (int(sp_idx.max()) >= n or int(sp_idx.min()) < 0):
-            raise ValueError("sparse index out of range for this layer")
-        self.register_buffer("sp_idx", sp_idx.to(torch.int64))
-        self.sp_val = torch.nn.Parameter(sp_val)
-
-    def weight(self) -> torch.Tensor:
-        # must check the cache here too, not just in TRLinear.weight(): materialize_() caches
-        # THIS class's weight() (ring + sparse, since self.weight() dispatches polymorphically
-        # from inside materialize_()), so without this check a later call would go
-        # super().weight() -> see the cache -> return the already-sparse-augmented tensor ->
-        # and then add the sparse residual a second time on top of it.
-        if self._cached_weight is not None:
-            return self._cached_weight
-        W = super().weight()
-        if self.sp_idx.numel():
-            W = W.reshape(-1).index_add(0, self.sp_idx, self.sp_val.to(W.dtype)).view_as(W)
-        return W
-
-    def extra_repr(self) -> str:
-        return super().extra_repr() + f", nnz={self.sp_idx.numel()}"
-
-
-def layer_self_test() -> None:
-    torch.manual_seed(0)
-    osp, isp = (4, 3, 2), (2, 5, 3)
-    dims = [o * i for o, i in zip(osp, isp)]
-    cores = [torch.randn(3, dims[0], 4), torch.randn(4, dims[1], 2), torch.randn(2, dims[2], 3)]
-    lay = TRLinear([c.clone() for c in cores], osp, isp, chunk_elems=10 ** 9)
-    ref = inv_rearrange_n(contract_tr(cores), osp, isp)
-    assert torch.allclose(lay.weight(), ref, atol=1e-5), "TRLinear.weight() != contract+inverse"
-    # chunking must not change the answer
-    lay_c = TRLinear([c.clone() for c in cores], osp, isp, chunk_elems=1)
-    assert lay_c._chunk < dims[0], "chunk did not engage"
-    assert torch.allclose(lay_c.weight(), ref, atol=1e-5), "chunked forward differs"
-    x = torch.randn(7, lay.in_features)
-    assert torch.allclose(lay(x), x @ ref.T, atol=1e-4), "forward != x W^T"
-    idx = torch.tensor([0, 5, lay.out_features * lay.in_features - 1], dtype=torch.int64)
-    val = torch.tensor([1.0, -2.0, 3.0])
-    sl = SparseAugTRLinear([c.clone() for c in cores], osp, isp, idx, val, chunk_elems=10 ** 9)
-    d = (sl.weight() - ref).reshape(-1)
-    assert torch.allclose(d[idx], val, atol=1e-5), "sparse residual not applied at sp_idx"
-    assert abs(float(d.abs().sum()) - float(val.abs().sum())) < 1e-4, "residual leaked"
-    try:
-        TRLinear([torch.zeros(2, 4, 2), torch.zeros(2, 3, 2), torch.zeros(2, 6, 2)],
-                 (4, 1, 2), (1, 3, 3))
-    except ValueError:
-        pass
-    else:
-        raise AssertionError("a degenerate split (mode = 1) was accepted")
-    # gradients must reach every core
-    lay.zero_grad()
-    lay(x).sum().backward()
-    for k, c in enumerate(lay.cores):
-        assert c.grad is not None and float(c.grad.abs().sum()) > 0, f"core {k} got no gradient"
-    # materialize_() must be a pure inference-speed cache: same weight, same forward output,
-    # and -- the one bug this is actually guarding against -- the sparse residual must not
-    # get added twice once SparseAugTRLinear.weight() is short-circuited by the cache.
-    pre = lay.weight().clone()
-    lay.materialize_()
-    assert lay._cached_weight is not None, "materialize_() did not populate the cache"
-    assert torch.allclose(lay.weight(), pre, atol=1e-6), "materialize_() changed weight()"
-    assert torch.allclose(lay(x), x @ ref.T, atol=1e-4), "forward != x W^T after materialize_()"
-    lay.unmaterialize_()
-    assert lay._cached_weight is None, "unmaterialize_() left the cache populated"
-    sl_ref = sl.weight().clone()
-    sl.materialize_()
-    assert torch.allclose(sl.weight(), sl_ref, atol=1e-6), \
-        "materialized SparseAugTRLinear.weight() differs -- sparse residual likely double-added"
-    print("TRLinear self-test PASSED (weight matches contraction, chunking exact, forward "
-          "correct, sparse localised, degenerate split refused, all cores differentiable, "
-          "materialize_() cache exact and idempotent)")