From 8a396aff6f6b95a83864fc572843deb94c3cfc8c Mon Sep 17 00:00:00 2001 From: Lysandre Date: Fri, 29 Aug 2025 10:38:03 +0200 Subject: [PATCH 1/3] DeepSeek-v3 Shared Shared --- .../deepseek_v3/modeling_deepseek_v3.py | 51 +++++++------- .../models/deepseek_v3/modular_deepseek_v3.py | 68 ++++++++----------- 2 files changed, 56 insertions(+), 63 deletions(-) diff --git a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py index f82806e4412d..1342d3cea676 100644 --- a/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modeling_deepseek_v3.py @@ -86,12 +86,11 @@ def forward(self, x, position_ids): class DeepseekV3MLP(nn.Module): - def __init__(self, config, hidden_size=None, intermediate_size=None): + def __init__(self, config, intermediate_size=None): super().__init__() self.config = config - self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) @@ -149,36 +148,25 @@ def forward(self, hidden_states): return topk_indices, topk_weights -class DeepseekV3MoE(nn.Module): - """ - A mixed expert module containing shared experts. - """ - +class DeepseekV3NaiveMoe(nn.ModuleList): def __init__(self, config): super().__init__() - self.config = config - self.experts = nn.ModuleList( - [ - DeepseekV3MLP(config, intermediate_size=config.moe_intermediate_size) - for _ in range(config.n_routed_experts) - ] - ) - self.gate = DeepseekV3TopkRouter(config) - self.shared_experts = DeepseekV3MLP( - config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts - ) + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_local_experts + for _ in range(self.num_experts): + self += [DeepseekV3MLP(config, intermediate_size=config.moe_intermediate_size)] - def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): + def forward(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): r""" CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused to not have to do a loop here (deepseek has 256 experts soooo yeah). """ final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) - expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) + expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self)) expert_mask = expert_mask.permute(2, 0, 1) - for expert_idx in range(len(self.experts)): - expert = self.experts[expert_idx] + for expert_idx in range(len(self)): + expert = self[expert_idx] mask = expert_mask[expert_idx] token_indices, weight_indices = torch.where(mask) @@ -194,12 +182,27 @@ def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weig # and all expert are "local" meaning we shard but we don't gather return final_hidden_states.type(hidden_states.dtype) + +class DeepseekV3MoE(nn.Module): + """ + A mixed expert module containing shared experts. + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.experts = DeepseekV3NaiveMoe(config) + self.gate = DeepseekV3TopkRouter(config) + self.shared_experts = DeepseekV3MLP( + config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts + ) + def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape topk_indices, topk_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) return hidden_states diff --git a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py index 2f16ad7f4af0..b866ef2fa11f 100644 --- a/src/transformers/models/deepseek_v3/modular_deepseek_v3.py +++ b/src/transformers/models/deepseek_v3/modular_deepseek_v3.py @@ -6,6 +6,7 @@ import torch.utils.checkpoint from torch import nn +from ..qwen2_moe.modeling_qwen2_moe import Qwen2MoeMLP from ...activations import ACT2FN from ...cache_utils import Cache from ...modeling_flash_attention_utils import FlashAttentionKwargs @@ -39,6 +40,10 @@ class DeepseekV3RotaryEmbedding(LlamaRotaryEmbedding): pass +class DeepseekV3MLP(Qwen2MoeMLP): + pass + + def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): r""" TODO let's just use the original freqcis computation to not have the view @@ -83,23 +88,6 @@ def yarn_get_mscale(scale=1, mscale=1): return 0.1 * mscale * math.log(scale) + 1.0 -class DeepseekV3MLP(nn.Module): - def __init__(self, config, hidden_size=None, intermediate_size=None): - super().__init__() - self.config = config - self.hidden_size = config.hidden_size if hidden_size is None else hidden_size - self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size - - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) - self.act_fn = ACT2FN[config.hidden_act] - - def forward(self, x): - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - return down_proj - - class DeepseekV3TopkRouter(nn.Module): def __init__(self, config): super().__init__() @@ -147,37 +135,25 @@ def forward(self, hidden_states): return topk_indices, topk_weights -# FIXME: update the experts -class DeepseekV3MoE(nn.Module): - """ - A mixed expert module containing shared experts. - """ - +class DeepseekV3NaiveMoe(nn.ModuleList): def __init__(self, config): super().__init__() - self.config = config - self.experts = nn.ModuleList( - [ - DeepseekV3MLP(config, intermediate_size=config.moe_intermediate_size) - for _ in range(config.n_routed_experts) - ] - ) - self.gate = DeepseekV3TopkRouter(config) - self.shared_experts = DeepseekV3MLP( - config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts - ) + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_local_experts + for _ in range(self.num_experts): + self += [DeepseekV3MLP(config, intermediate_size=config.moe_intermediate_size)] - def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): + def forward(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): r""" CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused to not have to do a loop here (deepseek has 256 experts soooo yeah). """ final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) - expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) + expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self)) expert_mask = expert_mask.permute(2, 0, 1) - for expert_idx in range(len(self.experts)): - expert = self.experts[expert_idx] + for expert_idx in range(len(self)): + expert = self[expert_idx] mask = expert_mask[expert_idx] token_indices, weight_indices = torch.where(mask) @@ -193,12 +169,26 @@ def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weig # and all expert are "local" meaning we shard but we don't gather return final_hidden_states.type(hidden_states.dtype) +class DeepseekV3MoE(nn.Module): + """ + A mixed expert module containing shared experts. + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.experts = DeepseekV3NaiveMoe(config) + self.gate = DeepseekV3TopkRouter(config) + self.shared_experts = DeepseekV3MLP( + config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts + ) + def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape topk_indices, topk_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) return hidden_states From 06ce53723e32baa91891487d951ebe75af100895 Mon Sep 17 00:00:00 2001 From: Lysandre Date: Fri, 29 Aug 2025 11:05:01 +0200 Subject: [PATCH 2/3] Dependents of DS3 --- .../models/dots1/modeling_dots1.py | 49 ++++++++++-------- .../models/glm4_moe/modeling_glm4_moe.py | 51 ++++++++++--------- 2 files changed, 54 insertions(+), 46 deletions(-) diff --git a/src/transformers/models/dots1/modeling_dots1.py b/src/transformers/models/dots1/modeling_dots1.py index f462e9d169fa..28aa28744a16 100644 --- a/src/transformers/models/dots1/modeling_dots1.py +++ b/src/transformers/models/dots1/modeling_dots1.py @@ -246,12 +246,11 @@ def forward( class Dots1MLP(nn.Module): - def __init__(self, config, hidden_size=None, intermediate_size=None): + def __init__(self, config, intermediate_size=None): super().__init__() self.config = config - self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) @@ -262,34 +261,25 @@ def forward(self, x): return down_proj -# FIXME: update to latest mlp interface -class Dots1MoE(nn.Module): - """ - A mixed expert module containing shared experts. - """ - +class Dots1NaiveMoe(nn.ModuleList): def __init__(self, config): super().__init__() - self.config = config - self.experts = nn.ModuleList( - [Dots1MLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.n_routed_experts)] - ) - self.gate = Dots1TopkRouter(config) - self.shared_experts = Dots1MLP( - config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts - ) + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_local_experts + for _ in range(self.num_experts): + self += [Dots1MLP(config, intermediate_size=config.moe_intermediate_size)] - def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): + def forward(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): r""" CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused to not have to do a loop here (deepseek has 256 experts soooo yeah). """ final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) - expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) + expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self)) expert_mask = expert_mask.permute(2, 0, 1) - for expert_idx in range(len(self.experts)): - expert = self.experts[expert_idx] + for expert_idx in range(len(self)): + expert = self[expert_idx] mask = expert_mask[expert_idx] token_indices, weight_indices = torch.where(mask) @@ -305,12 +295,27 @@ def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weig # and all expert are "local" meaning we shard but we don't gather return final_hidden_states.type(hidden_states.dtype) + +class Dots1MoE(nn.Module): + """ + A mixed expert module containing shared experts. + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.experts = Dots1NaiveMoe(config) + self.gate = Dots1TopkRouter(config) + self.shared_experts = Dots1MLP( + config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts + ) + def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape topk_indices, topk_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) return hidden_states diff --git a/src/transformers/models/glm4_moe/modeling_glm4_moe.py b/src/transformers/models/glm4_moe/modeling_glm4_moe.py index cb695ffbe638..99edbbd4a1ae 100644 --- a/src/transformers/models/glm4_moe/modeling_glm4_moe.py +++ b/src/transformers/models/glm4_moe/modeling_glm4_moe.py @@ -208,12 +208,11 @@ def forward( class Glm4MoeMLP(nn.Module): - def __init__(self, config, hidden_size=None, intermediate_size=None): + def __init__(self, config, intermediate_size=None): super().__init__() self.config = config - self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) @@ -292,36 +291,25 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" -class Glm4MoeMoE(nn.Module): - """ - A mixed expert module containing shared experts. - """ - +class Glm4MoeNaiveMoe(nn.ModuleList): def __init__(self, config): super().__init__() - self.config = config - self.experts = nn.ModuleList( - [ - Glm4MoeMLP(config, intermediate_size=config.moe_intermediate_size) - for _ in range(config.n_routed_experts) - ] - ) - self.gate = Glm4MoeTopkRouter(config) - self.shared_experts = Glm4MoeMLP( - config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts - ) + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_local_experts + for _ in range(self.num_experts): + self += [Glm4MoeMLP(config, intermediate_size=config.moe_intermediate_size)] - def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): + def forward(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): r""" CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused to not have to do a loop here (deepseek has 256 experts soooo yeah). """ final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) - expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) + expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self)) expert_mask = expert_mask.permute(2, 0, 1) - for expert_idx in range(len(self.experts)): - expert = self.experts[expert_idx] + for expert_idx in range(len(self)): + expert = self[expert_idx] mask = expert_mask[expert_idx] token_indices, weight_indices = torch.where(mask) @@ -337,12 +325,27 @@ def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weig # and all expert are "local" meaning we shard but we don't gather return final_hidden_states.type(hidden_states.dtype) + +class Glm4MoeMoE(nn.Module): + """ + A mixed expert module containing shared experts. + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.experts = Glm4MoeNaiveMoe(config) + self.gate = Glm4MoeTopkRouter(config) + self.shared_experts = Glm4MoeMLP( + config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts + ) + def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape topk_indices, topk_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) return hidden_states From cec498dfd7c893c7164af792037bc9aa67b94dc7 Mon Sep 17 00:00:00 2001 From: Lysandre Debut Date: Fri, 29 Aug 2025 11:54:20 +0200 Subject: [PATCH 3/3] Standardize GLM4V MoE (#40539) --- .../models/glm4v_moe/modeling_glm4v_moe.py | 50 +++++++++---------- .../models/glm4v_moe/modular_glm4v_moe.py | 25 +++++++--- 2 files changed, 41 insertions(+), 34 deletions(-) diff --git a/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py b/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py index bfa106324dc4..6068cb7f891d 100644 --- a/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py @@ -289,37 +289,25 @@ def forward(self, hidden_states): return topk_indices, topk_weights -# FIXME: refactor moe -class Glm4vMoeTextMoE(nn.Module): - """ - A mixed expert module containing shared experts. - """ - - def __init__(self, config: Glm4vMoeTextConfig): +class Glm4vMoeTextNaiveMoe(nn.ModuleList): + def __init__(self, config): super().__init__() - self.config = config - self.experts = nn.ModuleList( - [ - Glm4vMoeTextMLP(config, intermediate_size=config.moe_intermediate_size) - for _ in range(config.n_routed_experts) - ] - ) - self.gate = Glm4vMoeTextTopkRouter(config) - self.shared_experts = Glm4vMoeTextMLP( - config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts - ) + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_local_experts + for _ in range(self.num_experts): + self += [Glm4vMoeTextMLP(config, intermediate_size=config.moe_intermediate_size)] - def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): + def forward(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor): r""" CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused to not have to do a loop here (deepseek has 256 experts soooo yeah). """ final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype) - expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts)) + expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self)) expert_mask = expert_mask.permute(2, 0, 1) - for expert_idx in range(len(self.experts)): - expert = self.experts[expert_idx] + for expert_idx in range(len(self)): + expert = self[expert_idx] mask = expert_mask[expert_idx] token_indices, weight_indices = torch.where(mask) @@ -335,23 +323,33 @@ def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weig # and all expert are "local" meaning we shard but we don't gather return final_hidden_states.type(hidden_states.dtype) + +class Glm4vMoeTextMoE(nn.Module): + def __init__(self, config: Glm4vMoeTextConfig): + super().__init__(config) + self.config = config + self.experts = Glm4vMoeTextNaiveMoe(config) + self.gate = Glm4vMoeTextTopkRouter(config) + self.shared_experts = Glm4vMoeTextMLP( + config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts + ) + def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape topk_indices, topk_weights = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) return hidden_states class Glm4vMoeTextMLP(nn.Module): - def __init__(self, config, hidden_size=None, intermediate_size=None): + def __init__(self, config, intermediate_size=None): super().__init__() self.config = config - self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py index 4f333c2ef8b2..572756680985 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -17,6 +17,7 @@ import torch import torch.nn as nn +from ..deepseek_v3.modeling_deepseek_v3 import DeepseekV3NaiveMoe from ...cache_utils import Cache from ...configuration_utils import PretrainedConfig from ...modeling_flash_attention_utils import FlashAttentionKwargs @@ -408,22 +409,30 @@ def __init__(self, config: Glm4vMoeTextConfig): super().__init__(config) -# FIXME: update the expert class -class Glm4vMoeTextMoE(Glm4MoeMoE): +class Glm4vMoeTextNaiveMoe(DeepseekV3NaiveMoe): + pass + + +class Glm4vMoeTextMoE(nn.Module): def __init__(self, config: Glm4vMoeTextConfig): super().__init__(config) self.config = config - self.experts = nn.ModuleList( - [ - Glm4vMoeTextMLP(config, intermediate_size=config.moe_intermediate_size) - for _ in range(config.n_routed_experts) - ] - ) + self.experts = Glm4vMoeTextNaiveMoe(config) self.gate = Glm4vMoeTextTopkRouter(config) self.shared_experts = Glm4vMoeTextMLP( config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts ) + def forward(self, hidden_states): + residuals = hidden_states + orig_shape = hidden_states.shape + topk_indices, topk_weights = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = hidden_states + self.shared_experts(residuals) + return hidden_states + + class Glm4vMoeTextMLP(Glm4MoeMLP): pass