From 3ef6b66f39a339cf17d98fcbd85f14d07ff34dd0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 30 Nov 2025 23:37:32 -0800 Subject: [PATCH] Update qwen3_moe.py --- unsloth/models/qwen3_moe.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/unsloth/models/qwen3_moe.py b/unsloth/models/qwen3_moe.py index e80a87a950..f0b2b811ba 100644 --- a/unsloth/models/qwen3_moe.py +++ b/unsloth/models/qwen3_moe.py @@ -59,12 +59,12 @@ def Qwen3MoeSparseMoeBlock_fast_forward(self, X, temp_gate = None, temp_up = Non self.gate_proj, X, out = temp_gate ) # pretty much the only change from transformers implementation. - routing_weights = torch_nn_functional_softmax(router_logits, dim = -1) + routing_weights = torch_nn_functional_softmax(router_logits, dim = -1, dtype = torch.float32) routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim = -1) routing_weights /= routing_weights.sum(dim = -1, keepdim = True) # we cast back to the input dtype routing_weights = routing_weights.to(X.dtype) - final_X = torch.zeros((bsz * seq_len, hd), dtype = X.dtype, device = X.device) + final_X = torch.zeros((bsz * seq_len, hd), dtype = torch.float32, device = X.device) # One hot encode the selected experts to create an expert mask # this will be used to easily index which expert is going to be sollicitated @@ -128,7 +128,7 @@ def Qwen3MoeDecoderLayer_fast_forward( position_embeddings = position_embeddings, _flag_for_generation = self._flag_for_generation, ) - hidden_states = residual + hidden_states + hidden_states += hidden_states # MoE Router MLP residual = hidden_states @@ -138,7 +138,7 @@ def Qwen3MoeDecoderLayer_fast_forward( hidden_states, router_logits = Qwen3MoeSparseMoeBlock_fast_forward( self.mlp, hidden_states ) - hidden_states = residual + hidden_states + hidden_states += hidden_states else: residual = hidden_states hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states)