diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index 06e170dfd00e..c96a274e1a4f 100644 --- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py @@ -32,7 +32,7 @@ class GlmMoeDsaConfig(PreTrainedConfig): n_group (`int`, *optional*, defaults to 1): Number of groups for routed experts. mlp_layer_types (`list`, *optional*): - MLP type pattern for each layer (`"dense"` or `"sparse"`). Defaults to 3 dense + rest sparse. + MLP type pattern for each layer (`"dense"` or `"sparse"`). Defaults to `3` dense layers and then every `moe_layer_freq`-th layer sparse. index_topk (`int`, *optional*, defaults to 2048): Number of top tokens selected by the indexer for sparse attention. index_head_dim (`int`, *optional*, defaults to 128): @@ -40,7 +40,7 @@ class GlmMoeDsaConfig(PreTrainedConfig): index_n_heads (`int | None`, *optional*, defaults to 32): Number of heads for the indexer projections (DSA). indexer_types (`list[str]`, *optional*): - Indexer mode for each layer (`"full"` or `"shared"`). Defaults to first layer full, then every `index_topk_freq`-th layer full, rest shared. + Indexer mode for each layer (`"full"` or `"shared"`). Defaults to the pattern derived from `index_topk_freq` and `index_skip_topk_offset`. ```python >>> from transformers import GlmMoeDsaConfig, GlmMoeDsaModel @@ -78,6 +78,7 @@ class GlmMoeDsaConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + attribute_map = { "num_local_experts": "n_routed_experts", } @@ -122,25 +123,23 @@ class GlmMoeDsaConfig(PreTrainedConfig): def __post_init__(self, **kwargs): self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim - - # MLP layer types: first 3 dense, rest sparse if self.mlp_layer_types is None: - self.mlp_layer_types = ["dense"] * min(3, self.num_hidden_layers) + ["sparse"] * ( - self.num_hidden_layers - 3 - ) + moe_layer_freq = kwargs.get("moe_layer_freq", 1) + self.mlp_layer_types = [ + "sparse" if i >= 3 and i % moe_layer_freq == 0 else "dense" for i in range(self.num_hidden_layers) + ] - # Indexer layer types if self.indexer_types is None: - pattern = kwargs.pop("index_topk_pattern", None) - freq = kwargs.pop("index_topk_freq", 1) + pattern = kwargs.get("index_topk_pattern") if pattern is not None: self.indexer_types = ( [{"F": "full", "S": "shared"}[c] for c in pattern] if isinstance(pattern, str) else list(pattern) ) else: - # First layer full, then every freq-th layer full, rest shared + freq = max(kwargs.get("index_topk_freq", 1), 1) + offset = kwargs.get("index_skip_topk_offset", 2) self.indexer_types = [ - "full" if (max(i - 1, 0) % freq) == 0 else "shared" for i in range(self.num_hidden_layers) + "full" if (max(i - offset + 1, 0) % freq) == 0 else "shared" for i in range(self.num_hidden_layers) ] super().__post_init__(**kwargs) diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index d5bc5e4aa08b..d6a7d5a16859 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -77,10 +77,12 @@ def apply_rotary_pos_emb( unsqueeze_dim: int = 1, ) -> torch.Tensor: """ - Applies Rotary Position Embedding to a single tensor. + Applies Rotary Position Embedding to a single tensor (query, key, or the indexer's q/k stream). - This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved)`. - Instead of using complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding. + This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved=True)`. + Instead of complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding. + Rotary pairs are always interpreted as adjacent elements (GPT-J style), so we first de-interleave `x` into + halves before applying the standard `rotate_half`. Args: x (`torch.Tensor`): Input tensor of shape `[..., head_dim]`. @@ -94,11 +96,9 @@ def apply_rotary_pos_emb( """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - - # Split-half (NeoX/Llama style): (x[:d/2], x[d/2:]) - # This matches llama's apply_rotary_pos_emb logic. - x_rotated = (x * cos) + (rotate_half(x) * sin) - return x_rotated + *leading_dims, head_dim = x.shape + x = x.view(*leading_dims, head_dim // 2, 2).transpose(-1, -2).reshape(*leading_dims, head_dim) + return (x * cos) + (rotate_half(x) * sin) class GlmMoeDsaIndexer(nn.Module): @@ -106,8 +106,8 @@ class GlmMoeDsaIndexer(nn.Module): DeepSeek Sparse Attention (DSA) indexer for selecting top-k tokens. The Indexer has its own lightweight projections (wq_b, wk) separate from the - main MLA attention. It uses non-interleaved (NeoX/Llama) RoPE, unlike the main attention - which uses interleaved RoPE. + main MLA attention. RoPE uses the same interleaved pair layout as main MLA + attention. **Cache strategy**: The Indexer manages its own key cache (`_cached_keys`) separately from the DynamicCache used by MLA attention, since DynamicCache is sized for exactly @@ -281,6 +281,9 @@ class GlmMoeDsaAttention(nn.Module): and SDPA backends. The reference's compressed-cache decode path (which avoids the kv_b_proj expansion at decode time) is a future optimization that would require a dedicated MLA cache class. + **DSA layer sharing**: full layers run the indexer; shared layers reuse the previous full + layer's top-k indices (`skip_topk` / `next_skip_topk`, derived from `config.indexer_types`). + **FP8 compatibility**: all weight accesses use standard nn.Linear forward calls (never raw `.weight` access), so FP8-quantized checkpoints work transparently. """ @@ -332,15 +335,15 @@ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int): self.scaling = self.qk_head_dim ** (-0.5) - self.indexer = GlmMoeDsaIndexer(config, layer_idx) - # Refer: https://arxiv.org/abs/2603.12201 for more details. - # skip_topk: when True, this layer will skip computation and reuse previous layer's topk indices. - # next_skip_topk: when True, the next layer will skip computation and reuse this layer's topk indices. + # skip_topk: when True, this layer reuses the previous full indexer layer's top-k indices. + # next_skip_topk: when True, the next layer reuses this layer's top-k indices. + # Shared layers have no indexer of their own. self.skip_topk = config.indexer_types[layer_idx] == "shared" self.next_skip_topk = ( config.indexer_types[layer_idx + 1] == "shared" if layer_idx < len(config.indexer_types) - 1 else False ) + self.indexer = None if self.skip_topk else GlmMoeDsaIndexer(config, layer_idx) def forward( self, @@ -364,7 +367,6 @@ def forward( query_states = query_states.view(batch_size, seq_length, -1, self.qk_head_dim).transpose(1, 2) # Split nope/rope, apply RoPE, recombine — layout: [B, H, S, D] q_nope, q_pe = torch.split(query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # BHSD format # ===== KV path ===== compressed_kv = self.kv_a_proj_with_mqa(hidden_states) # [B, S, kv_rank + rope_D] @@ -378,9 +380,10 @@ def forward( k_nope = k_nope.transpose(1, 2) # [B, H, S, nope_D] value_states = value_states.transpose(1, 2) # [B, H, S, v_D] - # RoPE on k_pe (single-head rope stream) + # RoPE on q_pe / k_pe (single-head rope stream for k), using interleaved pair layout. k_pe = k_pe.view(batch_size, 1, seq_length, self.qk_rope_head_dim) # [B, 1, S, rope_D] - k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # BHSD format + q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) + k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) k_pe = k_pe.expand(-1, k_nope.shape[1], -1, -1) # [B, H, S, rope_D] # Assemble full Q and K @@ -394,6 +397,8 @@ def forward( # ===== Indexer (DSA sparse mask) ===== # attention_mask is [B, 1, S, T] (4D) for eager and (2D) otherwise but indexer works with [B, S, T] (3D) if not self.skip_topk or prev_topk_indices is None: + if self.indexer is None: + raise ValueError("Shared DSA layers require top-k indices from a previous full indexer layer.") indexer_mask = ( attention_mask[:, 0, :, :] if attention_mask is not None and attention_mask.dim() == 4 diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index 7937be10e606..28d8dd6e17f6 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -53,10 +53,12 @@ def apply_rotary_pos_emb( unsqueeze_dim: int = 1, ) -> torch.Tensor: """ - Applies Rotary Position Embedding to a single tensor. + Applies Rotary Position Embedding to a single tensor (query, key, or the indexer's q/k stream). - This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved)`. - Instead of using complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding. + This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved=True)`. + Instead of complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding. + Rotary pairs are always interpreted as adjacent elements (GPT-J style), so we first de-interleave `x` into + halves before applying the standard `rotate_half`. Args: x (`torch.Tensor`): Input tensor of shape `[..., head_dim]`. @@ -70,11 +72,9 @@ def apply_rotary_pos_emb( """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - - # Split-half (NeoX/Llama style): (x[:d/2], x[d/2:]) - # This matches llama's apply_rotary_pos_emb logic. - x_rotated = (x * cos) + (rotate_half(x) * sin) - return x_rotated + *leading_dims, head_dim = x.shape + x = x.view(*leading_dims, head_dim // 2, 2).transpose(-1, -2).reshape(*leading_dims, head_dim) + return (x * cos) + (rotate_half(x) * sin) @auto_docstring(checkpoint="zai-org/GLM-5") @@ -84,7 +84,7 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig): n_group (`int`, *optional*, defaults to 1): Number of groups for routed experts. mlp_layer_types (`list`, *optional*): - MLP type pattern for each layer (`"dense"` or `"sparse"`). Defaults to 3 dense + rest sparse. + MLP type pattern for each layer (`"dense"` or `"sparse"`). Defaults to `3` dense layers and then every `moe_layer_freq`-th layer sparse. index_topk (`int`, *optional*, defaults to 2048): Number of top tokens selected by the indexer for sparse attention. index_head_dim (`int`, *optional*, defaults to 128): @@ -92,7 +92,7 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig): index_n_heads (`int | None`, *optional*, defaults to 32): Number of heads for the indexer projections (DSA). indexer_types (`list[str]`, *optional*): - Indexer mode for each layer (`"full"` or `"shared"`). Defaults to first layer full, then every `index_topk_freq`-th layer full, rest shared. + Indexer mode for each layer (`"full"` or `"shared"`). Defaults to the pattern derived from `index_topk_freq` and `index_skip_topk_offset`. ```python >>> from transformers import GlmMoeDsaConfig, GlmMoeDsaModel @@ -123,6 +123,10 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig): "layers.*.mlp.down_proj": "rowwise", } + attribute_map = { + "num_local_experts": "n_routed_experts", + } + hidden_size: int = 6144 intermediate_size: int = 12288 moe_intermediate_size: int = 2048 @@ -136,34 +140,29 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig): index_topk: int = 2048 index_head_dim: int = 128 index_n_heads: int = 32 + indexer_types: list[str] | None = None pretraining_tp = AttributeError() rope_interleave = AttributeError() - indexer_types: list[str] | None = None - attribute_map = { - "num_local_experts": "n_routed_experts", - } def __post_init__(self, **kwargs): self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim - - # MLP layer types: first 3 dense, rest sparse if self.mlp_layer_types is None: - self.mlp_layer_types = ["dense"] * min(3, self.num_hidden_layers) + ["sparse"] * ( - self.num_hidden_layers - 3 - ) + moe_layer_freq = kwargs.get("moe_layer_freq", 1) + self.mlp_layer_types = [ + "sparse" if i >= 3 and i % moe_layer_freq == 0 else "dense" for i in range(self.num_hidden_layers) + ] - # Indexer layer types if self.indexer_types is None: - pattern = kwargs.pop("index_topk_pattern", None) - freq = kwargs.pop("index_topk_freq", 1) + pattern = kwargs.get("index_topk_pattern") if pattern is not None: self.indexer_types = ( [{"F": "full", "S": "shared"}[c] for c in pattern] if isinstance(pattern, str) else list(pattern) ) else: - # First layer full, then every freq-th layer full, rest shared + freq = max(kwargs.get("index_topk_freq", 1), 1) + offset = kwargs.get("index_skip_topk_offset", 2) self.indexer_types = [ - "full" if (max(i - 1, 0) % freq) == 0 else "shared" for i in range(self.num_hidden_layers) + "full" if (max(i - offset + 1, 0) % freq) == 0 else "shared" for i in range(self.num_hidden_layers) ] PreTrainedConfig.__post_init__(self, **kwargs) @@ -177,8 +176,8 @@ class GlmMoeDsaIndexer(nn.Module): DeepSeek Sparse Attention (DSA) indexer for selecting top-k tokens. The Indexer has its own lightweight projections (wq_b, wk) separate from the - main MLA attention. It uses non-interleaved (NeoX/Llama) RoPE, unlike the main attention - which uses interleaved RoPE. + main MLA attention. RoPE uses the same interleaved pair layout as main MLA + attention. **Cache strategy**: The Indexer manages its own key cache (`_cached_keys`) separately from the DynamicCache used by MLA attention, since DynamicCache is sized for exactly @@ -315,6 +314,9 @@ class GlmMoeDsaAttention(nn.Module): and SDPA backends. The reference's compressed-cache decode path (which avoids the kv_b_proj expansion at decode time) is a future optimization that would require a dedicated MLA cache class. + **DSA layer sharing**: full layers run the indexer; shared layers reuse the previous full + layer's top-k indices (`skip_topk` / `next_skip_topk`, derived from `config.indexer_types`). + **FP8 compatibility**: all weight accesses use standard nn.Linear forward calls (never raw `.weight` access), so FP8-quantized checkpoints work transparently. """ @@ -366,15 +368,15 @@ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int): self.scaling = self.qk_head_dim ** (-0.5) - self.indexer = GlmMoeDsaIndexer(config, layer_idx) - # Refer: https://arxiv.org/abs/2603.12201 for more details. - # skip_topk: when True, this layer will skip computation and reuse previous layer's topk indices. - # next_skip_topk: when True, the next layer will skip computation and reuse this layer's topk indices. + # skip_topk: when True, this layer reuses the previous full indexer layer's top-k indices. + # next_skip_topk: when True, the next layer reuses this layer's top-k indices. + # Shared layers have no indexer of their own. self.skip_topk = config.indexer_types[layer_idx] == "shared" self.next_skip_topk = ( config.indexer_types[layer_idx + 1] == "shared" if layer_idx < len(config.indexer_types) - 1 else False ) + self.indexer = None if self.skip_topk else GlmMoeDsaIndexer(config, layer_idx) def forward( self, @@ -398,7 +400,6 @@ def forward( query_states = query_states.view(batch_size, seq_length, -1, self.qk_head_dim).transpose(1, 2) # Split nope/rope, apply RoPE, recombine — layout: [B, H, S, D] q_nope, q_pe = torch.split(query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # BHSD format # ===== KV path ===== compressed_kv = self.kv_a_proj_with_mqa(hidden_states) # [B, S, kv_rank + rope_D] @@ -412,9 +413,10 @@ def forward( k_nope = k_nope.transpose(1, 2) # [B, H, S, nope_D] value_states = value_states.transpose(1, 2) # [B, H, S, v_D] - # RoPE on k_pe (single-head rope stream) + # RoPE on q_pe / k_pe (single-head rope stream for k), using interleaved pair layout. k_pe = k_pe.view(batch_size, 1, seq_length, self.qk_rope_head_dim) # [B, 1, S, rope_D] - k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # BHSD format + q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) + k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) k_pe = k_pe.expand(-1, k_nope.shape[1], -1, -1) # [B, H, S, rope_D] # Assemble full Q and K @@ -428,6 +430,8 @@ def forward( # ===== Indexer (DSA sparse mask) ===== # attention_mask is [B, 1, S, T] (4D) for eager and (2D) otherwise but indexer works with [B, S, T] (3D) if not self.skip_topk or prev_topk_indices is None: + if self.indexer is None: + raise ValueError("Shared DSA layers require top-k indices from a previous full indexer layer.") indexer_mask = ( attention_mask[:, 0, :, :] if attention_mask is not None and attention_mask.dim() == 4 diff --git a/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py b/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py index 8b3a734af01a..e52c9aec4f13 100644 --- a/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py +++ b/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py @@ -103,6 +103,13 @@ def test_default_mlp_layer_types(self): config.mlp_layer_types, ["dense", "dense", "dense", "sparse", "sparse", "sparse", "sparse", "sparse"] ) + def test_indexer_types_respect_skip_topk_offset(self): + config = GlmMoeDsaConfig(num_hidden_layers=8, index_topk_freq=4, index_skip_topk_offset=3) + self.assertEqual( + config.indexer_types, + ["full", "full", "full", "shared", "shared", "shared", "full", "shared"], + ) + @parameterized.expand(TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION) @unittest.skip("Won't fix: Blip2 + T5 backbone needs custom input preparation for this test") def test_eager_matches_sdpa_inference(self, *args):