From b59be5f829d8e73a6873fc7a9cc6ce767b783510 Mon Sep 17 00:00:00 2001 From: datta0 Date: Wed, 2 Apr 2025 06:49:06 +0000 Subject: [PATCH] add comments and use modified function --- unsloth/models/qwen3_moe.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/unsloth/models/qwen3_moe.py b/unsloth/models/qwen3_moe.py index 055c2dc907..319e3bffe3 100644 --- a/unsloth/models/qwen3_moe.py +++ b/unsloth/models/qwen3_moe.py @@ -76,7 +76,7 @@ def Qwen3MoeSparseMoeBlock_fast_forward(self, X, temp_gate = None, temp_up = Non # the current expert. We need to make sure to multiply the output hidden # states by `routing_weights` on the corresponding tokens (top-1 and top-2) current_state = X[None, top_x].reshape(-1, hd) - current_X = expert_layer(current_state) * routing_weights[top_x, idx, None] + current_X = expert_layer(current_state) * routing_weights[top_x, idx, None] # Qwen3MoeMLP.forward = fast_swiglu_inference takes care of making this faster. Analogous to Dense models' MLP # However `index_add_` only support torch tensors for indexing so we'll use # the `top_x` tensor here. @@ -119,10 +119,10 @@ def Qwen3MoeDecoderLayer_fast_forward( ) hidden_states = residual + hidden_states - # Fully Connected + # MoE Router MLP residual = hidden_states hidden_states = fast_rms_layernorm_inference(self.post_attention_layernorm, hidden_states) - hidden_states = Qwen3MoeSparseMoeBlock_fast_forward(self.mlp, hidden_states) + hidden_states, router_logits = Qwen3MoeSparseMoeBlock_fast_forward(self.mlp, hidden_states) hidden_states = residual + hidden_states else: residual = hidden_states @@ -140,15 +140,16 @@ def Qwen3MoeDecoderLayer_fast_forward( ) hidden_states = residual + hidden_states - # Fully Connected + # MoE Router MLP residual = hidden_states hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states) - hidden_states = self.mlp(hidden_states) + hidden_states, router_logits = self.mlp(hidden_states) hidden_states = residual + hidden_states pass outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) + if output_router_logits: outputs += (router_logits,) if use_cache: outputs += (present_key_value,) return outputs @@ -173,7 +174,7 @@ class FastQwen3MoeModel(FastQwen3Model): # Qwen3FlashAttention2 .forward = Qwen3Attention_fast_forward Qwen3MoeSparseMoeBlock .forward = Qwen3MoeSparseMoeBlock_fast_forward Qwen3MoeMLP .forward = fast_swiglu_inference # This is analogous to Dense models' MLP - Qwen3MoeDecoderLayer .forward = LlamaDecoderLayer_fast_forward + Qwen3MoeDecoderLayer .forward = Qwen3MoeDecoderLayer_fast_forward Qwen3MoeModel .forward = LlamaModel_fast_forward Qwen3MoeForCausalLM .forward = CausalLM_fast_forward(LlamaModel_fast_forward_inference) PeftModelForCausalLM.forward = PeftModelForCausalLM_fast_forward