Underlining

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+global_loss_landscape.png filter=lfs diff=lfs merge=lfs -text

LICENSE.md

@@ -0,0 +1,40 @@
+# License
+
+This repository contains materials under two licenses. The licenses apply only
+to rights held by the repository's contributors. They do not relicense training
+datasets, source recordings, lyrics, compositions, images, documents,
+trademarks, personalities, or other third-party material.
+
+## Model weights, documentation, and visual assets
+
+The `model-*.safetensors` files, model card, configuration, architecture
+diagram, and other non-code original materials are licensed under the
+**Creative Commons Attribution-NonCommercial 4.0 International License**
+(CC BY-NC 4.0):
+
+https://creativecommons.org/licenses/by-nc/4.0/legalcode
+
+You must provide appropriate attribution, link to the license, and indicate
+whether changes were made. You may not use these materials for commercial
+purposes under this license.
+
+## Inference source code
+
+The Python source files in this repository are licensed under the
+**Apache License, Version 2.0**:
+
+https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed
+under that license is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+CONDITIONS OF ANY KIND, either express or implied. See the Apache License for
+the specific language governing permissions and limitations.
+
+## Training-data and output rights
+
+No training dataset is included in this repository. All upstream dataset terms
+and content rights remain in effect. In particular, the source audio and lyrics
+represented by `webshart/suno-various-94k` are marked source-rights-retained and
+remain subject to their creators' rights. Nothing in this license grants a
+right to reproduce protected material that may be recalled or generated by the
+model.

README.md

@@ -1,3 +1,210 @@
 ---
-license: apache-2.0
+library_name: custom
+pipeline_tag: text-generation
+license: other
+license_name: cc-by-nc-4.0-weights-apache-2.0-code
+license_link: LICENSE.md
+tags:
+  - mamba2
+  - byte-level
+  - state-space-model
+  - multimodal
+  - custom-code
+  - causal-lm
+language:
+  - en
 ---
