[BUG] OLMoE crashes on inference: fused_moe_ffn_into hidden dim mismatch - up_proj/down_proj swapped via wrong expert-type mapping

[danielesalpietro]

Description

allenai/OLMoE-1B-7B-0924-Instruct loads successfully (with the fix in #119) but crashes on the very first inference request. Every expert forward pass fails immediately (layer 0, expert 0, 1, 2, 3, 4, ... all fail within the same millisecond), and the process falls back to a degraded/errored state.

[WARN] GPUExecFunc: expert forward failed:  fused_moe_ffn_into: hidden dim mismatch
Exception raised from fused_moe_ffn_into at /workspace/MoE-Infinity/extensions/kernel/fused_moe_mlp.cu:123 (most recent call first):
frame #0: c10::Error::Error(...)
frame #1: c10::detail::torchCheckFail(...)
frame #2: fused_moe_ffn_into(at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, CUstream_st*) + 0x2309
frame #3: MoEMLP::ForwardHelper(CUstream_st*) + 0x186
frame #4: MoEMLP::forward(at::Tensor, CUstream_st*) + 0x2f2
frame #5: ExpertDispatcher::GPUExecFunc(int, int) + 0xa24
frame #6: base::detail::startThread(void*) + 0x8a
  (expert_idx=1 layer_idx=0) - expert_dispatcher.cpp:535

(repeats for every expert_idx at layer_idx=0)

Root cause

MODEL_MAPPING_TYPES["olmoe"] in moe_infinity/common/constants.py is set to 4 (MIXTRAL_MOE_DENSE_ACT_DENSE). That expert type tells MoEMLP::ForwardHelper to read the per-expert weight buffer (param_) in gate, down, up order:

auto& gate_proj = param_[0];
auto& up_proj   = (expert_type_ == DEEPSEEK_MOE_DENSE_ACT_DENSE) ? param_[1] : param_[2];
auto& down_proj = (expert_type_ == DEEPSEEK_MOE_DENSE_ACT_DENSE) ? param_[2] : param_[1];

That ordering matches HF Mixtral's own expert module (w1=gate, w2=down, w3=up):

# transformers.models.mixtral.modeling_mixtral.MixtralBlockSparseTop2MLP.__init__
self.w1 = nn.Linear(hidden_dim, ffn_dim, bias=False)   # gate
self.w2 = nn.Linear(ffn_dim, hidden_dim, bias=False)   # down
self.w3 = nn.Linear(hidden_dim, ffn_dim, bias=False)   # up

But OLMoE's expert module registers weights in gate, up, down order instead:

# transformers.models.olmoe.modeling_olmoe.OlmoeMLP.__init__
self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
self.up_proj   = nn.Linear(hidden_size, intermediate_size, bias=False)
self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

Since param_ is populated via model.named_parameters(recurse=True) (module registration order), OLMoE's buffers land in param_ = [gate_proj, up_proj, down_proj]. Interpreted through the Mixtral branch, up_proj and down_proj get swapped, so fused_moe_ffn_into receives down_proj = up_proj tensor (shape [intermediate_size, hidden_size] instead of [hidden_size, intermediate_size]), tripping the H == H_out check at fused_moe_mlp.cu:123 whenever hidden_size != intermediate_size (true for allenai/OLMoE-1B-7B-0924-Instruct: hidden_size=2048, intermediate_size=1024).

Notably, qwen3 — which shares the exact same gate_proj/up_proj/down_proj naming and registration order as OLMoE — is correctly mapped to 5 (DEEPSEEK_MOE_DENSE_ACT_DENSE, the gate, up, down branch) in the same table. This strongly suggests "olmoe": 4 was simply a copy/paste slip and should be 5.

Suggested fix

--- a/moe_infinity/common/constants.py
+++ b/moe_infinity/common/constants.py
@@
-    "olmoe": 4,
+    "olmoe": 5,

(mirroring "qwen3": 5 right above it). I haven't submitted a PR yet since I don't have a way to test the fused CUDA kernel path in my current environment, but I'm confident in the analysis above from reading the code paths — happy to verify further if useful.

Environment

• Model: allenai/OLMoE-1B-7B-0924-Instruct (64 experts, top-8, hidden_size=2048, intermediate_size=1024, 16 layers)
• Served via moe_infinity.entrypoints.openai.api_server_v2
• --device-memory-ratio 0.5 --kv-cache-ratio 0.15 --max-batch-size 8
• transformers 4.57.6
• Built on top of Fix OLMoE support and silent model-init failures #119 (which fixes model loading/registration for OLMoE; this issue is a separate runtime bug that only surfaces once the model actually loads and serves a request)

Related

• Fix OLMoE support and silent model-init failures #119 — fixes OLMoE loading (architecture registration + silent init failures). This issue is about the inference-time crash that happens after a model loads successfully.