Granite 4.0 1B — frozen-weight KV-sharing experiment

The frozen-weight modification severely damages language-model quality. This release is a negative research result, not an improved or production-ready Granite model.

Pairing adjacent layers halves allocated KV-cache storage. On the local WikiText-2 subset, perplexity changes from 19.55 to 707,495.20 (lower is better). HellaSwag subset accuracy with per-token normalization changes from 62.50% to 31.25%.

What this release changes

OVRLab's independent experimental derivative of the standard all-attention IBM Granite 4.0 1B instruct model (approximately 1.63B parameters). This is OVRLab's next model experiment after the Granite 3.1 concision experiment, and a separate model with a different base and intervention.

Adjacent layers share keys and values: zero-based layers 0/1 use layer 0's K/V, 2/3 use layer 2's K/V, and so on. Each layer keeps its own queries, output projection, and feed-forward network. All 40 layers still execute; 20 own KV caches. Consumer K/V projections are bypassed, but their original parameters remain in the checkpoint. The tokenizer, generation configuration, and BF16 checkpoint are unchanged. This is an architecture/inference-code modification with zero training, not a fine-tune, weight edit, or quantization. Checkpoint size is unchanged.

The design is related to Cross-Layer Attention. That paper studies trained architectures; it does not establish that this frozen Granite conversion should preserve quality. This experiment was motivated by reading DeepSeek's KV-cache compression work. It does not implement the full DeepSeek CSA2/CED design and does not claim architectural novelty or IBM endorsement.

Measured results

Context tokens Native KV MiB Paired KV MiB Native prefill s Paired prefill s Native decode tok/s Paired decode tok/s Decode change
512 40.0 20.0 0.914 0.884 18.58 19.64 +5.7%
2,048 160.0 80.0 4.062 4.043 9.35 9.67 +3.4%
4,096 320.0 160.0 9.751 9.534 3.75 5.85 +55.8%
Diagnostic Native Paired
WikiText-2 subset perplexity ↓ 19.55 707,495.20
HellaSwag subset accuracy, token normalized ↑ 62.50% 31.25%
HellaSwag subset accuracy, raw ↑ 40.62% 25.00%

Tested on Apple M1 Pro, 10 CPU cores, 14 GPU cores, 32 GiB unified memory; macOS-26.3.1-arm64-arm-64bit. Python 3.12.13, PyTorch 2.8.0, Transformers 4.57.1; MPS, BF16, SDPA, batch size 1. The same physical weights are used for both conditions. These are local measurements, not a prediction of CUDA or server performance.

Performance uses 3 measured repeats after one full warm-up per condition/context. Condition order alternates by repeat. The table reports medians; all samples are included in the JSON. Decoding uses 32 identical teacher-forced tokens, with no early stopping. Device work is synchronized at timing boundaries. Prefill excludes loading, tokenization, sampling, and network overhead. It is not end-to-end serving latency.

KV memory counts unique allocated key/value tensor storage, not total or peak memory. The stored weight file remains unchanged in size. The 50% cache reduction does not imply 50% total memory savings or 2× speed. At 32,768 tokens the formula predicts 2.5 GiB → 1.25 GiB of BF16 KV tensors at batch size 1; this is an extrapolation, not a tested 32K result.

WikiText-2 uses the first 16 non-overlapping windows of 512 tokens from the raw test split, scoring 8,176 next-token predictions, without special tokens and with context reset per window. HellaSwag uses 32 validation examples, sampled with seed 42, zero-shot raw and token-normalized continuation likelihood. Context and endings are tokenized separately; no harness-specific preprocessing is applied. These custom subset diagnostics are not standard full-benchmark scores. The HellaSwag sample is small and unsuitable for fine-grained ranking. Five fixed greedy examples illustrate chat behavior without constituting an instruction-following benchmark.

No training, fine-tuning, calibration, or quantization was performed. No original training corpus is required for this experiment. Evaluation text comes only from the pinned public datasets. No parameters were updated. Untested: long contexts beyond the lengths above, multilingual performance, safety behavior, and production serving frameworks.

Load the experimental model

The custom code is pinned to Transformers 4.57.1. Install the tested environment before loading; arbitrary newer versions are not supported by this release. Use trust_remote_code=True to load the included Python implementation.

pip install torch==2.8.0 transformers==4.57.1 accelerate==1.10.1 huggingface-hub==0.36.0 safetensors==0.6.2
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "OVRLab/granite-4.0-1b-kvshare-experiment"
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(
    repo, trust_remote_code=True, dtype=torch.bfloat16,
    attn_implementation="sdpa",
).eval().to(device)

prompt = tokenizer.apply_chat_template(
    [{"role": "user", "content": "What is 17 + 28?"}],
    tokenize=False, add_generation_prompt=True,
)
inputs = tokenizer(prompt, return_tensors="pt").to(device)
with torch.inference_mode():
    output = model.generate(**inputs, max_new_tokens=64, do_sample=False)
print(tokenizer.decode(output[0, inputs.input_ids.shape[1]:], skip_special_tokens=True))

The model may produce poor or incoherent answers, as shown by the evaluations. For normal use, load the original IBM model. Compatibility with Ollama, llama.cpp, vLLM, and hosted inference providers has not been established. This release is intended for reproducing and studying the experiment.

Reproduce and inspect

Base revision: 6a7381ba1f54d684ff508d991aeb7dc580157103. The release checkpoint is byte-identical to that revision's checkpoint; SHA-256 verification is included. Differences are in the custom configuration and model forward code. Dataset revisions and example indices are recorded in the raw results.

Attribution and license

Base weights and tokenizer: Granite Team, IBM. Custom experiment: OVRLab. Code reuses Hugging Face Transformers' Apache-2.0 implementation. See LICENSE, NOTICE, and the original model card.

Downloads last month
227
Safetensors
Model size
2B params
Tensor type
BF16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for OVRLab/granite-4.0-1b-kvshare-experiment

Finetuned
(12)
this model

Paper for OVRLab/granite-4.0-1b-kvshare-experiment