birgermoell/oellm-9b-128k-theta32m-prelude

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 (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.

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

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).