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.

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:

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:

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:

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:

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:

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

Expected cache shape:

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.

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