78 lines
3.8 KiB
Markdown
78 lines
3.8 KiB
Markdown
|
---
|
|||
|
|
license: apache-2.0
|
|||
|
|
language: [en, sv, de, fr, es, it, nl, pl, pt, cs, fi, da, el, bg, hr, hu, ro, sk, sl, et, lt, lv, ga, mt, eu, gl, is, nb, nn, sr, uk, ca, mk, sq, oc, lb, bs]
|
|||
|
|
tags: [long-context, rope, abf, niah, qwen3, openeurollm, base-model, 128k]
|
|||
|
|
pipeline_tag: text-generation
|
|||
|
|
library_name: transformers
|
|||
|
|
---
|
|||
|
|
|
|||
|
|
# OELLM 9B — 128K context (ABF, θ=32M) — v3 (full training)
|
|||
|
|
|
|||
|
|
**Full-training production version** of [`oellm-9b-128k-theta32m`](https://huggingface.co/birgermoell/oellm-9b-128k-theta32m).
|
|||
|
|
Long-context extension of OpenEuroLLM `baby_9b_dense` (Qwen3 dense 9B), 4K → **128K** via native
|
|||
|
|
ABF (RoPE θ-scaling). Same θ=32M config as the sweep arm but trained on **2B tokens** (≈6×).
|
|||
|
|
**Base model — not instruction-tuned.** Multilingual.
|
|||
|
|
|
|||
|
|
## Evaluation — verified, multilingual (this checkpoint)
|
|||
|
|
Base-LM forced-choice NIAH, depth-stratified, at **128K**, 12 European languages:
|
|||
|
|
|
|||
|
|
| metric | result |
|
|||
|
|
|---|---|
|
|||
|
|
| **Overall 128K** | **100% (341/341)** |
|
|||
|
|
| **depth-0 (far start)** | **100% in every language** (6/6 each) |
|
|||
|
|
| Languages | cs de el en es fi fr hu pl pt sv uk — **all 100%** |
|
|||
|
|
| 4K / 16K / 64K | 100% |
|
|||
|
|
|
|||
|
|
Full training lifted the two languages that were 96% in the sweep arm (el, sv) to 100%.
|
|||
|
|
|
|||
|
|
## Why θ=32M (the key finding)
|
|||
|
|
Far-position (depth-0) retrieval failure at long context is a **RoPE high-dimension OOD**
|
|||
|
|
problem, **not** a data problem — fixed by scaling θ to the target length. Critical θ ≈ doubles
|
|||
|
|
per length-octave: 64K→8M, 128K→16M→**32M (100%)**, 256K→~32M. θ ablation at 128K depth-0:
|
|||
|
|
2M/5M→0%, 8M→0%, 16M→90%, **32M→100%**. (cf. LongRoPE2, arXiv:2502.20082.)
|
|||
|
|
|
|||
|
|
## How to reproduce
|
|||
|
|
**Base:** OELLM `baby_9b_dense` (Qwen3 dense): 36 layers, hidden 4096, FFN 12288, 32 heads /
|
|||
|
|
8 KV (GQA), kv-channels 128, qk-layernorm, RMSNorm, SwiGLU, untied embeddings, vocab 262144,
|
|||
|
|
`openeurollm/tokenizer-256k`, native 4K.
|
|||
|
|
|
|||
|
|
**Recipe:** staged native ABF, `4K→16K→32K→64K(θ=2M)→128K(θ=32M)`; this model = the **128K@θ=32M
|
|||
|
|
stage trained to 2B tokens** from the 64K checkpoint (`--finetune`).
|
|||
|
|
|
|||
|
|
**Training (Megatron-LM, LUMI 16× MI250X):**
|
|||
|
|
```
|
|||
|
|
--rotary-base 32000000 --seq-length 131072 --max-position-embeddings 131072 --use-flash-attn
|
|||
|
|
--tensor-model-parallel-size 8 --context-parallel-size 8 --sequence-parallel --use-distributed-optimizer
|
|||
|
|
--micro-batch-size 1 --global-batch-size 64 --bf16 --train-iters ~238 (2B tokens)
|
|||
|
|
--lr 1e-5 --min-lr 1e-6 --lr-decay-style cosine --weight-decay 0.1 --clip-grad 1.0
|
|||
|
|
--recompute-activations --recompute-granularity selective --save-interval 100
|
|||
|
|
--qk-layernorm --normalization RMSNorm --swiglu --group-query-attention --num-query-groups 8
|
|||
|
|
```
|
|||
|
|
ROCm 6.4.4 / PyTorch 2.9 / TE 2.4 / FA 2.8; ~490 tok/s/GPU at 128K.
|
|||
|
|
|
|||
|
|
**Data:** Jouni Luoma's length-biased multilingual long-context mix (finepdfs all langs + edu,
|
|||
|
|
dclm, hplt3, multisynth, nemotron, megamath, starcoder, pes2o, arxiv, wiki; tiered
|
|||
|
|
short/medium/long). Note: depth-0 is **insensitive** to the data mix — θ is the lever.
|
|||
|
|
|
|||
|
|
**Reproduce the eval:**
|
|||
|
|
```
|
|||
|
|
python scripts/eval_base_lm_niah.py --model <this-model> \
|
|||
|
|
--context-lengths 4096 16384 65536 131072 --depths 0.0 0.25 0.5 0.75 1.0 \
|
|||
|
|
--languages en de fr es nl pl sv fi cs it pt el hu uk da --trials 6
|
|||
|
|
```
|
|||
|
|
Code / write-ups: https://github.com/BirgerMoell/openeuro-longctx-datamix
|
|||
|
|
|
|||
|
|
## Usage
|
|||
|
|
```python
|
|||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|||
|
|
import torch
|
|||
|
|
m = AutoModelForCausalLM.from_pretrained("birgermoell/oellm-9b-128k-theta32m-v3",
|
|||
|
|
torch_dtype=torch.bfloat16, device_map="auto")
|
|||
|
|
tok = AutoTokenizer.from_pretrained("birgermoell/oellm-9b-128k-theta32m-v3")
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
## Caveats
|
|||
|
|
- **Base model** (no instruction tuning) — use as a completion model.
|
|||
|
|
- Evaluated with **single-needle** forced-choice NIAH; broader multi-task RULER not yet run.
|
|||
|
|
- Keep `rope_theta=32000000`, `max_position_embeddings=131072` for 128K.
|