feat(model): OLMo 3 (Olmo3ForCausalLM) in the dense-transformer package #117

Merged

rcsheets merged 4 commits from feat/olmo3 into main 2026-10-04 06:25:29 +00:00

Conversation 1 Commits 4 Files changed 11 +3272 -32

rcsheets commented 2026-10-04 05:57:00 +00:00

Owner

Copy link

Adds OLMo 3 (Olmo3ForCausalLM, e.g. allenai/Olmo-3-7B-Think and its compressed-tensors NVFP4 quantization) to internal/model/mistral.

It builds on #114, #115 and #116 (generation-config stop ids, per-layer rope and layer_types, sliding-window attention), all now on main.

The architecture

OLMo 3 is the Llama block with its norms moved and two norms added. Its layers also differ in attention:

	Llama / Mistral / Granite	OLMo 3
norm before attention	input_layernorm	--
norm on attention output, before the residual add	--	post_attention_layernorm
norm before MLP	post_attention_layernorm	--
norm on MLP output, before the residual add	--	post_feedforward_layernorm
q / k	--	q_norm, k_norm: RMSNorm over the whole projection, before rope
attention	full, one rope	3 of every 4 layers: 4096-token sliding window, plain rope; the 4th: full attention, YaRN

post_attention_layernorm names a different position in each family.

What changes

• Norm positions: layer carries RMSNorm weights at four optional positions: each sublayer's input or its output. Forward applies whichever are non-nil, through a norm helper that passes the input through when there is no weight, so Forward has no OLMo branch. The loader assigns tensor names by architecture (Mistral.postNorm).
• q/k norms: q_norm and k_norm are applied in place to the q and k projections before rope. RMSNorm is row-local on both backends, so in place is safe.
• Per-layer attention: each layer carries its own rope table and window, from config.Model.RopeFreqsForLayer and SlidingWindowFor, and passes the window to PagedAttention. For every other architecture these equal the old model-wide values with no window.
• mistral.New no longer rejects sliding_window.
• AGENTS.md and the README describe the OLMo support, including the KV cost: an F32 cache with as many kv heads as q heads is 1 MiB per cached token for the 7B.

The chat template is not part of this PR. Chat completions still return 501 until OLMo's think-only template has a renderer, which comes next.

Verification

• go build ./..., go vet ./..., go test ./... pass; gofmt is clean.
• Tiny checkpoint against transformers (testdata/generate_olmo3.py, transformers 5.14.1):
  • 4 layers (3 sliding, 1 full), window 6, YaRN on the full layer.
  • Sequences up to 28 tokens, so the window is crossed many times.
  • TestForwardOlmo3 matches the logits at 1e-4 for prefill and for token-by-token decode. It loads the config both as saved (v5, rope_parameters keyed by layer type) and respelled as the published checkpoints write it (v4, rope_scaling).
• TestOlmo3FeaturesAreLoadBearing: removing any one of seven pieces breaks the match with the goldens: the window, YaRN on the full layer, plain rope on the sliding layers, q_norm, k_norm, post_attention_layernorm, post_feedforward_layernorm.
• TestForwardOlmo3CUDA: the same goldens pass on the RTX PRO 6000 Blackwell and the RTX 3070. The whole mistral CUDA package passes on both (25 tests).
• Real checkpoints on the RTX PRO 6000 (greedy raw completions; chat is not wired up yet):

check	Olmo-3-7B-Think (BF16, 13.6 GiB)	Olmo-3-7B-Think-nvfp4 (4.9 GiB)
weight load	4 s	1 s
"The capital of France is"	" Paris. The capital of the United States is Washington, D.C. ..."	same
"Question: What is 17*23?\nAnswer:"	" 17*23 = 391", finish_reason: stop	same
fact planted at token 0 of a 6,488-token prompt (past the 4096 window)	recalled exactly	recalled exactly

🤖 Generated with Claude Code

