Nested_Mamba_3Level / README.md
Alienanthony's picture
Update README.md
dd25e82 verified
|
Raw
History Blame Contribute Delete
10.4 kB
---
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.
![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`
- 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