A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 07:30 UTC
Fix `test_moe.py` for Transformers v5 (#33413)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Jan 30, 2026 at 14:03 UTC
a11bc12d53a55cfa395ace8af8c9956a045b1791
1 file changed
+19
-8
tests/kernels/moe/test_moe.py
+19
-8
@@ -680,13 +680,21 @@ def test_mixtral_moe(
680
681
# Load the weights
682
vllm_moe.gate.weight.data[:] = hf_moe.gate.weight.data
683
- for i in range(config.num_local_experts):
684
- weights = (
685
- hf_moe.experts[i].w1.weight.data,
686
- hf_moe.experts[i].w3.weight.data,
687
- )
688
- vllm_moe.experts.w13_weight[i][:] = torch.cat(weights, dim=0)
689
- vllm_moe.experts.w2_weight[i][:] = hf_moe.experts[i].w2.weight.data
683
+ if isinstance(hf_moe.experts, torch.nn.ModuleList):
684
+ # Transformers v4
685
+ for i in range(config.num_local_experts):
686
+ weights = (
687
+ hf_moe.experts[i].w1.weight.data,
688
+ hf_moe.experts[i].w3.weight.data,
689
+ )
690
+ vllm_moe.experts.w13_weight[i][:] = torch.cat(weights, dim=0)
691
+ vllm_moe.experts.w2_weight[i][:] = hf_moe.experts[i].w2.weight.data
692
+ else:
693
+ # Transformers v5
694
+ vllm_moe.experts.w13_weight.data[:] = hf_moe.experts.gate_up_proj.data
695
+ vllm_moe.experts.w2_weight.data[:] = hf_moe.experts.down_proj.data
696
+ # TODO: remove this line after https://github.com/huggingface/transformers/pull/43622
697
+ hf_moe.experts.config._experts_implementation = "eager"
698
699
# Generate input batch of dimensions [batch_size, seq_len, hidden_dim]
700
hf_inputs = torch.randn((1, 64, config.hidden_size)).to(dtype).to("cuda")
@@ -718,7 +726,10 @@ def test_mixtral_moe(
726
get_forward_context().all_moe_layers = None
727
728
# Run forward passes for both MoE blocks
721
- hf_states, _ = hf_moe.forward(hf_inputs)
729
+ hf_states = hf_moe.forward(hf_inputs)
730
+ if isinstance(hf_states, tuple):
731
+ # Transformers v4
732
+ hf_states = hf_states[0]
733
vllm_states = vllm_moe.forward(vllm_inputs)
734
735
mixtral_moe_tol = {