birgermoell/oellm-9b-128k-theta32m

OELLM 9B — 128K context (ABF, θ=32M)

Long-context extension of the OpenEuroLLM baby_9b_dense (Qwen3 dense 9B) base model, from 4K → 128K context via native ABF (RoPE base-frequency scaling; no YaRN). Base model — not instruction-tuned. Multilingual (37 European languages).

It achieves perfect needle-in-a-haystack retrieval at 128K, including the hardest far-start (depth-0) position, across languages — the failure mode that low-θ ABF and length-biased data could not fix.

• Full-training sibling (2B tokens, same θ): oellm-9b-128k-theta32m-v3
• Code / scripts / write-ups: https://github.com/BirgerMoell/openeuro-longctx-datamix

---

Key result: depth-0 is a RoPE θ problem, not a data problem

Extending context by raising RoPE θ uniformly (standard ABF) 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 θ to the target length — consistent with the high-RoPE-dimension OOD analysis of LongRoPE2.

The critical θ roughly doubles per context-length octave:

context	θ that fixes depth-0
64K	8M
128K	16M (90%) → 32M (100%)
256K	~32–64M (extrapolated)

θ ablation at 128K depth-0 (base-LM NIAH): θ=2M/5M → 0%, θ=8M → 0% (fixes ≤64K only), θ=16M → 90%, θ=32M → 100%. This checkpoint uses θ=32M.

---

Evaluation

Method

Base-LM forced-choice NIAH (no instruction-following required): a key → magic-number needle is inserted at a controlled depth in a filler context of N tokens; the model scores 4 candidate values by summed log-probability of the answer tokens; argmax = prediction. Distractor values appear in the 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.

Results — by context length (depth-averaged)

context	accuracy	depth-0
4K	100%	100%
16K	100%	100%
64K	100%	100%
128K	100%	100%

Results — 128K, multilingual (depth-stratified)

264 trials across en, fi, fr, hu, pl, pt, sv (6 trials × 5 depths × langs):

language	128K accuracy	depth-0
en, fi, fr, hu, pl, pt	100%	6/6
sv	96% (1 non-depth-0 miss)	6/6
total	99% (262/264)	depth-0 = 100% in every language

All depths (0.0–1.0) are 100% at 64K; at 128K depth-0 is 100% across all tested languages.

---

How to reproduce

1. Base model & architecture

OpenEuroLLM baby_9b_dense (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. Tokenizer: openeurollm/tokenizer-256k. Native context 4K (rope θ ≈ 100K).

2. Extension recipe (staged native ABF)

Continued-pretraining in stages, raising --seq-length and --rotary-base (θ) each stage, loading the previous stage with --finetune:

4K(native) → 16K → 32K → 64K(θ=2M) → 128K(θ=32M)   # this checkpoint

The decisive step is the 128K stage at θ=32M (not the ~500K–5M that NTK/formulas suggest — those under-scale the high dims and leave depth-0 broken). This checkpoint = the 64K checkpoint (ckpt_65536, θ=2M, length-biased data) → 128K stage, θ=32M, ~0.3B tokens. (The v3 sibling is the same stage trained to 2B tokens.)

3. Training config (Megatron-LM, LUMI / 16× MI250X nodes)

--position-embedding-type rope --rotary-base 32000000
--seq-length 131072 --max-position-embeddings 131072 --use-flash-attn
--tensor-model-parallel-size 8 --pipeline-model-parallel-size 1
--context-parallel-size 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
--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

Container: ROCm 6.4.4 / PyTorch 2.9 / TE 2.4 / FA 2.8. Throughput ≈ 490 tok/s/GPU at 128K. Megatron→HF conversion emits Qwen3Config + q_norm/k_norm (validated bit-identical).

4. Data

Length-biased multilingual long-context mix ("keep the domain mix, upsample long docs within domain"): finepdfs (all langs + edu), dclm, hplt3, multisynth, nemotron, megamath, starcoder, pes2o, arxiv, wiki — tiered short/medium/long. Depth-0 is insensitive to the data mix; θ is the lever. See docs/depth0_diagnosis_theta_sweep.md and docs/longctx_undersampling_analysis.md in the repo.

5. Reproduce the eval

python scripts/eval_base_lm_niah.py --model <this-model> \
    --ctx 4096 16384 65536 131072 --depths 0.0 0.25 0.5 0.75 1.0 \
    --langs en de fr es nl pl sv fi cs it --trials 6

---

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
m = AutoModelForCausalLM.from_pretrained("birgermoell/oellm-9b-128k-theta32m",
                                         torch_dtype=torch.bfloat16, device_map="auto")
tok = AutoTokenizer.from_pretrained("birgermoell/oellm-9b-128k-theta32m")
# Base model: use as a completion model. Keep rope_theta=32M / max_position=131072 for 128K.

Caveats

• Base model (no chat/instruction tuning) — evaluate/use as a completion model.
• This checkpoint is a short θ-adaptation (~0.3B tokens); the fuller 2B-token version is v3.
• Evaluated with single-needle NIAH; broader multi-task RULER (multi-key, aggregation, VT) not yet run.
• config.json sets rope_theta=32000000, max_position_embeddings=131072 — keep these for 128K.

Citation

OpenEuroLLM long-context extension (baby_9b_dense). RoPE θ-scaling diagnosis informed by LongRoPE2 (Gao et al., 2025, arXiv:2502.20082).