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

OELLM 9B — 128K context — Prelude anneal300b

This is a BF16 Hugging Face export of a 128K continued-pretraining checkpoint derived from openeurollm/prelude, revision f289699b246dba59907df27e743d4a433613175a (anneal300b_iter_0989075). It uses the Qwen3 dense 9B architecture, a 131,072-token maximum position length, and RoPE theta 32,000,000.

This is a base completion model. It is not instruction-tuned or safety-aligned.

What this checkpoint is

The goal of this run was to test whether the later anneal300b Prelude checkpoint could be extended from its native short context to 128K while preserving useful retrieval at shorter lengths. It follows the same staged absolute-base-frequency (ABF) recipe as the earlier birgermoell/oellm-9b-128k-theta32m-prelude, but it is a different lineage:

• Earlier release: the approximately 1T-token Prelude base.
• This release: openeurollm/prelude@anneal300b_iter_0989075, the approximately 300B-token annealed checkpoint at commit f289699b246dba59907df27e743d4a433613175a.

The separate repository name is intentional. This model is an experimental comparison point, not an in-place replacement for the earlier 1T-base release.

Architecture

Property	Value
Architecture	Qwen3-compatible dense decoder-only transformer
Parameters	approximately 9B
Transformer layers	36
Hidden size	4,096
Feed-forward size	12,288
Attention heads	32
Key/value groups	8 (grouped-query attention)
Head dimension	128
Vocabulary	262,144 tokens
Normalization	RMSNorm with Q/K layer normalization
MLP	SwiGLU
Embeddings	untied input/output embeddings
Maximum positions	131,072
RoPE theta	32,000,000
Weight dtype	BF16

The tokenizer is the OpenEuroLLM 256K-vocabulary tokenizer. The exported Transformers configuration preserves the Q/K normalization parameters and the 128K positional settings used during training.

Training lineage

The model was continued on LUMI using staged native ABF context extension:

Stage	Sequence length	RoPE theta	Iterations	Tokens	LUMI job
16K	16,384	500,000	953	999,292,928	21683851
32K	32,768	1,000,000	476	998,244,352	21683852
64K	65,536	2,000,000	715	2,998,927,360	21740962
128K	131,072	32,000,000	238	1,996,488,704	21683854

Total continued-pretraining volume through this checkpoint: 6,992,953,344 tokens.

The curriculum increases both sequence length and RoPE theta at each stage:

Prelude anneal300b
  -> 16K  (theta 500K)
  -> 32K  (theta 1M)
  -> 64K  (theta 2M)
  -> 128K (theta 32M)

The unusually large final theta is deliberate. At long retrieval distances, smaller theta values can leave high-frequency RoPE dimensions outside the range encountered during training. The theta schedule follows the behavior characterized in the earlier Prelude long-context work; the new quick evaluation below tests whether it transfers to this later base checkpoint.

Optimization and parallelism

• Hardware: 16 LUMI nodes / 128 AMD MI250X GCDs per stage.
• Tensor parallelism: 8; pipeline parallelism: 1.
• Context parallelism: 1 at 16K and 32K, 2 at 64K, and 8 at 128K.
• Micro-batch size: 1; global batch size: 64.
• Precision: BF16 with FlashAttention.
• Optimizer: distributed Adam, beta1 0.9, beta2 0.95, weight decay 0.1.
• Schedule: cosine decay from 1e-5 to 1e-6, with approximately 5% warmup.
• Stability controls: gradient clipping at 1.0, sequence parallelism, selective activation recomputation, and fresh checkpoint continuation at every context stage.
• Checkpoint format during training: Megatron-Core torch_dist; optimizer and RNG state were not retained in the release export.

The 128K stage completed all 238 iterations on its first launcher attempt, saved the final checkpoint successfully, reported validation loss 1.556005 (perplexity 4.73985), and had zero skipped and zero NaN iterations. These values are run-health diagnostics, not cross-model benchmark scores.

Continued-pretraining data

The 16K–128K curriculum used a frozen, token-proportional multilingual blend with 152 data prefixes. It is length-aware and draws from long and general-domain sources including FinePDFs, DCLM, HPLT3, multilingual synthetic material, Nemotron data, MegaMath, StarCoder, PeS2o, arXiv, and Wikipedia. Sources were grouped into short, medium, and long tiers so that longer natural documents were represented as sequence length increased.

The mixture spans European languages and code/scientific text, but it is not balanced equally by language. The model therefore should not be interpreted as having uniform capability across every language listed in the metadata. The frozen 16K–128K blend hash is provided below for exact run identity. The exact 128K training receipt and the machine-readable quick-evaluation records are included in this repository.

Megatron-to-Hugging-Face conversion

The iteration-238 Megatron checkpoint was converted directly to a Qwen3-compatible Transformers layout in BF16. Publication is gated on all of the following:

1. The Megatron checkpoint marker and iteration-238 distributed metadata must exist.
2. The converted config must report model_type=qwen3, max_position_embeddings=131072, and rope_theta=32000000.
3. The converted weights and tokenizer must load for the retrieval evaluation.
4. The complete 16K/32K/64K/128K quick-evaluation matrix must finish successfully.
5. The release job revalidates the files, authenticated namespace, and uploaded Hub contents.