+
+# Three-Level Nested Byte Mamba-2
+
+This repository contains a research checkpoint for a **2.478B-parameter causal byte model** with three nested Mamba-2 resolutions. It predicts raw bytes rather than tokenizer IDs and was trained on a mixture of web/PDF text, serialized image-text examples, and serialized audio.
+
+It is not a Transformers `AutoModel` checkpoint and is not instruction-formatted as a conventional chat model. Use the included cached inference script.
+
+![Three-level nested Mamba-2 architecture](architecture.svg)
+
+## Checkpoint contents
+
+The published weights are sharded SafeTensors containing only the 980 model tensors. The original optimizer, scaler, training phase, data cursor, dataset paths, source fingerprints, and other training-only checkpoint objects were removed.
+
+- Parameters: **2,478,820,575**
+- Weight precision on disk: **FP32**
+- Raw tensor size: **9,915,282,300 bytes**
+- Source checkpoint step: **889,000**
+- Recommended runtime precision: **BF16**
+- Recommended placement: fine/decoder on `cuda:0`, level 2 on `cuda:1`, level 3 on `cuda:2`
+
+The source checkpoint step is documentation only; it is not embedded in the SafeTensors weights or inference configuration.
+
+## Latest validation results
+
+The latest recorded validation event is step **890,000**, one scheduled validation event after the packaged `last.pt` weight step.
+
+| Validation stream | Cross entropy (nats/byte) | Bits per byte | Scored bytes |
+|---|---:|---:|---:|
+| Aggregate mixed validation | **3.828962** | **5.524025** | 14,530,840 |
+| JSONL text | 0.937809 | 1.352973 | 1,246,101 |
+| Parquet text | 0.820114 | 1.183175 | 1,929,612 |
+| Image + text multimodal | 1.397195 | 2.015726 | 1,626,324 |
+| Audio objectives | 5.189134 | 7.486338 | 9,824,803 |
+
+The aggregate should not be interpreted as a pure language score: audio accounts for most evaluated bytes and has a substantially different entropy scale. For text use, the JSONL and Parquet rows are the relevant measurements.
+
+## Architecture
+
+### Byte vocabulary
+
+There is no learned tokenizer:
+
+```text
+PAD=0, BOS=1, EOS=2, UNK=3
+raw byte 0..255 -> ID 4..259
+vocabulary size = 260
+```
+
+UTF-8 text and serialized binary modalities therefore share one next-byte objective.
+
+### Three causal resolutions
+
+1. **Fine level:** a local causal convolutional encoder and 6 Mamba-2 blocks operate at byte resolution. A learned causal boundary head closes variable pools between 1 and 96 bytes.
+2. **Level 2:** 20 Mamba-2 blocks consume completed fine-pool states. A learned boundary head groups 4–16 completed fine pools.
+3. **Level 3:** 30 Mamba-2 blocks consume completed level-2 states and group 2–16 level-2 pools.
+
+Every Mamba block uses model width 2,000, Mamba-2 `d_state=64`, and head dimension 100. A pool can use only states already available in its causal prefix. A closure never revises an earlier prediction.
+
+### Fusion decoder
+
+For each byte, the decoder concatenates four 2,000-dimensional signals:
+
+- byte-local contextual state;
+- current fine latent;
+- latest level-2 latent;
+- latest level-3 latent.
+
+The 10,000-dimensional concatenation is normalized, projected through an 8,000-wide GELU fusion layer, reduced to width 2,000, and mapped to 260 next-byte logits. This enlarged decoder was added to avoid choking the information arriving from three recurrent resolutions.
+
+### Pool-density fallback
+
+Fine pooling includes a rolling short-pool quota. Among the most recent 6,000 completed fine pools, at most 3,000 may be shorter than 6 bytes. When that quota fills, the next pool must reach the secondary minimum; short closures become eligible again as older short pools leave the rolling window. The quota counts completed pools, not raw bytes.
+
+### Delayed decoder controller
+
+The checkpoint includes an optional hold/refresh/compress controller. Its output at time `t` can influence closure only at `t+1`:
+
+```text
+decode byte t -> controller C[t] -> choose closure at t+1 -> decode byte t+1
+```
+
+The included cached inference path applies this without future leakage or a second full-model pass.
+
+### Parameter distribution
+
+| Component | Parameters |
+|---|---:|
+| Fine level, shared byte modules, decoder, and LM head | 411,954,571 |
+| Level 2 pooler and 20 Mamba-2 blocks | 831,559,202 |
+| Level 3 pooler and 30 Mamba-2 blocks | 1,235,306,802 |
+
+Pooling reduces sequence activations and recurrent update frequency, not layer-weight storage. This is why the deepest level remains the largest parameter group even though it updates least frequently.
+
+## Inference
+
+### Dependencies
+
+Use Linux, CUDA, and versions of PyTorch, `mamba-ssm`, Triton, and `causal-conv1d` that are mutually compatible:
+
+```bash
+pip install -r requirements.txt
+```
+
+BF16 is strongly recommended. FP16 cached rollouts can become numerically unstable on some Mamba-2 builds.
+
+### Three-GPU inference
+
+From the downloaded repository:
+
+```bash
+python infer_nested_model.py \
+  --checkpoint . \
+  --prompt "The history of state space models begins" \
+  --max-new-bytes 512 \
+  --precision bf16 \
+  --fine-device cuda:0 \
+  --nested-devices cuda:1 \
+  --tertiary-device cuda:2 \
+  --temperature 0.8 \
+  --top-p 0.9
+```
+
+The script accepts `--prompt-file` for arbitrary byte prefixes, `--output` for raw generated bytes, `--html-output` for hierarchy-attribution output, and `--image-output-dir` to extract complete generated P6 images.
+
+Single-GPU inference is supported when the GPU can hold the requested precision:
+
+```bash
+python infer_nested_model.py --checkpoint . --device cuda:0 \
+  --prompt "Once upon a time" --max-new-bytes 256 --precision bf16
+```
+
+### Stateful generation
+
+Generation prefills the prompt once, then caches the convolution and SSM states for the fine, level-2, and level-3 stacks. New bytes advance those caches token by token; the entire prefix is not reprocessed for every generated byte.
+
+## Training mixture and modality representation
+
+The training run mixed educational web text, PDF-derived text, image/question
+and instruction examples, and paired music/cover data from the following
+repositories:
+
+| Training source | Use in this model | Upstream licensing and rights notice |
+|---|---|---|
+| [HuggingFaceM4/FineVision](https://huggingface.co/datasets/HuggingFaceM4/FineVision) | Image, document, question, and instruction examples | FineVision is an aggregation. Each constituent dataset retains its own license; rights in prompts contributed by FineVision are offered under CC BY 4.0. Consult the license metadata for the constituent subsets. |
+| [HuggingFaceFW/finepdfs](https://huggingface.co/datasets/HuggingFaceFW/finepdfs) | PDF-derived document text | ODC-By 1.0; use is also subject to applicable Common Crawl terms and upstream-content rights. |
+| [HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) | Educational web text | ODC-By 1.0; source pages retain their applicable rights. |
+| [webshart/suno-various-94k](https://huggingface.co/datasets/webshart/suno-various-94k) | Music, captions, lyrics, and generated cover pairs | Marked `source-rights-retained`. Rights in source audio and lyrics remain with their creators; the dataset does not grant rights over the underlying content. |
+
+These datasets are not redistributed in this repository. Their upstream terms
+continue to apply independently and are not replaced by this repository's
+license.
+
+- Text and code are UTF-8 bytes.
+- Images are complete RGB PPM byte sequences plus associated text.
+- Audio uses 24 kHz EnCodec payloads with generation and detection objectives.
+- Instruction and dialogue fields present in source records were serialized in full rather than using assistant-response-only loss.
+
+This mixture makes the checkpoint experimental and general-purpose at the byte level; it does not guarantee strong image or audio generation quality.
+
+## Limitations
+
+- This is custom research code, not an official Mamba or Transformers architecture.
+- The model is not a safety-aligned chat assistant.
+- Raw-byte sampling can produce invalid UTF-8, malformed images, or incomplete audio containers.
+- Image training used small PPM rasters, limiting fine visual detail.
+- Audio validation remains much weaker than text validation.
+- The audio corpus includes third-party creator material whose source rights are retained. The model license does not grant rights to reproduce protected training content, lyrics, compositions, voices, or recordings.
+- The current weights are FP32 and large; practical use generally requires BF16 casting.
+- The delayed pooling controller is causal but makes exact routing inherently sequential.
+- The latest CSV validation event is at step 890,000, while the packaged `last.pt` weights identify step 889,000; the table must therefore be read as the latest run validation, not an evaluation re-run performed directly on this exported artifact.
+
+## Intended use
+
+Intended for research into byte-level modeling, hierarchical state-space models, adaptive causal pooling, long recurrent context, and mixed text/binary generation. Validate outputs independently before using them in downstream systems.
+
+## License
+
+The model weights, model card, and visual assets are available under
+[CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/). The Python
+inference source is available under
+[Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0). Training
+datasets and third-party content are not covered by either grant. See
+[LICENSE.md](LICENSE.md) for the precise repository scope and notices.
+
+## Repository files
+
+- `model-*.safetensors`: inference-only model shards
+- `model.safetensors.index.json`: tensor-to-shard map
+- `config.json`: architecture-only inference configuration
+- `modeling_nested_mamba.py`: custom model implementation
+- `nested_inference_tools.py`: SafeTensors loading, state caching, and sampling
+- `infer_nested_model.py`: command-line generator
+- `architecture.svg`: architecture visualization
+- `LICENSE.md`: weight, documentation, and code license scope

config.json

@@ -0,0 +1,49 @@
+{
+  "model_type": "nested_byte_mamba2",
+  "architectures": ["ForwardBackwardRepairModel"],
+  "architecture": "byte_latent_mamba_nested_jsonl",
+  "vocab_size": 260,
+  "byte_offset": 4,
+  "special_token_ids": {
+    "pad": 0,
+    "bos": 1,
+    "eos": 2,
+    "unk": 3
+  },
+  "dim": 2000,
+  "layers": 26,
+  "fine_layers": 6,
+  "nested_layers": 20,
+  "tertiary_layers": 30,
+  "decoder_dim": 8000,
+  "position_bins": 8192,
+  "mamba_version": 2,
+  "mamba_d_state": 64,
+  "mamba2_headdim": 100,
+  "no_mamba": false,
+  "blt_min_patch_bytes": 1,
+  "blt_max_patch_bytes": 96,
+  "blt_patch_change_threshold": 48,
+  "blt_close_threshold": 0.98,
+  "blt_mid_close_bonus": 0.05,
+  "blt_short_pool_budget": 3000,
+  "blt_short_pool_window": 6000,
+  "blt_secondary_min_patch_bytes": 6,
+  "nested_pool_factor": 16,
+  "nested_min_pool_factor": 4,
+  "nested_close_threshold": 0.98,
+  "tertiary_pool_factor": 16,
+  "tertiary_min_pool_factor": 2,
+  "tertiary_close_threshold": 0.98,
+  "detach_inactive_coarse_gradients": true,
+  "decoder_pool_controller": true,
+  "pool_controller_alpha": 1.0,
+  "pool_controller_beta": 0.5,
+  "pool_controller_gamma": 1.0,
+  "recommended_precision": "bf16",
+  "recommended_placement": {
+    "fine_device": "cuda:0",
+    "nested_devices": ["cuda:1"],
+    "tertiary_device": "cuda:2"
+  }
+}

infer_nested_model.py

@@ -0,0 +1,463 @@
+#!/usr/bin/env python3
+"""Cached byte generation for two- and three-level nested Mamba checkpoints."""
+from __future__ import annotations
+
+import argparse
+import base64
+import gc
+import io
+import json
+import time
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import torch
+
+from nested_inference_tools import generate_bytes, load_nested_checkpoint
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--checkpoint",
+        default=".",
+        help="Hugging Face model directory (default: current directory) or legacy last.pt.",
+    )
+    prompt = parser.add_mutually_exclusive_group()
+    prompt.add_argument("--prompt", default="", help="UTF-8 prompt text.")
+    prompt.add_argument("--prompt-file", help="Read prompt bytes from this file.")
+    parser.add_argument("--max-new-bytes", type=int, default=256)
+    parser.add_argument("--temperature", type=float, default=0.8)
+    parser.add_argument("--top-p", type=float, default=0.9)
+    parser.add_argument("--top-k", type=int, default=0)
+    parser.add_argument("--repeat-penalty", type=float, default=1.05)
+    parser.add_argument("--repeat-window", type=int, default=256)
+    parser.add_argument("--greedy", action="store_true")
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument(
+        "--precision",
+        choices=["fp16", "bf16", "fp32"],
+        default="bf16",
+        help=(
+            "Model weight/activation precision. BF16 is the safe default for "
+            "Mamba-2's long cached rollouts; use FP16 only for checkpoints and "
+            "GPUs verified to remain finite."
+        ),
+    )
+    parser.add_argument("--device", default="cuda:0", help="Single-device inference target.")
+    parser.add_argument("--fine-device", default=None)
+    parser.add_argument("--nested-devices", default=None)
+    parser.add_argument("--tertiary-device", default=None)
+    parser.add_argument("--use-saved-placement", action="store_true")
+    parser.add_argument("--output", help="Optional raw-byte output file.")
+    parser.add_argument("--stats-json", help="Optional JSON statistics output.")
+    parser.add_argument(
+        "--image-output-dir",
+        help=(
+            "Extract complete generated P6 PPM images, convert them to PNG, and "
+            "write them to this directory. Images are also embedded in --html-output."
+        ),
+    )
+    parser.add_argument("--html-output", "--output-html", dest="html_output", help="Write a self-contained prompt/response report with toggleable L2/L3 influence coloring.")
+    return parser.parse_args()
+
+
+def _ppm_token(data: bytes, position: int) -> tuple[bytes, int]:
+    """Read one whitespace/comment-delimited PPM header token."""
+    size = len(data)
+    while position < size:
+        if data[position] in b" \t\r\n":
+            position += 1
+            continue
+        if data[position] == ord("#"):
+            newline = data.find(b"\n", position)
+            if newline < 0:
+                raise ValueError("unterminated PPM header comment")
+            position = newline + 1
+            continue
+        break
+    start = position
+    while position < size and data[position] not in b" \t\r\n#":
+        position += 1
+    if position == start:
+        raise ValueError("missing PPM header token")
+    return data[start:position], position
+
+
+def extract_ppm_images(data: bytes) -> tuple[List[Dict[str, object]], List[str]]:
+    """Extract complete P6 RGB images from an arbitrary generated byte stream."""
+    try:
+        from PIL import Image
+    except ImportError as error:
+        return [], [f"Pillow is required to render generated images: {error}"]
+
+    images: List[Dict[str, object]] = []
+    warnings: List[str] = []
+    search_from = 0
+    while len(images) < 32:
+        start = data.find(b"P6", search_from)
+        if start < 0:
+            break
+        search_from = start + 2
+        if start > 0 and data[start - 1] not in b" \t\r\n>":
+            continue
+        if start + 2 >= len(data) or data[start + 2] not in b" \t\r\n":
+            continue
+        try:
+            width_token, position = _ppm_token(data, start + 2)
+            height_token, position = _ppm_token(data, position)
+            maximum_token, position = _ppm_token(data, position)
+            width, height, maximum = int(width_token), int(height_token), int(maximum_token)
+            if width < 1 or height < 1 or width * height > 16_777_216:
+                raise ValueError(f"unsafe dimensions {width}x{height}")
+            if maximum != 255:
+                raise ValueError(f"unsupported maximum channel value {maximum}; expected 255")
+            if position >= len(data) or data[position] not in b" \t\r\n":
+                raise ValueError("missing whitespace before PPM pixel payload")
+            # The delimiter is one whitespace unit. Treat CRLF as one unit so
+            # the first pixel is never shifted on Windows-produced headers.
+            pixel_start = position + 1
+            if data[position] == ord("\r") and pixel_start < len(data) and data[pixel_start] == ord("\n"):
+                pixel_start += 1
+            pixel_bytes = width * height * 3
+            pixel_end = pixel_start + pixel_bytes
+            if pixel_end > len(data):
+                warnings.append(
+                    f"Incomplete PPM at byte {start}: {width}x{height} needs "
+                    f"{pixel_bytes:,} RGB bytes, but only {max(0, len(data)-pixel_start):,} remain."
+                )
+                continue
+            image = Image.frombytes("RGB", (width, height), data[pixel_start:pixel_end])
+            encoded = io.BytesIO()
+            image.save(encoded, format="PNG", optimize=True)
+            images.append(
+                {
+                    "width": width,
+                    "height": height,
+                    "start": start,
+                    "end": pixel_end,
+                    "png": encoded.getvalue(),
+                }
+            )
+            search_from = pixel_end
+        except (TypeError, ValueError) as error:
+            warnings.append(f"Invalid PPM candidate at byte {start}: {error}.")
+    return images, warnings
+
+
+def image_report_payload(images: List[Dict[str, object]], warnings: List[str]) -> Dict[str, object]:
+    return {
+        "images": [
+            {
+                "width": int(item["width"]),
+                "height": int(item["height"]),
+                "start": int(item["start"]),
+                "end": int(item["end"]),
+                "data_uri": "data:image/png;base64," + base64.b64encode(item["png"]).decode("ascii"),
+            }
+            for item in images
+        ],
+        "warnings": list(warnings),
+    }
+
+
+def write_extracted_images(
+    directory: Path, images: List[Dict[str, object]], *, prefix: str
+) -> List[str]:
+    directory.mkdir(parents=True, exist_ok=True)
+    paths: List[str] = []
+    for index, item in enumerate(images, 1):
+        path = directory / f"{prefix}_{index:03d}_{item['width']}x{item['height']}.png"
+        path.write_bytes(item["png"])
+        paths.append(str(path))
+    return paths
+
+
+def attributed_utf8_segments(
+    data: bytes, attributions: List[Dict[str, object]]
+) -> List[Dict[str, object]]:
+    """Group byte attribution into valid UTF-8 characters without losing data."""
+    segments: List[Dict[str, object]] = []
+    index = 0
+    while index < len(data):
+        first = data[index]
+        width = (
+            1
+            if first < 0x80
+            else 2
+            if 0xC2 <= first <= 0xDF
+            else 3
+            if 0xE0 <= first <= 0xEF
+            else 4
+            if 0xF0 <= first <= 0xF4
+            else 1
+        )
+        chunk = data[index : index + width]
+        try:
+            text = chunk.decode("utf-8", "strict")
+        except UnicodeDecodeError:
+            width = 1
+            chunk = data[index : index + 1]
+            text = chunk.decode("utf-8", "replace")
+        values = attributions[index : index + width]
+        count = max(1, len(values))
+        segments.append(
+            {
+                "text": text,
+                "bytes": list(chunk),
+                "level2_active": any(bool(item.get("level2_active")) for item in values),
+                "level3_active": any(bool(item.get("level3_active")) for item in values),
+                "level2_delta_logp": sum(
+                    float(item.get("level2_delta_logp", 0.0)) for item in values
+                )
+                / count,
+                "level3_delta_logp": sum(
+                    float(item.get("level3_delta_logp", 0.0)) for item in values
+                )
+                / count,
+                "hierarchy_delta_logp": sum(
+                    float(item.get("hierarchy_delta_logp", 0.0)) for item in values
+                )
+                / count,
+            }
+        )
+        index += width
+    return segments
+
+
+def make_generation_html(
+    prompt: bytes,
+    generated: bytes,
+    attributions: List[Dict[str, object]],
+    stats: Dict[str, object],
+    comparisons: Optional[List[Dict[str, object]]] = None,
+    image_report: Optional[Dict[str, object]] = None,
+) -> str:
+    segments = attributed_utf8_segments(generated, attributions)
+    payload = json.dumps(
+        {
+            "prompt": prompt.decode("utf-8", "replace"),
+            "segments": segments,
+            "stats": stats,
+            "comparisons": comparisons or [],
+            "image_report": image_report or {"images": [], "warnings": []},
+        },
+        separators=(",", ":"),
+    ).replace("</", "<\\/")
+    return f"""<!doctype html>
+<html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
+<title>Nested Mamba generation influence</title>
+<style>
+:root{{--bg:#090d17;--panel:#111a2b;--line:#273652;--text:#e7edf9;--muted:#91a0b8;--prompt:#d6deec;--response:#63a4ff;--fine:#51d6ca;--l2:#ffbe55;--l3:#a78bfa;--negative:#fb7185}}
+*{{box-sizing:border-box}}body{{margin:0;background:linear-gradient(145deg,#080c15,#11192b);color:var(--text);font:14px/1.5 system-ui,sans-serif}}main{{max-width:1200px;margin:auto;padding:28px}}
+h1{{margin:0 0 6px;font-size:27px}}p{{color:var(--muted)}}.toolbar,.panel{{background:#111a2bf2;border:1px solid var(--line);border-radius:12px}}
+.toolbar{{display:flex;align-items:center;gap:18px;flex-wrap:wrap;padding:13px 16px;margin:18px 0}}label{{display:flex;align-items:center;gap:8px;font-weight:650}}input{{accent-color:var(--l3)}}
+.legend{{display:flex;gap:14px;flex-wrap:wrap;color:var(--muted);font-size:12px}}.dot{{width:10px;height:10px;border-radius:50%;display:inline-block;margin-right:5px}}
+.panel{{padding:20px}}.text{{white-space:pre-wrap;overflow-wrap:anywhere;font:16px/1.65 ui-monospace,SFMono-Regular,Menlo,Consolas,monospace}}.responses{{display:grid;grid-template-columns:repeat(3,minmax(0,1fr));gap:12px;margin-top:16px}}.response-card{{min-width:0;background:#0c1321;border:1px solid var(--line);border-radius:10px;padding:14px}}.response-card h2{{font:700 14px system-ui,sans-serif;margin:0 0 4px}}.response-meta{{color:var(--muted);font:11px system-ui,sans-serif;margin-bottom:11px}}
+.prompt{{color:var(--prompt)}}.generated{{color:var(--response);transition:color .15s,opacity .15s}}.divider{{display:block;color:var(--muted);font:12px system-ui,sans-serif;margin:18px 0 7px;border-top:1px solid var(--line);padding-top:10px}}
+.cards{{display:grid;grid-template-columns:repeat(auto-fit,minmax(170px,1fr));gap:10px;margin-top:14px}}.card{{background:#0c1321;border:1px solid var(--line);border-radius:9px;padding:11px}}.name{{color:var(--muted);font-size:11px;text-transform:uppercase}}.value{{font-size:18px;font-weight:700}}
+.note{{border-left:3px solid var(--l2);padding:9px 12px;background:#ffbe550d}}
+.image-panel{{margin-top:16px}}.image-galleries{{display:grid;grid-template-columns:repeat(3,minmax(0,1fr));gap:12px}}.image-gallery{{background:#0c1321;border:1px solid var(--line);border-radius:10px;padding:12px;min-width:0}}.image-gallery h2{{font-size:14px;margin:0 0 9px}}.image-item{{margin:0 0 12px}}.image-item img{{display:block;max-width:100%;height:auto;image-rendering:auto;border:1px solid #354766;background:#05070b}}.image-item figcaption,.image-status{{color:var(--muted);font-size:11px;margin-top:5px;white-space:pre-wrap}}
+@media(max-width:900px){{.responses{{grid-template-columns:1fr}}}}
+</style></head><body><main>
+<h1>Nested Mamba generation influence</h1><p id="subtitle"></p>
+<div class="toolbar"><label><input id="influence" type="checkbox" checked> Color generated text by hierarchy influence</label><div class="legend"><span><i class="dot" style="background:var(--fine)"></i>fine only</span><span><i class="dot" style="background:var(--l2)"></i>L2 dominant</span><span><i class="dot" style="background:var(--l3)"></i>L3 dominant</span><span><i class="dot" style="background:var(--negative)"></i>dominant branch reduced selected-byte probability</span></div></div>
+<p class="note">Color is a decoder counterfactual, not merely an activation marker. Positive delta means the cached dynamic level increased the selected byte's log probability relative to replacing that level with its beginning-of-document state. Hover text for exact values.</p>
+<section class="panel"><div class="text"><span class="divider">PROMPT SHARED BY ALL ROLLOUTS</span><span class="prompt" id="prompt"></span></div><div class="responses"><article class="response-card"><h2>Level 1 only</h2><div class="response-meta" id="level1Meta"></div><div class="text generated" id="level1Response"></div></article><article class="response-card"><h2>Levels 1 + 2</h2><div class="response-meta" id="level12Meta"></div><div class="text generated" id="level12Response"></div></article><article class="response-card"><h2>Levels 1 + 2 + 3</h2><div class="response-meta" id="fullMeta"></div><div class="text" id="response"></div></article></div><div class="cards" id="cards"></div></section>
+<section class="panel image-panel"><h1>Generated image rendering</h1><p>Complete P6 RGB byte sequences found in each prompt-plus-rollout stream are validated and embedded here as PNG.</p><div class="image-galleries"><div class="image-gallery"><h2>Level 1 only</h2><div id="level1Images"></div></div><div class="image-gallery"><h2>Levels 1 + 2</h2><div id="level12Images"></div></div><div class="image-gallery"><h2>Levels 1 + 2 + 3</h2><div id="fullImages"></div></div></div></section>
+<script>
+const R={payload},prompt=document.getElementById("prompt"),response=document.getElementById("response"),toggle=document.getElementById("influence");
+prompt.textContent=R.prompt;document.getElementById("subtitle").textContent=`${{R.stats.checkpoint}} • step ${{Number(R.stats.checkpoint_step).toLocaleString()}}`;
+const byMode=Object.fromEntries(R.comparisons.map(x=>[x.mode,x]));for(const [mode,id,color] of [["level1","level1Response","#51d6ca"],["level12","level12Response","#ffbe55"]]){{const item=byMode[mode]||{{text:"",generated_bytes:0,seconds:0}};document.getElementById(id).textContent=item.text;document.getElementById(id).dataset.color=color;document.getElementById(mode+"Meta").textContent=`${{Number(item.generated_bytes).toLocaleString()}} bytes • ${{Number(item.seconds).toFixed(2)}} s`;}}document.getElementById("fullMeta").textContent=`${{Number(R.stats.generated_bytes).toLocaleString()}} bytes • ${{Number(R.stats.generation_seconds_including_prefill).toFixed(2)}} s`;
+const magnitudes=R.segments.flatMap(x=>[Math.abs(x.level2_delta_logp),Math.abs(x.level3_delta_logp)]).filter(Number.isFinite).sort((a,b)=>a-b),scale=magnitudes[Math.floor(magnitudes.length*.9)]||.01;
+function classification(x){{const candidates=[];if(x.level2_active)candidates.push([Math.abs(x.level2_delta_logp),x.level2_delta_logp,"#ffbe55","L2"]);if(x.level3_active)candidates.push([Math.abs(x.level3_delta_logp),x.level3_delta_logp,"#a78bfa","L3"]);if(!candidates.length)return ["#51d6ca",.72,"fine only"];candidates.sort((a,b)=>b[0]-a[0]);const best=candidates[0],color=best[1]<0?"#fb7185":best[2],opacity=.48+.52*Math.min(1,best[0]/scale);return [color,opacity,best[1]<0?best[3]+" negative":best[3]+" dominant"]}}
+R.segments.forEach(x=>{{const span=document.createElement("span"),style=classification(x);span.className="generated";span.textContent=x.text;span.dataset.color=style[0];span.dataset.opacity=style[1];span.title=`${{style[2]}} • bytes ${{x.bytes.join(",")}} • L2 Δlogp ${{x.level2_delta_logp.toFixed(5)}} • L3 Δlogp ${{x.level3_delta_logp.toFixed(5)}} • combined Δlogp ${{x.hierarchy_delta_logp.toFixed(5)}}`;response.appendChild(span)}});
+function recolor(){{response.querySelectorAll(".generated").forEach(span=>{{span.style.color=toggle.checked?span.dataset.color:"#63a4ff";span.style.opacity=toggle.checked?span.dataset.opacity:"1"}});for(const id of ["level1Response","level12Response"]){{const node=document.getElementById(id);node.style.color=toggle.checked?node.dataset.color:"#63a4ff"}}}}toggle.addEventListener("change",recolor);recolor();
+const attrs=R.segments,mean=key=>attrs.length?attrs.reduce((s,x)=>s+Number(x[key]||0),0)/attrs.length:0,cards=[["Prompt bytes",R.stats.prompt_bytes],["Generated bytes",R.stats.generated_bytes],["Generation tok/s",Number(R.stats.generated_bytes_per_second_including_prefill).toFixed(1)],["Mean L2 Δlogp",mean("level2_delta_logp").toFixed(5)],["Mean L3 Δlogp",mean("level3_delta_logp").toFixed(5)],["Mean combined Δlogp",mean("hierarchy_delta_logp").toFixed(5)]];
+document.getElementById("cards").innerHTML=cards.map(x=>`<div class="card"><div class="name">${{x[0]}}</div><div class="value">${{x[1]}}</div></div>`).join("");
+function renderImages(target,report){{const root=document.getElementById(target),images=report?.images||[],warnings=report?.warnings||[];if(!images.length&&!warnings.length){{root.innerHTML='<div class="image-status">No complete P6 image detected.</div>';return}}images.forEach((item,index)=>{{const figure=document.createElement("figure");figure.className="image-item";const image=document.createElement("img");image.src=item.data_uri;image.alt=`Generated image ${{index+1}}, ${{item.width}} by ${{item.height}} pixels`;const caption=document.createElement("figcaption");caption.textContent=`Image ${{index+1}} • ${{item.width}}×${{item.height}} • byte range ${{item.start.toLocaleString()}}–${{(item.end-1).toLocaleString()}}`;figure.append(image,caption);root.appendChild(figure)}});if(warnings.length){{const status=document.createElement("div");status.className="image-status";status.textContent=warnings.join("\\n");root.appendChild(status)}}}}
+renderImages("fullImages",R.image_report);renderImages("level1Images",byMode.level1?.image_report);renderImages("level12Images",byMode.level12?.image_report);
+</script></main></body></html>"""
+
+
+def main() -> None:
+    args = parse_args()
+    if args.max_new_bytes < 0 or args.temperature <= 0 or not 0 < args.top_p <= 1:
+        raise ValueError("max-new-bytes must be non-negative, temperature positive, and top-p in (0, 1]")
+    if args.prompt_file:
+        prompt = Path(args.prompt_file).expanduser().read_bytes()
+    else:
+        prompt = args.prompt.encode("utf-8")
+
+    load_started = time.perf_counter()
+    loaded = load_nested_checkpoint(
+        args.checkpoint,
+        precision=args.precision,
+        device=args.device,
+        fine_device=args.fine_device,
+        nested_devices=args.nested_devices,
+        tertiary_device=args.tertiary_device,
+        use_saved_placement=args.use_saved_placement,
+    )
+    torch.cuda.synchronize()
+    load_seconds = time.perf_counter() - load_started
+    generation_started = time.perf_counter()
+    generated, stream = generate_bytes(
+        loaded,
+        prompt,
+        max_new_bytes=args.max_new_bytes,
+        temperature=args.temperature,
+        top_p=args.top_p,
+        top_k=args.top_k,
+        repeat_penalty=args.repeat_penalty,
+        repeat_window=args.repeat_window,
+        greedy=args.greedy,
+        seed=args.seed,
+        collect_hierarchy_attribution=bool(args.html_output),
+    )
+    torch.cuda.synchronize()
+    generation_seconds = time.perf_counter() - generation_started
+    combined = prompt + generated
+    full_images, full_image_warnings = extract_ppm_images(combined)
+    full_image_report = image_report_payload(full_images, full_image_warnings)
+    print(combined.decode("utf-8", "replace"))
+    attributions = list(stream.get("generated_attribution", []))
+    stats = {
+        "checkpoint": str(loaded.checkpoint_path),
+        "checkpoint_step": loaded.checkpoint_step,
+        "trained_tokens": loaded.trained_tokens,
+        "precision": loaded.precision,
+        "model_parallel": loaded.model_parallel,
+        "prompt_bytes": len(prompt),
+        "generated_bytes": len(generated),
+        "load_seconds": load_seconds,
+        "generation_seconds_including_prefill": generation_seconds,
+        "generated_bytes_per_second_including_prefill": (
+            len(generated) / generation_seconds if generation_seconds else 0.0
+        ),
+        "fine_patches": int(stream["completed_patches"]),
+        "level2_pools": int(stream["completed_nested_patches"]),
+        "level3_pools": int(stream.get("completed_tertiary_patches", 0)),
+        "hierarchy_attribution": bool(args.html_output),
+        "rendered_images": len(full_images),
+        "image_render_warnings": full_image_warnings,
+    }
+    if attributions:
+        stats.update(
+            {
+                "attributed_bytes": len(attributions),
+                "level2_active_generated_bytes": sum(
+                    bool(item.get("level2_active")) for item in attributions
+                ),
+                "level3_active_generated_bytes": sum(
+                    bool(item.get("level3_active")) for item in attributions
+                ),
+                "mean_level2_delta_logp": sum(
+                    float(item.get("level2_delta_logp", 0.0))
+                    for item in attributions
+                )
+                / len(attributions),
+                "mean_level3_delta_logp": sum(
+                    float(item.get("level3_delta_logp", 0.0))
+                    for item in attributions
+                )
+                / len(attributions),
+                "mean_hierarchy_delta_logp": sum(
+                    float(item.get("hierarchy_delta_logp", 0.0))
+                    for item in attributions
+                )
+                / len(attributions),
+            }
+        )
+    comparisons: List[Dict[str, object]] = []
+    del stream
+    gc.collect()
+    if args.html_output:
+        for mode, label in (("level1", "Level 1 only"), ("level12", "Levels 1 + 2")):
+            comparison_started = time.perf_counter()
+            comparison_bytes, comparison_stream = generate_bytes(
+                loaded,
+                prompt,
+                max_new_bytes=args.max_new_bytes,
+                temperature=args.temperature,
+                top_p=args.top_p,
+                top_k=args.top_k,
+                repeat_penalty=args.repeat_penalty,
+                repeat_window=args.repeat_window,
+                greedy=args.greedy,
+                seed=args.seed,
+                hierarchy_mode=mode,
+            )
+            torch.cuda.synchronize()
+            comparison_seconds = time.perf_counter() - comparison_started
+            comparisons.append(
+                {
+                    "mode": mode,
+                    "label": label,
+                    "text": comparison_bytes.decode("utf-8", "replace"),
+                    "generated_bytes": len(comparison_bytes),
+                    "seconds": comparison_seconds,
+                    "fine_patches": int(comparison_stream["completed_patches"]),
+                    "level2_pools": int(
+                        comparison_stream["completed_nested_patches"]
+                    ),
+                    "level3_pools": int(
+                        comparison_stream.get("completed_tertiary_patches", 0)
+                    ),
+                    "image_report": image_report_payload(
+                        *extract_ppm_images(prompt + comparison_bytes)
+                    ),
+                }
+            )
+            del comparison_stream
+            gc.collect()
+        # Keep base64 PNG payloads in the HTML only; terminal/stats JSON should
+        # stay compact even when a rollout contains a large rendered image.
+        stats["comparison_rollouts"] = [
+            {key: value for key, value in item.items() if key != "image_report"}
+            for item in comparisons
+        ]
+    if args.output:
+        Path(args.output).expanduser().write_bytes(combined)
+    if args.image_output_dir:
+        image_paths = write_extracted_images(
+            Path(args.image_output_dir).expanduser(), full_images, prefix="full"
+        )
+        stats["image_output_files"] = image_paths
+        for item in comparisons:
+            report = item.get("image_report") or {}
+            comparison_images = []
+            for encoded in report.get("images", []):
+                raw = base64.b64decode(str(encoded["data_uri"]).split(",", 1)[1])
+                comparison_images.append({**encoded, "png": raw})
+            stats.setdefault("comparison_image_output_files", {})[item["mode"]] = (
+                write_extracted_images(
+                    Path(args.image_output_dir).expanduser(),
+                    comparison_images,
+                    prefix=str(item["mode"]),
+                )
+            )
+    print(json.dumps(stats, indent=2))
+    if args.stats_json:
+        Path(args.stats_json).expanduser().write_text(
+            json.dumps(stats, indent=2), encoding="utf-8"
+        )
+    if args.html_output:
+        html_path = Path(args.html_output).expanduser()
+        html_path.parent.mkdir(parents=True, exist_ok=True)
+        html_path.write_text(
+            make_generation_html(
+                prompt,
+                generated,
+                attributions,
+                stats,
+                comparisons=comparisons,
+                image_report=full_image_report,
+            ),
+            encoding="utf-8",
+        )
+        print(f"wrote hierarchy influence report: {html_path}")
+
+
+if __name__ == "__main__":
+    main()

model.safetensors.index.json

@@ -0,0 +1,987 @@
+{
+  "metadata": {
+    "total_size": 9915282300
+  },
+  "weight_map": {
+    "forward_model.boe_patch": "model-00001-of-00003.safetensors",
+    "forward_model.byte_emb.weight": "model-00001-of-00003.safetensors",
+    "forward_model.decoder.0.bias": "model-00003-of-00003.safetensors",
+    "forward_model.decoder.0.weight": "model-00003-of-00003.safetensors",
+    "forward_model.decoder.1.bias": "model-00003-of-00003.safetensors",
+    "forward_model.decoder.1.weight": "model-00003-of-00003.safetensors",
+    "forward_model.decoder.3.bias": "model-00003-of-00003.safetensors",
+    "forward_model.decoder.3.weight": "model-00003-of-00003.safetensors",
+    "forward_model.global_blocks.0.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.0.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.1.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.2.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.3.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.4.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.global_blocks.5.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.lm_head.weight": "model-00003-of-00003.safetensors",
+    "forward_model.local_conv.0.conv.bias": "model-00001-of-00003.safetensors",
+    "forward_model.local_conv.0.conv.weight": "model-00001-of-00003.safetensors",
+    "forward_model.local_conv.2.conv.bias": "model-00001-of-00003.safetensors",
+    "forward_model.local_conv.2.conv.weight": "model-00001-of-00003.safetensors",
+    "forward_model.local_norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.local_norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.0.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.1.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.10.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.11.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.12.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.13.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.14.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.15.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.16.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.17.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.18.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.19.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.2.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.3.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.4.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.5.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.6.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.7.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.8.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.ff.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.ff.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.ff.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.ff.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.ff.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.ff.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.mixer.A_log": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.mixer.D": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.mixer.conv1d.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.mixer.conv1d.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.mixer.dt_bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.mixer.in_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.mixer.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.mixer.out_proj.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_global_blocks.9.norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_len_emb.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.close.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.close.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.close.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.close.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.close.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.close.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.close_value.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.context.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.context.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.context_norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.context_norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.value.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.value.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.value.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.value.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.value.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.value.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.weight.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.weight.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.weight.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.nested_pool_encoder.weight.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.patch_len_emb.weight": "model-00001-of-00003.safetensors",
+    "forward_model.patch_pos_emb.weight": "model-00001-of-00003.safetensors",
+    "forward_model.pool_controller.bias": "model-00003-of-00003.safetensors",
+    "forward_model.pool_controller.weight": "model-00003-of-00003.safetensors",
+    "forward_model.pool_encoder.close.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.close.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.close.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.close.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.close.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.close.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.close_value.weight": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.context.bias": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.context.weight": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.context_norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.context_norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.value.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.value.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.value.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.value.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.value.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.value.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.weight.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.weight.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.weight.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.pool_encoder.weight.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.0.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.1.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.10.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.11.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.12.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.13.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.14.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.15.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.16.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.17.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.18.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.19.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.2.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.ff.0.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.ff.0.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.ff.1.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.ff.1.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.ff.3.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.ff.3.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.mixer.conv1d.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.mixer.conv1d.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.mixer.in_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.mixer.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.mixer.out_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.20.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.ff.0.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.ff.0.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.ff.1.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.ff.1.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.ff.3.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.ff.3.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.mixer.A_log": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.mixer.D": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.mixer.conv1d.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.mixer.conv1d.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.mixer.dt_bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.mixer.in_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.mixer.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.mixer.out_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.norm.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.21.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.ff.0.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.ff.0.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.ff.1.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.ff.1.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.ff.3.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.ff.3.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.mixer.A_log": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.mixer.D": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.mixer.conv1d.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.mixer.conv1d.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.mixer.dt_bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.mixer.in_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.mixer.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.mixer.out_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.norm.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.22.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.ff.0.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.ff.0.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.ff.1.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.ff.1.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.ff.3.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.ff.3.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.mixer.A_log": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.mixer.D": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.mixer.conv1d.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.mixer.conv1d.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.mixer.dt_bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.mixer.in_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.mixer.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.mixer.out_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.norm.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.23.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.ff.0.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.ff.0.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.ff.1.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.ff.1.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.ff.3.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.ff.3.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.mixer.A_log": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.mixer.D": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.mixer.conv1d.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.mixer.conv1d.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.mixer.dt_bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.mixer.in_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.mixer.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.mixer.out_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.norm.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.24.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.ff.0.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.ff.0.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.ff.1.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.ff.1.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.ff.3.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.ff.3.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.mixer.A_log": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.mixer.D": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.mixer.conv1d.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.mixer.conv1d.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.mixer.dt_bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.mixer.in_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.mixer.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.mixer.out_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.norm.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.25.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.ff.0.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.ff.0.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.ff.1.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.ff.1.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.ff.3.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.ff.3.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.mixer.A_log": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.mixer.D": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.mixer.conv1d.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.mixer.conv1d.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.mixer.dt_bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.mixer.in_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.mixer.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.mixer.out_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.norm.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.26.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.ff.0.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.ff.0.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.ff.1.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.ff.1.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.ff.3.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.ff.3.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.mixer.A_log": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.mixer.D": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.mixer.conv1d.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.mixer.conv1d.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.mixer.dt_bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.mixer.in_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.mixer.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.mixer.out_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.norm.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.27.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.ff.0.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.ff.0.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.ff.1.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.ff.1.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.ff.3.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.ff.3.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.mixer.A_log": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.mixer.D": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.mixer.conv1d.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.mixer.conv1d.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.mixer.dt_bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.mixer.in_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.mixer.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.mixer.out_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.norm.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.28.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.ff.0.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.ff.0.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.ff.1.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.ff.1.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.ff.3.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.ff.3.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.mixer.A_log": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.mixer.D": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.mixer.conv1d.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.mixer.conv1d.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.mixer.dt_bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.mixer.in_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.mixer.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.mixer.out_proj.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.norm.bias": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.29.norm.weight": "model-00003-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.3.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.4.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.5.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.6.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.7.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.8.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.ff.0.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.ff.0.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.ff.1.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.ff.1.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.ff.3.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.ff.3.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.mixer.A_log": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.mixer.D": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.mixer.conv1d.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.mixer.conv1d.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.mixer.dt_bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.mixer.in_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.mixer.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.mixer.out_proj.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.norm.bias": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_global_blocks.9.norm.weight": "model-00002-of-00003.safetensors",
+    "forward_model.tertiary_len_emb.weight": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.close.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.close.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.close.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.close.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.close.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.close.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.close_value.weight": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.context.bias": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.context.weight": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.context_norm.bias": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.context_norm.weight": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.value.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.value.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.value.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.value.1.weight": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.value.3.bias": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.value.3.weight": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.weight.0.bias": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.weight.0.weight": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.weight.1.bias": "model-00001-of-00003.safetensors",
+    "forward_model.tertiary_pool_encoder.weight.1.weight": "model-00001-of-00003.safetensors"
+  }
+}

nested_inference_tools.py

@@ -0,0 +1,557 @@
+#!/usr/bin/env python3
+"""Shared loading, streaming, and sampling helpers for nested byte Mamba tools."""
+from __future__ import annotations
+
+import gc
+import json
+import math
+import random
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from modeling_nested_mamba import (
+    ForwardBackwardRepairModel,
+)
+
+
+PAD = 0
+BOS = 1
+EOS = 2
+UNK = 3
+BYTE_OFFSET = 4
+VOCAB_SIZE = 260
+
+
+@dataclass
+class LoadedNestedModel:
+    model: ForwardBackwardRepairModel
+    config: Dict[str, object]
+    checkpoint_step: int
+    trained_tokens: int
+    checkpoint_path: Path
+    precision: str
+    fine_device: torch.device
+    model_parallel: bool
+
+    @property
+    def core(self):
+        return self.model.forward_model
+
+
+def _dtype_for_precision(precision: str) -> torch.dtype:
+    normalized = precision.lower()
+    if normalized == "fp16":
+        return torch.float16
+    if normalized == "bf16":
+        return torch.bfloat16
+    if normalized == "fp32":
+        return torch.float32
+    raise ValueError(f"unsupported precision {precision!r}; choose fp16, bf16, or fp32")
+
+
+def _mamba2_fused_causal_conv_available() -> bool:
+    """Whether Mamba-2's combined full-sequence kernel can call causal-conv1d."""
+    try:
+        from mamba_ssm.ops.triton import ssd_combined
+
+        return getattr(ssd_combined, "causal_conv1d_fwd_function", None) is not None
+    except Exception:
+        return False
+
+
+def _configure_mamba2_inference_kernels(model: ForwardBackwardRepairModel) -> bool:
+    """Select the portable full-sequence path when fused causal-conv1d is absent."""
+    fused_available = _mamba2_fused_causal_conv_available()
+    changed = False
+    for block_group in (
+        model.forward_model.global_blocks,
+        model.forward_model.nested_global_blocks,
+        model.forward_model.tertiary_global_blocks,
+    ):
+        for block in block_group:
+            if int(getattr(block, "mamba_version", 1)) != 2:
+                continue
+            if not fused_available and bool(getattr(block.mixer, "use_mem_eff_path", False)):
+                block.mixer.use_mem_eff_path = False
+                changed = True
+    return changed
+
+
+def _build_model(config: Dict[str, object]) -> ForwardBackwardRepairModel:
+    if config.get("training_method") == "diffusionblocks":
+        raise ValueError(
+            "this checkpoint declares an architecture that is incompatible with "
+            "the nested Mamba inference implementation"
+        )
+    model = ForwardBackwardRepairModel(
+        vocab_size=int(config.get("vocab_size", VOCAB_SIZE)),
+        dim=int(config["dim"]),
+        layers=int(config["layers"]),
+        position_bins=int(config.get("position_bins", 8192)),
+        use_mamba=not bool(config.get("no_mamba", False)),
+        # Old checkpoints predate these fields and are Mamba-1/d_state=16.
+        mamba_version=int(config.get("mamba_version", 1)),
+        mamba_d_state=int(
+            config.get(
+                "mamba_d_state",
+                64 if int(config.get("mamba_version", 1)) == 2 else 16,
+            )
+        ),
+        mamba2_headdim=int(config.get("mamba2_headdim", 0)),
+        num_sections=1,
+        min_patch_bytes=int(config.get("blt_min_patch_bytes", 16)),
+        max_patch_bytes=int(config.get("blt_max_patch_bytes", 96)),
+        patch_change_threshold=int(config.get("blt_patch_change_threshold", 48)),
+        close_threshold=float(config.get("blt_close_threshold", 0.98)),
+        mid_close_bonus=float(config.get("blt_mid_close_bonus", 0.05)),
+        nested_pool_factor=int(config.get("nested_pool_factor", 16)),
+        nested_min_pool_factor=int(config.get("nested_min_pool_factor", 0)),
+        nested_close_threshold=(
+            None if config.get("nested_close_threshold") is None
+            else float(config["nested_close_threshold"])
+        ),
+        nested_layers=int(config.get("nested_layers", 0)),
+        tertiary_pool_factor=int(config.get("tertiary_pool_factor", 0)),
+        tertiary_min_pool_factor=int(config.get("tertiary_min_pool_factor", 0)),
+        tertiary_close_threshold=(
+            None if config.get("tertiary_close_threshold") is None
+            else float(config["tertiary_close_threshold"])
+        ),
+        tertiary_layers=int(config.get("tertiary_layers", 0)),
+        decoder_dim=int(config.get("decoder_dim", 0)),
+        detach_inactive_coarse_gradients=bool(
+            config.get("detach_inactive_coarse_gradients", True)
+        ),
+        decoder_pool_controller=bool(config.get("decoder_pool_controller", False)),
+        pool_controller_alpha=float(config.get("pool_controller_alpha", 1.0)),
+        pool_controller_beta=float(config.get("pool_controller_beta", 0.5)),
+        pool_controller_gamma=float(config.get("pool_controller_gamma", 0.5)),
+        # Missing fields identify older nested checkpoints whose routing did
+        # not include a short-pool budget.
+        short_pool_budget=int(config.get("blt_short_pool_budget", 0)),
+        short_pool_window=int(config.get("blt_short_pool_window", 0)),
+        secondary_min_patch_bytes=int(config.get("blt_secondary_min_patch_bytes", 16)),
+    )
+    model.forward_model.mamba2_unfused_inference = _configure_mamba2_inference_kernels(model)
+    return model
+
+
+def _parse_device_list(value: Optional[str]) -> List[torch.device]:
+    if not value:
+        return []
+    return [torch.device(item.strip()) for item in value.split(",") if item.strip()]
+
+
+def _validate_cuda_devices(devices: Sequence[torch.device]) -> None:
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "CUDA is unavailable. mamba_ssm selective-scan inference requires CUDA in this environment."
+        )
+    count = torch.cuda.device_count()
+    for device in devices:
+        if device.type != "cuda" or device.index is None or device.index >= count:
+            raise ValueError(f"requested device {device} is unavailable; CUDA device count is {count}")
+
+
+def load_nested_checkpoint(
+    checkpoint_path: str,
+    *,
+    precision: str = "fp16",
+    device: str = "cuda:0",
+    fine_device: Optional[str] = None,
+    nested_devices: Optional[str] = None,
+    tertiary_device: Optional[str] = None,
+    use_saved_placement: bool = False,
+) -> LoadedNestedModel:
+    """Load weights without materializing checkpoint optimizer tensors.
+
+    By default all inference runs on ``--device``. Model-parallel placement is
+    enabled by supplying ``fine_device`` and ``nested_devices``, or by opting
+    into the placement recorded in the checkpoint.
+    """
+    path = Path(checkpoint_path).expanduser().resolve()
+    checkpoint = None
+    if path.is_dir():
+        config_path = path / "config.json"
+        if not config_path.is_file():
+            raise FileNotFoundError(f"model config not found: {config_path}")
+        config = json.loads(config_path.read_text(encoding="utf-8"))
+    elif path.is_file():
+        # Backward-compatible local loading for an original training checkpoint.
+        warnings.warn(
+            "Loading a PyTorch .pt checkpoint requires pickle deserialization. "
+            "Only load .pt files that you created or obtained from a trusted source; "
+            "use the published SafeTensors directory for untrusted downloads.",
+            UserWarning,
+            stacklevel=2,
+        )
+        checkpoint = torch.load(path, map_location="cpu", weights_only=False, mmap=True)
+        config = dict(checkpoint.get("config") or {})
+    else:
+        raise FileNotFoundError(f"model directory or checkpoint not found: {path}")
+    if config.get("architecture") not in (None, "byte_latent_mamba_nested_jsonl") and config.get(
+        "architecture_label"
+    ) != "byte_latent_mamba_nested_jsonl":
+        raise ValueError(f"{path} is not identified as a nested JSONL checkpoint")
+
+    with torch.device("meta"):
+        model = _build_model(config)
+    if bool(getattr(model.forward_model, "mamba2_unfused_inference", False)):
+        print(
+            "Mamba-2 fused causal-conv1d is unavailable; using the portable "
+            "batched convolution + SSD scan path."
+        )
+    if path.is_dir():
+        from safetensors.torch import load_file
+
+        index_path = path / "model.safetensors.index.json"
+        single_path = path / "model.safetensors"
+        if index_path.is_file():
+            index = json.loads(index_path.read_text(encoding="utf-8"))
+            weight_map = dict(index.get("weight_map") or {})
+            expected = set(model.state_dict().keys())
+            published = set(weight_map.keys())
+            if expected != published:
+                missing = sorted(expected - published)[:8]
+                unexpected = sorted(published - expected)[:8]
+                raise RuntimeError(
+                    f"SafeTensors index does not match architecture; "
+                    f"missing={missing}, unexpected={unexpected}"
+                )
+            for filename in dict.fromkeys(weight_map.values()):
+                shard_path = path / filename
+                shard = load_file(str(shard_path), device="cpu")
+                model.load_state_dict(shard, strict=False, assign=True)
+                del shard
+        elif single_path.is_file():
+            state = load_file(str(single_path), device="cpu")
+            model.load_state_dict(state, strict=True, assign=True)
+            del state
+        else:
+            raise FileNotFoundError(
+                f"no model.safetensors or model.safetensors.index.json in {path}"
+            )
+        meta_names = [
+            name for name, value in model.state_dict().items() if value.device.type == "meta"
+        ]
+        if meta_names:
+            raise RuntimeError(f"unloaded model tensors remain: {meta_names[:8]}")
+        step = 0
+        trained_tokens = 0
+    else:
+        model.load_state_dict(checkpoint["model"], strict=True, assign=True)
+        step = int(checkpoint.get("step", 0))
+        trained_tokens = int(checkpoint.get("trained_tokens", 0))
+        del checkpoint
+    gc.collect()
+
+    dtype = _dtype_for_precision(precision)
+    if int(config.get("mamba_version", 1)) == 2 and dtype == torch.float16:
+        print(
+            "Mamba-2 FP16 inference: cached SSM accumulators will remain FP32; "
+            "BF16 is recommended when the GPU supports it."
+        )
+    if use_saved_placement:
+        fine_device = fine_device or str(config.get("fine_device") or "cuda:0")
+        if nested_devices is None:
+            saved_nested = config.get("nested_devices") or []
+            nested_devices = ",".join(str(item) for item in saved_nested)
+        tertiary_device = tertiary_device or (
+            str(config["tertiary_device"]) if config.get("tertiary_device") else None
+        )
+
+    coarse = _parse_device_list(nested_devices)
+    if coarse:
+        if not fine_device:
+            raise ValueError("--fine-device is required with --nested-devices")
+        fine = torch.device(fine_device)
+        tertiary = torch.device(tertiary_device) if tertiary_device else None
+        requested = [fine, *coarse, *([tertiary] if tertiary else [])]
+        _validate_cuda_devices(requested)
+        # Cast on CPU first so a large FP32 checkpoint is never temporarily
+        # placed in full on the root GPU.
+        model.to(dtype=dtype)
+        model.configure_model_parallel(fine, coarse, tertiary_device=tertiary)
+        root = fine
+        parallel = True
+    else:
+        root = torch.device(device)
+        _validate_cuda_devices([root])
+        model.to(root, dtype=dtype)
+        parallel = False
+
+    model.eval()
+    return LoadedNestedModel(
+        model=model,
+        config=config,
+        checkpoint_step=step,
+        trained_tokens=trained_tokens,
+        checkpoint_path=path,
+        precision=precision,
+        fine_device=root,
+        model_parallel=parallel,
+    )
+
+
+def new_stream(loaded: LoadedNestedModel, maximum_input_bytes: int) -> Dict[str, object]:
+    minimum = max(1, int(loaded.config.get("blt_min_patch_bytes", 16)))
+    max_patches = math.ceil((int(maximum_input_bytes) + 1) / minimum) + 8
+    return loaded.core.new_stream_state({}, max_patches=max_patches)
+
+
+def stream_token(
+    loaded: LoadedNestedModel,
+    stream: Dict[str, object],
+    token_id: int,
+) -> torch.Tensor:
+    token = torch.tensor(
+        [[int(token_id)]], dtype=torch.long, device=loaded.fine_device
+    )
+    return loaded.core.stream_step(stream, {"x": token})[0, 0]
+
+
+def apply_sampling_filters(
+    logits: torch.Tensor,
+    *,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    repeat_penalty: float,
+    recent_tokens: Sequence[int],
+) -> torch.Tensor:
+    filtered = logits.float().clone()
+    filtered[PAD] = filtered[BOS] = filtered[UNK] = -torch.inf
+    if repeat_penalty > 1.0:
+        for token_id in set(int(value) for value in recent_tokens):
+            if 0 <= token_id < filtered.numel():
+                value = filtered[token_id]
+                filtered[token_id] = (
+                    value / repeat_penalty if value >= 0 else value * repeat_penalty
+                )
+    filtered /= max(1e-5, float(temperature))
+    if top_k > 0 and top_k < filtered.numel():
+        threshold = torch.topk(filtered, int(top_k)).values[-1]
+        filtered[filtered < threshold] = -torch.inf
+    if 0.0 < top_p < 1.0:
+        probabilities = torch.softmax(filtered, dim=-1)
+        sorted_probabilities, sorted_indices = torch.sort(probabilities, descending=True)
+        remove = torch.cumsum(sorted_probabilities, dim=0) > float(top_p)
+        remove[0] = False
+        filtered[sorted_indices[remove]] = -torch.inf
+    return filtered
+
+
+def choose_token(filtered_logits: torch.Tensor, greedy: bool = False) -> int:
+    finite = torch.isfinite(filtered_logits)
+    if not bool(finite.any().item()):
+        raise FloatingPointError(
+            "sampling has no finite logits; the recurrent inference state became "
+            "non-finite. Retry with --precision bf16 (recommended for Mamba-2) "
+            "or --precision fp32."
+        )
+    if greedy:
+        return int(filtered_logits.argmax())
+    probabilities = torch.softmax(filtered_logits, dim=-1)
+    if not bool(torch.isfinite(probabilities).all().item()):
+        raise FloatingPointError(
+            "sampling probabilities became non-finite. Retry with --precision "
+            "bf16 (recommended for Mamba-2) or --precision fp32."
+        )
+    return int(torch.multinomial(probabilities, 1))
+
+
+def hierarchy_token_attribution(
+    loaded: LoadedNestedModel,
+    stream: Dict[str, object],
+    logits: torch.Tensor,
+    token_id: int,
+) -> Dict[str, object]:
+    """Measure direct L2/L3 decoder influence on one selected next byte.
+
+    This reuses the cached streaming state and only reruns the small decoder
+    and LM head. Positive deltas mean the dynamic hierarchy increased the
+    selected token's log probability relative to its BOE counterfactual.
+    """
+    parts = list(loaded.core.last_stream_decode_parts)
+    decoder_device = parts[0].device
+    normal_log_probability = float(
+        F.log_softmax(logits.float(), dim=-1)[int(token_id)].detach().cpu().item()
+    )
+
+    def counterfactual(*, remove_l2: bool, remove_l3: bool) -> float:
+        altered = list(parts)
+        if remove_l2:
+            initial_nested = stream["initial_nested_global"]
+            if initial_nested.device != decoder_device:
+                initial_nested = initial_nested.to(decoder_device, non_blocking=True)
+            altered[3] = initial_nested
+        if remove_l3 and len(altered) > 4:
+            initial_tertiary = stream["initial_tertiary_global"]
+            if initial_tertiary.device != decoder_device:
+                initial_tertiary = initial_tertiary.to(
+                    decoder_device, non_blocking=True
+                )
+            altered[4] = initial_tertiary
+        altered_logits = loaded.core.lm_head(
+            loaded.core.decoder(torch.cat(altered, dim=-1))
+        )[0, 0].float()
+        return float(
+            F.log_softmax(altered_logits, dim=-1)[int(token_id)].detach().cpu().item()
+        )
+
+    level2_active = int(stream["completed_nested_patches"]) > 0
+    level3_active = int(stream.get("completed_tertiary_patches", 0)) > 0
+    without_l2 = (
+        counterfactual(remove_l2=True, remove_l3=False)
+        if level2_active
+        else normal_log_probability
+    )
+    without_l3 = (
+        counterfactual(remove_l2=False, remove_l3=True)
+        if level3_active
+        else normal_log_probability
+    )
+    without_hierarchy = (
+        counterfactual(remove_l2=True, remove_l3=True)
+        if level3_active
+        else without_l2
+    )
+    return {
+        "level2_active": level2_active,
+        "level3_active": level3_active,
+        "level2_delta_logp": normal_log_probability - without_l2,
+        "level3_delta_logp": normal_log_probability - without_l3,
+        "hierarchy_delta_logp": normal_log_probability - without_hierarchy,
+        "selected_logp": normal_log_probability,
+    }
+
+
+def hierarchy_mode_logits(
+    loaded: LoadedNestedModel,
+    stream: Dict[str, object],
+    logits: torch.Tensor,
+    mode: str,
+) -> torch.Tensor:
+    """Return next-token logits with selected hierarchy readouts disabled."""
+    normalized = str(mode).lower()
+    if normalized == "full":
+        return logits
+    if normalized not in {"level1", "level12"}:
+        raise ValueError("hierarchy mode must be 'level1', 'level12', or 'full'")
+    parts = list(loaded.core.last_stream_decode_parts)
+    decoder_device = parts[0].device
+    if normalized == "level1":
+        initial_nested = stream["initial_nested_global"]
+        if initial_nested.device != decoder_device:
+            initial_nested = initial_nested.to(decoder_device, non_blocking=True)
+        parts[3] = initial_nested
+    if len(parts) > 4:
+        initial_tertiary = stream["initial_tertiary_global"]
+        if initial_tertiary.device != decoder_device:
+            initial_tertiary = initial_tertiary.to(decoder_device, non_blocking=True)
+        parts[4] = initial_tertiary
+    return loaded.core.lm_head(
+        loaded.core.decoder(torch.cat(parts, dim=-1))
+    )[0, 0]
+
+
+def generate_bytes(
+    loaded: LoadedNestedModel,
+    prompt: bytes,
+    *,
+    max_new_bytes: int,
+    temperature: float = 0.8,
+    top_p: float = 0.9,
+    top_k: int = 0,
+    repeat_penalty: float = 1.05,
+    repeat_window: int = 256,
+    greedy: bool = False,
+    seed: int = 1234,
+    collect_hierarchy_attribution: bool = False,
+    hierarchy_mode: str = "full",
+) -> Tuple[bytes, Dict[str, object]]:
+    if collect_hierarchy_attribution and hierarchy_mode != "full":
+        raise ValueError("hierarchy attribution is defined for the full hierarchy rollout")
+    torch.manual_seed(seed)
+    random.seed(seed)
+    stream = new_stream(loaded, len(prompt) + max_new_bytes + 2)
+    with torch.inference_mode():
+        logits = stream_token(loaded, stream, BOS)
+        for value in prompt:
+            logits = stream_token(loaded, stream, BYTE_OFFSET + int(value))
+        output = bytearray()
+        attributions: List[Dict[str, object]] = []
+        recent: List[int] = [BYTE_OFFSET + int(value) for value in prompt[-repeat_window:]]
+        for generated_position in range(max(0, int(max_new_bytes))):
+            sampling_logits = hierarchy_mode_logits(
+                loaded, stream, logits, hierarchy_mode
+            )
+            if not bool(torch.isfinite(sampling_logits).all().item()):
+                finite_count = int(torch.isfinite(sampling_logits).sum().item())
+                raise FloatingPointError(
+                    "non-finite generation logits before sampling: "
+                    f"generated_byte={generated_position} mode={hierarchy_mode} "
+                    f"precision={loaded.precision} "
+                    f"finite_logits={finite_count}/{sampling_logits.numel()}. "
+                    "Retry with --precision bf16 (recommended for Mamba-2) or "
+                    "--precision fp32."
+                )
+            filtered = apply_sampling_filters(
+                sampling_logits,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                repeat_penalty=repeat_penalty,
+                recent_tokens=recent[-repeat_window:],
+            )
+            token_id = choose_token(filtered, greedy=greedy)
+            if token_id == EOS:
+                break
+            if not BYTE_OFFSET <= token_id < BYTE_OFFSET + 256:
+                continue
+            if collect_hierarchy_attribution:
+                attributions.append(
+                    hierarchy_token_attribution(
+                        loaded, stream, sampling_logits, token_id
+                    )
+                )
+            output.append(token_id - BYTE_OFFSET)
+            recent.append(token_id)
+            logits = stream_token(loaded, stream, token_id)
+    stream["generated_attribution"] = attributions
+    stream["hierarchy_mode"] = hierarchy_mode
+    return bytes(output), stream
+
+
+def load_jsonl_texts(
+    path: str,
+    *,
+    text_field: str = "text",
+    max_documents: Optional[int] = None,
+) -> Iterable[Tuple[int, str]]:
+    jsonl = Path(path).expanduser().resolve()
+    if not jsonl.is_file():
+        raise FileNotFoundError(f"JSONL file not found: {jsonl}")
+    yielded = 0
+    with jsonl.open("r", encoding="utf-8") as handle:
+        for line_number, line in enumerate(handle, 1):
+            if not line.strip():
+                continue
+            try:
+                record = json.loads(line)
+            except json.JSONDecodeError as error:
+                raise ValueError(f"invalid JSON at {jsonl}:{line_number}: {error.msg}") from error
+            text = record.get(text_field)
+            if not isinstance(text, str):
+                raise ValueError(
+                    f"{jsonl}:{line_number} must contain a string field {text_field!r}"
+                )
+            yield line_number, text
+            yielded += 1
+            if max_documents is not None and yielded >= int(max_documents):
+                return

requirements.txt

@@ -0,0 +1,5 @@
+torch
+safetensors
+mamba-ssm
+causal-conv1d
+Pillow