birgermoell/oellm-9b-256k-theta64m-prelude-anneal300b

OELLM 9B — 256K context — Prelude anneal300b

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

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

Relationship to the earlier release

The earlier birgermoell/oellm-9b-256k-theta64m-prelude uses the approximately 1T-token Prelude base. This repository instead starts from the approximately 300B-token annealed checkpoint. It is published separately as an experimental comparison lineage, not as an overwrite or drop-in quality replacement.

Architecture

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

The export preserves the Q/K normalization parameters, tokenizer, and positional configuration used in Megatron-LM.

Training lineage

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

Total continued-pretraining volume: 7,982,809,088 tokens.

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

The large final theta is deliberate: the staged ABF schedule keeps RoPE dimensions usable at retrieval distances far beyond the original short-context training range.

Optimization and systems configuration

• 16 LUMI nodes / 128 AMD MI250X GCDs at every production stage.
• Tensor parallelism 8, pipeline parallelism 1, sequence parallelism enabled.
• Context parallelism increased to 16 at 256K.
• Micro-batch size 1 and global batch size 64.
• BF16, FlashAttention, distributed Adam (beta1 0.9, beta2 0.95), weight decay 0.1.
• Cosine learning-rate decay; 256K used 8e-6 to 8e-7 with approximately 5% warmup.
• Gradient clipping at 1.0 and selective activation recomputation.
• Megatron-Core torch_dist checkpoints; optimizer and RNG state are not part of this HF release.

The 256K production stage completed 59/59 iterations on launcher attempt 2, trained on 989,855,744 tokens, saved the final checkpoint successfully, and recorded zero skipped and zero NaN iterations. Attempt 1 ended during cold startup with a context-parallel NCCL timeout; the fresh second initialization then ran the entire stage stably. The last periodic training log at iteration 58 reported loss 1.427106. This is a run-health value, not a comparable benchmark score.

Continued-pretraining data

The 16K–128K curriculum used a frozen token-proportional multilingual blend with 152 prefixes. The 256K stage used a related 157-prefix blend targeted at very long sequences. Sources include FinePDFs, DCLM, HPLT3, multilingual synthetic text, Nemotron material, MegaMath, StarCoder, PeS2o, arXiv, and Wikipedia, grouped into length-aware tiers. The mix covers European languages plus code and scientific text, but is not uniform by language.

Conversion and release gates

The iteration-59 Megatron checkpoint is converted directly to a Qwen3-compatible Transformers layout in BF16. Publication requires the exact iteration marker, distributed metadata, non-empty weights and tokenizer, and a config with model_type=qwen3, max_position_embeddings=262144, and rope_theta=64000000. The Hub upload is verified for public visibility and required release files.

Evaluation status — partial

The complete English/Swedish 16K–256K retrieval evaluation is still in progress. The first evaluation attempt exhausted one 64 GiB GPU when beginning the first 256K case. The completed records below were preserved and published before the resumable two-GPU retry.

These numbers are partial and must not be interpreted as evidence of 256K retrieval quality. English 16K–128K main-grid cases are complete (24/24 correct); English 256K, all Swedish cases, and the control conditions had not yet run when this snapshot was made.

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

The raw snapshot is available at evaluation/niah-partial-en-16k-128k.jsonl. A separate evaluation/niah-partial-status.json records its scope and the interrupted job. Both files and this section will be superseded by the complete evaluation after the retry finishes.

Evaluation method

The forced-choice base-LM harness inserts numeric key/value facts into a long context. The queried fact is placed at a selected depth, while the four candidate answers include adversarial distractor values that genuinely occur elsewhere in the same context. Candidates are ranked by answer-token log likelihood, avoiding any dependency on chat instruction following. Actual constructed contexts can be slightly shorter than the nominal boundary because complete fact lines are used and space is reserved for the query/candidate suffix.

RULER long-context evaluation

Results supplied by Jouni Luoma, using NVIDIA/RULER. Scores range from 0 to 100, with higher being better. Context-length columns are nominal lengths; — means that task was not run.

Single-needle retrieval

Task	4K	8K	16K	32K	64K	128K
niah_single_1	100.00	100.00	100.00	100.00	99.80	99.80
niah_single_2	—	—	—	—	—	—
niah_single_3	—	—	—	—	—	—

Multi-key, multi-value, and multi-query retrieval

Task	4K	8K	16K	32K	64K	128K
niah_multikey_1	—	—	—	—	—	—
niah_multikey_2	99.60	96.80	92.60	88.80	69.40	41.40
niah_multikey_3	99.40	88.40	65.80	51.40	29.40	7.60
niah_multivalue	99.00	99.05	98.35	89.80	83.05	82.75
niah_multiquery	—	—	—	—	—	—

Variable tracking, common/frequent words, and question answering

Task	4K	8K	16K	32K	64K	128K
ruler_vt	—	—	—	—	—	—
ruler_cwe	—	—	—	—	—	—
ruler_fwe	—	—	—	—	—	—
ruler_qa_hotpot	47.80	47.60	45.80	44.40	39.80	38.60
ruler_qa_squad	66.22	52.62	52.45	50.53	49.90	44.95

No RULER results beyond 128K are included in this table.

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

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

Keep rope_theta=64000000 and max_position_embeddings=262144 unchanged. Full-length inference requires substantial accelerator memory, memory-efficient attention, careful KV-cache sizing, and usually tensor parallelism.

Intended use

This checkpoint is intended for base-model and long-context research, continued pretraining, retrieval experiments, and as a starting point for later supervised or preference-based post-training. Use completion-style prompts; it is not a chat or instruction-following model.

Limitations

• The initial quick evaluation is intentionally small and covers only English and Swedish.
• Single-needle retrieval does not establish long-document reasoning, summarization, multi-needle retrieval, or robust generation throughout the entire 256K window.
• General and short-context regression suites remain necessary before broad quality claims.
• This 300B-base lineage should not be assumed to match the earlier 1T-base release.
• The model inherits limitations, biases, and uneven language representation from the Prelude base and continued-pretraining mixture.

Reproducibility

• Base revision: f289699b246dba59907df27e743d4a433613175a
• Megatron-LM commit: b359462c12858cedd2238a22eca0dca7aa6b8872
• Tokenizer SHA256: ac214401227404105817141c21c6693e3b17c878465a42765d2c345310e003ce
• 16K–128K blend SHA256: c25d345c2c1969be4308fab2c5423584ef2fade72b1e711558fd459936b7ac6a
• 256K blend SHA256: 18b8d441cbf0306a1a5935cec17f63a17d45ec9b71282675e594a3d1a1d51f3a
• 256K checkpoint iteration: 59
• 256K training job: 21684001

Related resources

• Earlier 1T-base 256K checkpoint: birgermoell/oellm-9b-256k-theta64m-prelude
• 128K checkpoint from this anneal300b lineage: birgermoell/oellm-9b-128k-theta32m-prelude-anneal300b
• Implementation and evaluation code: BirgerMoell/openeuro-longctx-datamix