Quick retrieval evaluation

The release job ran a deliberately small, base-LM forced-choice needle-in-a-haystack smoke test at 16K, 32K, 64K, and 128K. Each language/context cell contains three needle depths and two trials per depth (six cases). This verifies export and retrieval behavior; it is too small for a publication-grade accuracy estimate.

Context	English	Swedish	Combined
16,384	100.0% (6/6)	100.0% (6/6)	100.0% (12/12)
32,768	100.0% (6/6)	100.0% (6/6)	100.0% (12/12)
65,536	100.0% (6/6)	100.0% (6/6)	100.0% (12/12)
131,072	100.0% (6/6)	100.0% (6/6)	100.0% (12/12)

Depth breakdown

Depth 0.0 places the needle at the far beginning (maximum retrieval distance); depth 1.0 places it nearest the query.

Language	Context	Depth 0.0	Depth 0.5	Depth 1.0
en	16,384	100.0% (2/2)	100.0% (2/2)	100.0% (2/2)
en	32,768	100.0% (2/2)	100.0% (2/2)	100.0% (2/2)
en	65,536	100.0% (2/2)	100.0% (2/2)	100.0% (2/2)
en	131,072	100.0% (2/2)	100.0% (2/2)	100.0% (2/2)
sv	16,384	100.0% (2/2)	100.0% (2/2)	100.0% (2/2)
sv	32,768	100.0% (2/2)	100.0% (2/2)	100.0% (2/2)
sv	65,536	100.0% (2/2)	100.0% (2/2)	100.0% (2/2)
sv	131,072	100.0% (2/2)	100.0% (2/2)	100.0% (2/2)

Controls

Language	Control	Accuracy
en	short_ctx	100.0% (2/2)
en	shuffled	100.0% (2/2)
en	no_context	50.0% (1/2)
sv	short_ctx	100.0% (2/2)
sv	shuffled	100.0% (2/2)
sv	no_context	0.0% (0/2)

Main-condition failures

• None in this small smoke test.

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "birgermoell/oellm-9b-128k-theta32m-prelude-anneal300b"
tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = AutoModelForCausalLM.from_pretrained(
    repo_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

Long-context inference requires substantial accelerator memory. Keep rope_theta=32000000 and max_position_embeddings=131072 unchanged for the trained 128K configuration.

This checkpoint is best used for base-model research, continued pretraining, long-context retrieval experiments, and as a starting point for subsequent supervised or preference-based post-training. It is not a chat model; use raw completion prompts or a base-LM evaluation harness.

For shorter prompts, no special scaling override is needed: the model should be loaded with its published 128K configuration and evaluated directly at the desired shorter length.

Evaluation method

The included smoke test uses forced-choice needle-in-a-haystack scoring designed for base models. Each context contains key/value facts. A target fact is placed at a selected relative depth, and the model scores four candidate values by answer-token log likelihood. Distractors are values of other keys that really occur in the same context, making the test stricter than choosing against arbitrary out-of-context numbers. The highest-scoring candidate is the prediction.

The test covers English and Swedish at nominal 16K, 32K, 64K, and 128K lengths. Actual context lengths can be slightly below the nominal boundary because the generator fills contexts using complete fact lines and leaves room for the query and candidate suffix. Depth 0.0 is the most distant case: the needle is at the beginning and the query is at the end.

Controls check that short-context scoring works, that the model follows the presented binding rather than a memorized key/value association, and that the no-context condition remains near chance. All per-case records are published under evaluation/.

Limitations

• Base model only; prompts should use completion-style evaluation rather than assuming chat behavior.
• The included retrieval run is a small smoke test in English and Swedish, not a comprehensive benchmark.
• General and short-context regression suites have not yet been completed for this checkpoint.
• Single-needle retrieval does not establish multi-hop reasoning, long-document summarization, multi-needle retrieval, or robust generation across the entire 128K window.
• This run begins from the 300B-token annealed Prelude checkpoint and should not be assumed to match the quality profile of the earlier 1T-base 128K Prelude release.
• Very long inference may require tensor parallelism, memory-efficient attention, and careful KV-cache sizing.
• The training data is multilingual and inherits the limitations and biases of the Prelude base and continued-pretraining blend.

Reproducibility

• Base revision: f289699b246dba59907df27e743d4a433613175a
• Megatron-LM commit: b359462c12858cedd2238a22eca0dca7aa6b8872
• Tokenizer SHA256: ac214401227404105817141c21c6693e3b17c878465a42765d2c345310e003ce
• 16K–128K blend SHA256: c25d345c2c1969be4308fab2c5423584ef2fade72b1e711558fd459936b7ac6a
• Training hardware: 16 LUMI nodes / 128 AMD MI250X GCDs

Related resources

• Earlier 1T-base 128K checkpoint: birgermoell/oellm-9b-128k-theta32m-prelude
• OpenEuroLLM Prelude base: openeurollm/prelude
• Long-context implementation and evaluation code: BirgerMoell/openeuro-longctx-datamix