Add model card

README.md

@@ -0,0 +1,132 @@
+---
+language:
+- en
+library_name: transformers
+pipeline_tag: text-generation
+datasets:
+- HuggingFaceFW/fineweb
+tags:
+- llama
+- pretraining
+- linear-attention
+- kimi-delta-attention
+- custom_code
+---
+
+# Llama 1B with Kimi Delta Attention, 6B tokens
+
+A dense ~1B-parameter Llama 3-style decoder where one attention layer in every four is Kimi Delta Attention (KDA), a linear-attention layer from [Kimi Linear](https://arxiv.org/abs/2510.26692). It is otherwise identical to the [QK-norm baseline](https://huggingface.co/Mercity/pretrain-baseline-qknorm) and was trained on the same 6B FineWeb tokens.
+
+This is a base model. It is not instruction-tuned or safety-tuned.
+
+## Results
+
+### Training
+
+| Metric | Value | vs. baseline |
+| --- | ---: | ---: |
+| Final train loss (step 3,053) | 2.5634 | −0.0064 |
+| Final eval loss (step 3,000) | 2.5848 | −0.0059 |
+| Final grad norm (step 3,053) | 0.0514 | +0.0060 |
+| Peak grad norm after step 200 | 0.559 | +0.020 |
+| Tokens / steps | 6B / 3,053 | same |
+
+### Zero-shot benchmarks
+
+Scores from `lm-eval` on each task's full split. Shared-9 is the unweighted mean of the nine tasks.
+
+| Benchmark | Metric | Score | vs. baseline |
+| --- | --- | ---: | ---: |
+| HellaSwag | acc_norm | 39.65 | +0.67 |
+| WinoGrande | acc | 51.46 | −0.16 |
+| ARC-Easy | acc_norm | 39.60 | −0.55 |
+| ARC-Challenge | acc_norm | 24.06 | +0.43 |
+| PIQA | acc_norm | 67.79 | +1.20 |
+| OpenBookQA | acc_norm | 27.60 | −1.40 |
+| CommonsenseQA | acc | 20.07 | +0.25 |
+| SciQ | acc_norm | 63.80 | +0.30 |
+| LAMBADA | acc | 38.29 | +0.54 |
+| **Shared-9 average** | | **41.37** | **+0.14** |
+| Shared-9 average, 4-bit NF4 | | 40.64 | −0.18 |
+
+The +0.14 Shared-9 gap is within single-seed noise, so treat KDA as matching the baseline, not beating it.
+
+## Usage
+
+The architecture class ships with the checkpoint, so load it with `trust_remote_code=True`. The KDA layers need [flash-linear-attention](https://github.com/fla-org/flash-linear-attention) and a CUDA GPU (its kernels are written in Triton). Tested with `transformers==5.8.0` and `flash-linear-attention==0.5.2`.
+
+```bash
+pip install "transformers==5.8.0" "flash-linear-attention==0.5.2"
+```
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo = "Mercity/pretrain-kda-1b"
+
+tokenizer = AutoTokenizer.from_pretrained(repo)
+model = AutoModelForCausalLM.from_pretrained(
+    repo,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+    device_map="cuda",
+)
+
+inputs = tokenizer("The capital of France is", return_tensors="pt").to(model.device)
+# use_cache=False is required: the KDA layers keep no recurrent state between
+# decoding steps, so cached generation would feed them one token at a time.
+output = model.generate(**inputs, max_new_tokens=32, do_sample=False, use_cache=False)
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+Generation without the cache re-reads the full sequence at every step, so it is slow for long outputs. Scoring text in one forward pass has no such cost:
+
+```python
+batch = tokenizer("FineWeb is a large web-text dataset.", return_tensors="pt").to(model.device)
+with torch.no_grad():
+    loss = model(**batch, labels=batch["input_ids"]).loss
+print(f"loss={loss.item():.3f}  ppl={loss.exp().item():.1f}")
+```
+
+The first forward pass on a new machine spends about 90 seconds compiling the KDA kernels.
+
+## Model details
+
+| Setting | Value |
+| --- | --- |
+| Architecture | `LlamaKDA` (Llama 3-style dense decoder, hybrid KDA + softmax attention) |
+| Total parameters | 1.056B |
+| Layers | 32: 24 GQA + 8 KDA |
+| KDA layer positions | 0, 4, 8, 12, 16, 20, 24, 28 |
+| Hidden size | 1,536 |
+| Intermediate size (SwiGLU) | 5,120 |
+| Attention heads / KV heads | 12 / 6 (GQA layers) |
+| QK normalization | On (GQA layers) |
+| Max sequence length | 8,192 |
+| Tokenizer | Llama 2, 32,000 tokens |
+| Embeddings | Tied input and output |
+
+## Training
+
+| Setting | Value |
+| --- | --- |
+| Data | FineWeb `sample-10BT`, packed 8,192-token sequences |
+| Tokens / steps | 6B / 3,053 |
+| Batch | 10 per device × 24 gradient accumulation (~1.97M tokens per step) |
+| Optimizer | Muon (LR 0.02, momentum 0.95, 5 Newton-Schulz steps, WD 0.1) + AdamW (LR 3e-4, β 0.9/0.95, WD 0.1) |
+| Schedule | Cosine, 150 warmup steps |
+| Hardware | 1 × NVIDIA B200, ~18 hours |
+| Stack | TorchTitan, FlashAttention 4, Liger kernels, flash-linear-attention |
+
+## Related checkpoints
+
+| Model | Change from the baseline | Shared-9 |
+| --- | --- | ---: |
+| [Baseline (QK-norm)](https://huggingface.co/Mercity/pretrain-baseline-qknorm) | Reference model | 41.23 |
+| [N-gram 25%](https://huggingface.co/Mercity/pretrain-longcat-ngram-25pct) | ~25% of parameters moved into LongCat n-gram tables, 23 layers | 40.56 |
+| [N-gram 50%](https://huggingface.co/Mercity/pretrain-longcat-ngram-50pct) | ~48% of parameters moved into LongCat n-gram tables, 16 layers | 39.54 |
+
+## Limitations
+
+Trained on 6B English web tokens only, a small budget for a 1B model. Benchmark scores are single-seed. KDA inference needs a CUDA GPU and does not support cached generation through `transformers`. The model will repeat or make up facts and has had no alignment training.