| --- |
| 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 |
| - multimodal |
| - custom-code |
| - base_model |
| language: |
| - en |
| datasets: |
| - HuggingFaceM4/FineVision |
| - HuggingFaceFW/finepdfs |
| - HuggingFaceFW/fineweb-edu |
| - webshart/suno-various-94k |
| --- |
| |
| # 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. |
|
|
|  |
|
|
| ## 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` |
| - Last training count: 15GB **Absoloutly undertrained** |
|
|
| 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 |