124 lines
6.3 KiB
Markdown
124 lines
6.3 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) — prelude (1T-token base)
|
||
|
||
128K long-context extension of the OpenEuroLLM **prelude** base — Qwen3 dense 9B
|
||
(`qwen3_9b_hf_baby` at **iter_0124800 ≈ 1 trillion tokens**, the strongest available base
|
||
checkpoint). Same validated θ=32M recipe as
|
||
[`oellm-9b-128k-theta32m-v3`](https://huggingface.co/birgermoell/oellm-9b-128k-theta32m-v3)
|
||
(which used the earlier ~0.6T checkpoint), now on the upgraded 1T base.
|
||
**Base model — not instruction-tuned.** Multilingual (37 European languages).
|
||
|
||
---
|
||
|
||
## Key result: depth-0 is a RoPE-θ problem, not data
|
||
Standard ABF (uniform θ-scaling) leaves the model unable to retrieve from the **far start** of a
|
||
long window ("depth-0"). We showed this is **not** fixable with more long-range data (two
|
||
length-biased datasets gave depth-0 ≈ 0%), but **is** fixed by scaling RoPE θ to the target length
|
||
— the high RoPE dimensions are out-of-distribution at long range (cf. LongRoPE2, arXiv:2502.20082).
|
||
|
||
**Critical θ ≈ doubles per context-length octave:**
|
||
| context | θ |
|
||
|---|---|
|
||
| 64K | 8M |
|
||
| 128K | 16M (90%) → **32M (100%)** |
|
||
| 256K | ~32–64M |
|
||
θ ablation @128K depth-0: 2M/5M=0%, 8M=0% (fixes ≤64K), 16M=90%, **32M=100%**. This model uses **θ=32M**.
|
||
|
||
## Evaluation
|
||
**Method:** base-LM **forced-choice NIAH** (4-way, log-likelihood of answer tokens; no
|
||
instruction-following needed). Distractor values are placed *in-context* (adversarial), so the
|
||
chance floor is below 25%. Depth 0.0 = needle at the far start (max query distance); 1.0 = most
|
||
recent. Script: [`scripts/eval_base_lm_niah.py`](https://github.com/BirgerMoell/openeuro-longctx-datamix/blob/main/scripts/eval_base_lm_niah.py).
|
||
|
||
**Results @128K** — 15 languages, 900 trials. **Overall 96% (864/900).**
|
||
|
||
By depth (needle position; 0.0 = far start):
|
||
| depth | 0.0 | 0.25 | 0.5 | 0.75 | 1.0 |
|
||
|---|---|---|---|---|---|
|
||
| accuracy | **97%** | 96% | **88%** | 96% | 100% |
|
||
|
||
By language: cs, da, en, es, pt, uk = 100%; el, fr, hu = 96%; it, nl, pl = 93%; de, fi, sv = 90%.
|
||
depth-0 is 12/12 in every language except fi (8/12).
|
||
|
||
**Read:** the θ=32M fix is confirmed — **far-position (depth-0) retrieval is solid (97%)**. The
|
||
model is strong overall (96%) with a **mild "lost-in-the-middle"** (depth 0.5 = 88%). This is a
|
||
slightly softer profile than the 0.6T-base sibling **v3** (which scored ~100% across depths on 12
|
||
languages) — an honest, characterized difference, concentrated in the middle of the window rather
|
||
than the far end. Candidate follow-ups: a larger 128K token budget, or investigating whether the
|
||
mid-depth softness is a property of the longer-pretrained 1T base. (4K/16K/64K remain ~100%.)
|
||
|
||
## How to reproduce
|
||
|
||
### 1. Base model & architecture
|
||
OpenEuroLLM `prelude` = `qwen3_9b_hf_baby` **iter_0124800** (~1T tokens). Qwen3 dense:
|
||
36 layers, hidden 4096, FFN 12288, 32 attention heads / 8 KV groups (GQA), kv-channels 128,
|
||
**qk-layernorm**, RMSNorm, SwiGLU, untied embeddings, vocab 262144, `openeurollm/tokenizer-256k`,
|
||
native context 4K (rope θ=100000).
|
||
|
||
### 2. Conversion HF → Megatron
|
||
The prelude base ships in HF safetensors; converted to Megatron-core (torch_dist) via
|
||
**Megatron-Bridge** (`convert_checkpoints.py import`) before extension.
|
||
|
||
### 3. Extension recipe (staged native ABF)
|
||
Continued-pretraining, raising `--seq-length` and `--rotary-base` each stage, `--finetune` from
|
||
the previous stage:
|
||
```
|
||
prelude(4K, θ=100k) → 16K(θ=500k) → 32K(θ=1M) → 64K(θ=2M) → 128K(θ=32M)
|
||
```
|
||
Token budgets: 16K/32K ≈ 1B each, 64K = 3B, 128K = 2B.
|
||
|
||
### 4. Training config (Megatron-LM, LUMI / 16× MI250X)
|
||
```
|
||
--rotary-base {500k|1M|2M|32M per stage} --seq-length {16384..131072} --use-flash-attn
|
||
--tensor-model-parallel-size 8 --pipeline-model-parallel-size 1
|
||
--context-parallel-size {1|1|2|8} --sequence-parallel --use-distributed-optimizer
|
||
--micro-batch-size 1 --global-batch-size 64 --bf16
|
||
--lr 1e-5 --min-lr 1e-6 --lr-decay-style cosine --weight-decay 0.1 --clip-grad 1.0
|
||
--adam-beta1 0.9 --adam-beta2 0.95 --recompute-activations --recompute-granularity selective
|
||
--qk-layernorm --normalization RMSNorm --swiglu --group-query-attention --num-query-groups 8
|
||
--ckpt-format torch_dist --finetune --no-save-optim --no-save-rng --save-interval 100
|
||
```
|
||
Container: ROCm 6.4.4 / PyTorch 2.9 / TE 2.4 / FA 2.8. Throughput ≈ 500 tok/s/GPU at 128K.
|
||
Megatron→HF export emits Qwen3Config + q_norm/k_norm.
|
||
|
||
### 5. Data
|
||
**Jouni Luoma's length-biased multilingual long-context mix** (token-proportional blend over 152
|
||
sources): finepdfs (all langs + edu), dclm, hplt3 (38 langs), multisynth (synthetic multilingual),
|
||
nemotron, megamath, starcoder, pes2o, arxiv, wiki — tiered short/medium/long (long_threshold=64K).
|
||
The same data as v3. Note: **depth-0 is insensitive to the data mix — θ is the lever.** Data hosted
|
||
on stable `/scratch` (the burst-buffer `/flash` proved unreliable mid-run).
|
||
|
||
### 6. 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
|
||
```
|
||
|
||
## Usage
|
||
```python
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
import torch
|
||
m = AutoModelForCausalLM.from_pretrained("openeurollm/oellm-9b-128k-theta32m-prelude",
|
||
torch_dtype=torch.bfloat16, device_map="auto")
|
||
tok = AutoTokenizer.from_pretrained("openeurollm/oellm-9b-128k-theta32m-prelude")
|
||
# Base completion model. Keep rope_theta=32M / max_position=131072 for 128K.
|
||
```
|
||
|
||
## Caveats
|
||
- **Base model** (no instruction/chat 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.
|
||
|
||
## Links / citation
|
||
Code & write-ups: https://github.com/BirgerMoell/openeuro-longctx-datamix
|
||
(`docs/depth0_diagnosis_theta_sweep.md` for the θ analysis). Sibling (0.6T base):
|
||
`oellm-9b-128k-theta32m-v3`. RoPE-θ diagnosis informed by LongRoPE2 (arXiv:2502.20082).
|