ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update qwen3_moe.py

Browse source
This commit is contained in:
Daniel Han 2025-11-30 23:37:32 -08:00
commit 3ef6b66f39
1 changed files with 4 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

8
unsloth/models/qwen3_moe.py
View file

@ -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)