---
language:
- en
library_name: transformers
pipeline_tag: text-generation
license: apache-2.0
tags:
- causal-lm
- text-generation
- pytorch
- safetensors
- custom-code
- trust-remote-code
- kv-cache
- kv-cache-quantization
- kv-cache-aware-training
- 2-bit-kv-cache
- experimental
datasets:
- HuggingFaceFW/fineweb-edu
model-index:
- name: BananaMind-KV1-8M-2Bit-Experimental
  results:
  - task:
      type: text-generation
      name: Text Generation
    dataset:
      type: wikitext
      name: WikiText-2 Raw
      config: wikitext-2-raw-v1
      split: test
    metrics:
    - type: mean_kld
      name: Mean KLD vs 16-bit KV cache reference
      value: 0.09162049349112282
      unit: nats/token
    - type: kv_cache_compression
      name: Average KV cache shrink vs FP16
      value: 5.333333333333333
      unit: x
---

# BananaMind-KV1-8M-2Bit-Experimental

BananaMind-KV1-8M-2Bit-Experimental is a Hugging Face Transformers causal language model repository for the first Project KV1 checkpoint.

This is an 8.13M parameter BananaMind KV1 model trained for text generation with KV-cache-aware 2-bit K/V quantization. The model weights are stored normally in `model.safetensors`; the experimental part is the KV1 low-bit K/V path used during training and generation, with generation-time cache entries affine-quantized to 2 bits and bit-packed.

## Model Details

- **Model type:** decoder-only causal language model
- **Architecture:** `KV1ForCausalLM`
- **Model type ID:** `banana_kv1`
- **Parameters:** 8,130,816
- **Weights:** FP32 `safetensors`
- **Context length:** 1024 tokens
- **Vocabulary size:** 8192
- **Hidden size:** 256
- **Layers:** 8
- **Attention heads:** 8
- **KV heads:** 2
- **Intermediate size:** 768
- **KV cache:** KV-cache-aware 2-bit affine K/V path, packed as a 2-bit cache when `use_cache=True`
- **Training data:** `HuggingFaceFW/fineweb-edu`, `sample-10BT` subset

## Repository Files

This folder is structured as a Hugging Face model repo:

- `config.json` maps `AutoConfig` and `AutoModelForCausalLM` to the custom KV1 code.
- `model.safetensors` contains the model weights.
- `tokenizer.json` and `tokenizer_config.json` contain the tokenizer.
- `generation_config.json` contains default generation settings.
- `configuration_kv1.py` and `modeling_kv1.py` implement the custom model.
- `train_args.json` records the local training configuration.

Because this model uses custom architecture code, load it with `trust_remote_code=True`.


## Quick Start

Use the local folder path, or replace it with the full Hub repo ID after upload.

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "./BananaMind-KV1-8M-2Bit-Experimental"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to(device)
model.eval()

prompt = "The color of a banana is yellow. The color of the sky is"
inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(device)

with torch.inference_mode():
    output_ids = model.generate(
        **inputs,
        max_new_tokens=50,
        do_sample=False,
    )

print(tokenizer.decode(output_ids[0], skip_special_tokens=True))
```

## Generate With Repo Script

From the repo root:

```bash
python generate_kv1_2bit.py --prompt "The color of a banana is yellow. The color of the sky is" --max-new-tokens 50 --greedy
```

For sampled output:

```bash
python generate_kv1_2bit.py --prompt "The color of a banana is yellow. The color of the sky is" --max-new-tokens 100 --temperature 0.8 --top-p 0.95
```

## 2-Bit KV-Cache-Aware Training And Cache

This checkpoint is KV-cache-aware trained, not just a model with a post-training KV-cache wrapper. During non-cache forward passes, KV1 uses an STE-style 2-bit affine quantization path for K/V tensors so the model learns under the low-bit K/V constraint. During generation with `use_cache=True`, those K/V tensors are stored as packed 2-bit cache entries.

When `use_cache=True`, Project KV1 quantizes each K/V vector to four affine levels and packs four 2-bit codes into one `torch.uint8` byte:

```text
byte = code0 + (code1 << 2) + (code2 << 4) + (code3 << 6)
```

For this model, `head_dim=32`, so packed K/V cache tensors use last dimension `8` instead of `32`.

The first layer cache tuple is:

```python
k_pack, k_min, k_scale, v_pack, v_min, v_scale = outputs.past_key_values[0]
```

`k_pack` and `v_pack` are packed `uint8` tensors. The `min` and `scale` tensors are FP16 metadata used to dequantize the packed 2-bit codes for attention.

Check the packed cache shape:

```bash
python generate_kv1_2bit.py --prompt "The color of the sky is" --max-new-tokens 8 --greedy --show-cache
```

Expected cache shape:

```text
k_pack=(1, 2, seq_len, 8) torch.uint8
v_pack=(1, 2, seq_len, 8) torch.uint8
```

That confirms the KV cache is bit-packed instead of stored as FP16 cache tensors.

## Evaluation

The KV-cache-aware 2-bit checkpoint was compared against the 16-bit KV-cache counterpart on WikiText-2 Raw test data.

| Metric | Value |
| --- | ---: |
| Mean KLD, `KL(P_16bit_KV || P_2bit_packed_KV)` | 0.0916204935 nats/token |
| Mean KLD | 0.1321804316 bits/token |
| Evaluated token positions | 372,675 |
| Sequence length | 1024 |
| Average KV cache shrink vs FP16 | 5.3333x |

The local evaluation artifact is `wikitext2_kld_2bit_packed_vs_fp16.json`.

## Training

Training used the `sample-10BT` subset of `HuggingFaceFW/fineweb-edu`, sequence length 1024, seed 1337, and the local 8k tokenizer recorded in `train_args.json`. The run used `kv_cache_bits=2`, so this is a KV-cache-aware trained checkpoint rather than a post-training KV-cache-only conversion.

Key training arguments:

- `max_steps`: 10204
- `batch_size`: 72
- `grad_accum`: 16
- `lr`: 0.0005
- `min_lr`: 0.00005
- `warmup_steps`: 500
- `weight_decay`: 0.1
- `kv_cache_bits`: 2

## License

Apache 2.0