Repository navigation
fix(olmoe): apply q_norm/k_norm before the head reshape - #317
Merged
Merged
Conversation
OLMoE normalizes Q/K over the full num_heads*head_dim projection (the q_norm/k_norm weights are hidden-sized, 2048 for OLMoE-1B-7B), matching the upstream order where the RMSNorm runs on the q_proj/k_proj output before splitting into heads. The code applied the norm after reshaping to [batch, seq_len, heads, head_dim], so rms_norm received a 2048-element weight against a head_dim-wide (128) last dimension and aborted with std::invalid_argument. OLMoE-1B-7B was non-functional as a result. Move q_norm/k_norm ahead of the head reshape. Validated on mlx-community/OLMoE-1B-7B-0125-Instruct-4bit: loads and generates coherent greedy temp-0 output with no crash. The change is confined to olmoe.rs attention. Closes #316
inureyes
added
type:bug
7 tasks
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
OLMoE normalizes Q/K over the full num_heads*head_dim projection (the q_norm/k_norm weights are hidden-sized, 2048 for OLMoE-1B-7B), matching the upstream order where the RMSNorm runs on the q_proj/k_proj output before splitting into heads. The code applied the norm after reshaping to [batch, seq_len, heads, head_dim], so rms_norm received a 2048-element weight against a head_dim-wide (128) last dimension and aborted with std::invalid_argument. OLMoE-1B-7B was non-functional as a result. Move q_norm/k_norm ahead of the head reshape. Validated on mlx-community/OLMoE-1B-7B-0125-Instruct-4bit: loads and generates coherent greedy temp-0 output with no crash. The change is confined to olmoe.rs attention. Closes lablup#316
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
OLMoE normalizes Q/K over the full num_heads*head_dim projection (the q_norm/k_norm weights are hidden-sized, 2048 for OLMoE-1B-7B), matching the upstream order where the RMSNorm runs on the q_proj/k_proj output before the head split. The code applied the norm after reshaping to
[batch, seq_len, heads, head_dim], sorms_normreceived a 2048-element weight against a head_dim-wide (128) last dimension and aborted withstd::invalid_argument; OLMoE-1B-7B was non-functional as a result. This moves q_norm/k_norm ahead of the head reshape. Validated onmlx-community/OLMoE-1B-7B-0125-Instruct-4bit: loads and generates coherent greedy temp-0 output (e.g. "What is the capital of France? Paris, France.") with no crash. The change is confined to olmoe.rs attention; no other model is affected. Found while validating epic #307 (#302 wires olmoe to the fused decode-MoE kernel, which could not be runtime-validated until this crash was fixed). Closes #316