Adds OLMo 3 (`Olmo3ForCausalLM`, e.g. `allenai/Olmo-3-7B-Think` and its compressed-tensors NVFP4 quantization) to `internal/model/mistral`. It builds on #114, #115 and #116 (generation-config stop ids, per-layer rope and `layer_types`, sliding-window attention), all now on main. ## The architecture OLMo 3 is the Llama block with its norms moved and two norms added. Its layers also differ in attention: | | Llama / Mistral / Granite | OLMo 3 | | --- | --- | --- | | norm before attention | `input_layernorm` | -- | | norm on attention output, before the residual add | -- | `post_attention_layernorm` | | norm before MLP | `post_attention_layernorm` | -- | | norm on MLP output, before the residual add | -- | `post_feedforward_layernorm` | | q / k | -- | `q_norm`, `k_norm`: RMSNorm over the whole projection, before rope | | attention | full, one rope | 3 of every 4 layers: 4096-token sliding window, plain rope; the 4th: full attention, YaRN | `post_attention_layernorm` names a different position in each family. ## What changes - **Norm positions:** `layer` carries RMSNorm weights at four optional positions: each sublayer's input or its output. `Forward` applies whichever are non-nil, through a `norm` helper that passes the input through when there is no weight, so `Forward` has no OLMo branch. The loader assigns tensor names by architecture (`Mistral.postNorm`). - **q/k norms:** `q_norm` and `k_norm` are applied in place to the q and k projections before rope. RMSNorm is row-local on both backends, so in place is safe. - **Per-layer attention:** each `layer` carries its own rope table and window, from `config.Model.RopeFreqsForLayer` and `SlidingWindowFor`, and passes the window to `PagedAttention`. For every other architecture these equal the old model-wide values with no window. - `mistral.New` no longer rejects `sliding_window`. - AGENTS.md and the README describe the OLMo support, including the KV cost: an F32 cache with as many kv heads as q heads is 1 MiB per cached token for the 7B. The chat template is not part of this PR. Chat completions still return 501 until OLMo's think-only template has a renderer, which comes next. ## Verification - `go build ./...`, `go vet ./...`, `go test ./...` pass; gofmt is clean. - **Tiny checkpoint against transformers** (`testdata/generate_olmo3.py`, transformers 5.14.1): - 4 layers (3 sliding, 1 full), window 6, YaRN on the full layer. - Sequences up to 28 tokens, so the window is crossed many times. - `TestForwardOlmo3` matches the logits at 1e-4 for prefill and for token-by-token decode. It loads the config both as saved (v5, `rope_parameters` keyed by layer type) and respelled as the published checkpoints write it (v4, `rope_scaling`). - `TestOlmo3FeaturesAreLoadBearing`: removing any one of seven pieces breaks the match with the goldens: the window, YaRN on the full layer, plain rope on the sliding layers, `q_norm`, `k_norm`, `post_attention_layernorm`, `post_feedforward_layernorm`. - `TestForwardOlmo3CUDA`: the same goldens pass on the RTX PRO 6000 Blackwell and the RTX 3070. The whole `mistral` CUDA package passes on both (25 tests). - **Real checkpoints** on the RTX PRO 6000 (greedy raw completions; chat is not wired up yet): | check | Olmo-3-7B-Think (BF16, 13.6 GiB) | Olmo-3-7B-Think-nvfp4 (4.9 GiB) | | --- | --- | --- | | weight load | 4 s | 1 s | | "The capital of France is" | " Paris. The capital of the United States is Washington, D.C. ..." | same | | "Question: What is 17*23?\nAnswer:" | " 17*23 = 391", `finish_reason: stop` | same | | fact planted at token 0 of a 6,488-token prompt (past the 4096 window) | recalled exactly | recalled exactly | 🤖 Generated with [Claude Code](https://claude.com/claude-code)

rcsheets added 6 commits 2026-10-04 05:57:00 +00:00

feat: stop on generation_config's eos_token_id as well as the tokenizer's EOS ... efe739d64d

The engine and the grammar ended a turn only on the tokenizer's single
EOS id. A checkpoint's generation_config.json can name others, and HF
generate and vLLM stop on those: OLMo 3's tokenizer names <|endoftext|>
while its chat turns end in <|im_end|>, so gllm would run on past every
reply; GLM ships three ids; Nemotron 3 adds </s> beside <|im_end|>.

config.Load now reads eos_token_id (one id or a list) into
StopTokenIDs. The engine stops on any of them, masks all of them while
a grammar is incomplete, and grammar.NewVocab takes them as extra stops
that are allowed exactly when the value is complete.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

feat(config): v4 rope_scaling, per-layer-type rope, and layer_types ... 3d2b2fe6a7

config.Load ignored rope_scaling, the transformers v4 key for rope
scaling, so a YaRN checkpoint written that way got plain rope with no
error. It also could not represent models whose layers differ in rope:
OLMo 3 gives its sliding-window layers plain rope and its full-attention
layers YaRN, which transformers v5 writes as rope_parameters keyed by
layer type -- a shape the flat decode read as all zeros (theta 0).

