- in
- [B batch, T positions, 6144 d]
- out
- [B batch, T positions, 6144 d]
- parameters
- 9,702,998,272 · × 75 MoE blocks = 727,724,870,400
Parameters
9,702,998,272 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Router | 1,573,120 across its components | 1,573,120 |
| Routed expert | 9,663,676,416 | |
| Shared expert | 37,748,736 across its components | 37,748,736 |
| Per-instance total | 9,702,998,272 | |
Across 75 independent MoE blocks: parameters.
Derived from the pinned configuration and transformers implementation, and equal to the count of the checkpoint’s tensors read from its safetensors headers, less the multi-token-prediction layer the implementation does not load. The total includes the routers’ 19,200 correction-bias values, which the implementation holds as buffers.
In numbers
Weights · one MoE block
Held: 256 routed experts × 37,748,736 parameters, the shared expert’s 37,748,736 and the router’s 1,573,120, 9,702,998,272 in all. Used per token: 8 routed experts, the shared expert and the router, 341,311,744 (3.5%).
bfloat16 experts; the router’s matrix in bfloat16 and its bias in float32.
config.json 7modeling_glm_moe_dsa.py 487–492modeling_glm_moe_dsa.py 647modeling_glm_moe_dsa.py 535–559modeling_glm_moe_dsa.py 576–583Explanation
The implementation flattens batch and positions into one token axis for the experts and restores the shape (modeling_glm_moe_dsa.py 580–581).
Codemodeling_glm_moe_dsa.py 535–583
class GlmMoeDsaMoE(nn.Module):
"""
A mixed expert module containing shared experts.
"""
def __init__(self, config: GlmMoeDsaConfig):
super().__init__()
self.config = config
self.experts = GlmMoeDsaExperts(config)
self.gate = GlmMoeDsaTopkRouter(config)
self.shared_experts = GlmMoeDsaMLP(
config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts
) def forward(
self,
hidden_states: torch.Tensor,
top_k_index: torch.Tensor,
top_k_weights: torch.Tensor,
) -> torch.Tensor:
final_hidden_states = torch.zeros_like(hidden_states)
with torch.no_grad():
expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts + 1)
expert_mask = expert_mask.permute(2, 1, 0)
expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
for expert_idx in expert_hit:
expert_idx = expert_idx[0]
if expert_idx == self.num_experts:
continue
top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
current_state = hidden_states[token_idx]
gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
current_hidden_states = self.act_fn(gate) * up
current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
return final_hidden_states def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
residuals = hidden_states
orig_shape = hidden_states.shape
_, topk_weights, topk_indices = 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_statesSources2
Direct evidence
The MoE layer holds the routed experts, the router and one shared expert MLP of n_shared_experts × moe_intermediate_size features.
moe-init · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 562–574
Read code excerpt
class GlmMoeDsaMoE(nn.Module):
"""
A mixed expert module containing shared experts.
"""
def __init__(self, config: GlmMoeDsaConfig):
super().__init__()
self.config = config
self.experts = GlmMoeDsaExperts(config)
self.gate = GlmMoeDsaTopkRouter(config)
self.shared_experts = GlmMoeDsaMLP(
config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts
)Direct evidence
Each token’s output is the router-weighted sum of its selected experts plus the shared expert applied to every token.
moe-forward · huggingface/transformers / modeling_glm_moe_dsa.py
Pinned revision 5e4d6304de55 ↗ · lines 535–583
Read code excerpt
def forward(
self,
hidden_states: torch.Tensor,
top_k_index: torch.Tensor,
top_k_weights: torch.Tensor,
) -> torch.Tensor:
final_hidden_states = torch.zeros_like(hidden_states)
with torch.no_grad():
expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts + 1)
expert_mask = expert_mask.permute(2, 1, 0)
expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
for expert_idx in expert_hit:
expert_idx = expert_idx[0]
if expert_idx == self.num_experts:
continue
top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
current_state = hidden_states[token_idx]
gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
current_hidden_states = self.act_fn(gate) * up
current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
return final_hidden_states… lines 560–575 omitted …
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
residuals = hidden_states
orig_shape = hidden_states.shape
_, topk_weights, topk_indices = 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