K2-Horizon-7B-GDN-Hybrid-Exp
EXPERIMENTAL. This is a research prototype. Quality drops on reasoning tasks.
A hybrid version of IFM/K2-Horizon-7B in which half of the full-attention layers are replaced with Gated DeltaNet (GDN) layers to reduce KV cache size.
Architecture
- 36 layers: the 18 even-indexed layers (0-based) are GDN, and the 18 odd-indexed layers keep the original full attention.
- GDN layers: 16 key heads and 32 value heads (head dimension 128), plus a short convolution with kernel size 4.
- KV cache: 73,728 bytes per token instead of 147,456. Each sequence also requires a fixed-size GDN state of about 19 MiB in bf16.
The GDN layers were initialized from the original attention weights using a method adapted from Taylor-Calibrate. They were then distilled from the base model: first layer by layer, then end to end. Only the GDN parameters were trained; all other weights are unchanged from the base model.
Usage
This model requires custom code (trust_remote_code=True). Tested with transformers==5.17.0.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "TFM-ML/K2-Horizon-7B-GDN-Hybrid-Exp"
tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
repo, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto")
License
Apache-2.0