Load now folds rope_scaling into RopeParameters (or, for OLMo 3, into
the full-attention entry of RopeByLayerType, as Olmo3Config does),
recognizes the nested rope_parameters form, accepts the old "type" key,
and reads layer_types. Models get per-layer settings through
RopeFreqsForLayer and SlidingWindowFor. No checkpoint gllm serves today
uses rope_scaling or sliding layer types, so their rope is unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

feat(backend): sliding-window PagedAttention ... f2a07123aa

PagedAttention gains a window argument: when positive, the token at
position p attends to positions p-window+1 through p only, the masking
OLMo 3's sliding_attention layers (and Mistral 7B v0.1's) use; 0 keeps
today's unlimited causal attention. The CPU reference starts each
token's visible range at the window, and both CUDA kernels start their
position loop there; the launcher's flash dispatch measures the longest
windowed span rather than the longest context.

Every caller passes 0, so behavior is unchanged until a model asks for
a window. The KV cache still keeps every position; evicting the ones a
window no longer reaches is a separate memory optimization.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Merge remote-tracking branch 'origin/feat/config-rope-scaling-layer-types' into feat/olmo3 ... c6dc48332d

# Conflicts:
#	internal/config/config.go

Merge remote-tracking branch 'origin/feat/sliding-window-attention' into feat/olmo3 71975e85e2

feat(model): OLMo 3 (Olmo3ForCausalLM) in the dense-transformer package ... 4c531d3ce4

OLMo 3 is the Llama block with its norms moved and two added, so it is
served by internal/model/mistral rather than a package of its own:

- Each layer carries RMSNorm weights at four optional positions, each
  sublayer's input (Llama, Mistral, Granite) or its output before the
  residual add (OLMo), and Forward applies whichever are present. The
  loader assigns post_attention_layernorm by architecture, since it is
  the pre-MLP norm in Llama and the attention-output norm in OLMo.
- q_norm and k_norm RMS-normalize the whole q and k projections before
  rope.
- Rope tables and attention windows are per layer, from
  config.RopeFreqsForLayer and SlidingWindowFor: plain rope and a
  4096-token window on three of every four layers, YaRN and full
  attention on the fourth.

mistral.New no longer rejects sliding_window, now that PagedAttention
honors one.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

pr-reviewer-bot commented 2026-10-04 05:57:00 +00:00

Collaborator

Copy link

Automated review by pr-reviewer v0.54.0 | Safety Check | Nemotron 3 Nano | tracking id r-c1f10b-6f1e37
This is an AI-generated review and may contain mistakes.

Status: ❌ Failed

---

Review failed. Tracking id r-c1f10b-6f1e37 — see logs for details.

Comment @pr-reviewer-bot retry to try again.

<!-- pr-reviewer:review --> *Automated review by [pr-reviewer](https://git.brooktrails.org/brooktrails/pr-reviewer) v0.54.0 | Safety Check | Nemotron 3 Nano | tracking id `r-c1f10b-6f1e37`* *This is an AI-generated review and may contain mistakes.* **Status:** ❌ Failed --- Review failed. Tracking id `r-c1f10b-6f1e37` — see logs for details. Comment `@pr-reviewer-bot retry` to try again.

rcsheets referenced this pull request 2026-10-04 06:01:02 +00:00

feat(tokenizer): OLMo 3 Think chat template #118

rcsheets added 1 commit 2026-10-04 06:24:09 +00:00

Merge remote-tracking branch 'origin/main' into feat/olmo3 00c909679b

rcsheets merged commit da882f2a4d into main 2026-10-04 06:25:29 +00:00

rcsheets deleted branch feat/olmo3 2026-10-04 06:25:29 +00:00

rcsheets referenced this pull request from a commit 2026-10-04 06:25:30 +00:00

Merge pull request 'feat(model): OLMo 3 (Olmo3ForCausalLM) in the dense-transformer package' (#117) from feat/olmo3 into main

Sign in to join this conversation.

Reviewers

No reviewers

Labels

Clear labels

No items

No labels

Milestone

Clear milestone

No items

No milestone

Projects

Clear projects

No items

No project

Assignees

Clear assignees

No assignees

2 participants

Notifications

Subscribe

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set

Reference

brooktrails/gllm!117

No description provided.