Hugging Face's logo

Phase-Technologies
/
forge-500m-base

Text Generation
PyTorch
English
nlp
causal-lm
ssm
mamba
Mixture of Experts
custom-architecture
Model card Files
xet
 Community
forge-500m-base / model /hse_layer.py
Phase-Technologies's picture
Phase-Technologies
Add model architecture component: hse_layer.py
b56490b verified
Raw History Blame Contribute Delete
11.5 kB
"""
HSE (Hierarchical Sparse Expert) FFN Layer.
Two-tier MoE:
Tier 1: Select 1 semantic domain from K1=4 domains
Tier 2: Select top-2 specialists from K2=8 experts within that domain
Total: 32 experts, 2 active per token
Load balanced at both tiers independently for stable training.
"""
from __future__ import annotations
import math
from typing import Optional, Tuple, Dict
import torch
import torch.nn as nn
import torch.nn.functional as F
from .dgn_norm import build_norm
from .triton_kernels import fused_swiglu
class FeedForwardExpert(nn.Module):
"""Single SwiGLU expert in the HSE layer."""
def __init__(self, d_model: int, d_ff: int):
super().__init__()
self.gate = nn.Linear(d_model, d_ff, bias=False)
self.up = nn.Linear(d_model, d_ff, bias=False)
self.down = nn.Linear(d_ff, d_model, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.down(fused_swiglu(self.gate(x), self.up(x)))
class DenseSwiGLUFFN(nn.Module):
"""Standard Dense SwiGLU FFN (used in alternating layers)."""
def __init__(
self,
d_model: int = 1280,
d_ff: int = 3200,
norm_type: str = "dgn",
dgn_n_groups: int = 16,
norm_eps: float = 1e-6,
):
super().__init__()
self.norm = build_norm(norm_type, d_model, dgn_n_groups, norm_eps)
self.gate = nn.Linear(d_model, d_ff, bias=False)
self.up = nn.Linear(d_model, d_ff, bias=False)
self.down = nn.Linear(d_ff, d_model, bias=False)
nn.init.normal_(self.gate.weight, std=0.02)
nn.init.normal_(self.up.weight, std=0.02)
nn.init.normal_(self.down.weight, std=0.02 / math.sqrt(2))
def forward(self, x: torch.Tensor) -> torch.Tensor:
residual = x
h = self.norm(x)
return residual + self.down(fused_swiglu(self.gate(h), self.up(h)))
class HSELayer(nn.Module):
"""
Hierarchical Sparse Expert FFN Layer.
Architecture:
pre-norm β†’ Tier1-router β†’ select domain d*
β†’ Tier2-router[d*] β†’ select top-2 experts
β†’ weighted sum of 2 expert outputs
β†’ residual add
"""
def __init__(
self,
d_model: int = 1280,
n_domains: int = 4,
n_experts_per_domain: int = 8,
top_k: int = 2,
d_ff_expert: int = 288,
capacity_factor: float = 1.25,
expert_dropout: float = 0.1,
aux_loss_alpha: float = 0.01,
gumbel_tau_init: float = 1.0,
norm_type: str = "dgn",
dgn_n_groups: int = 16,
norm_eps: float = 1e-6,
layer_idx: int = 0,
):
super().__init__()
assert top_k <= n_experts_per_domain
self.d_model = d_model
self.n_domains = n_domains
self.n_experts_per_domain = n_experts_per_domain
self.n_experts_total = n_domains * n_experts_per_domain
self.top_k = top_k
self.d_ff_expert = d_ff_expert
self.capacity_factor = capacity_factor
self.expert_dropout = expert_dropout
self.aux_loss_alpha = aux_loss_alpha
self.layer_idx = layer_idx
# Pre-norm
self.norm = build_norm(norm_type, d_model, dgn_n_groups, norm_eps)
# ── Tier-1 Router ────────────────────────────────────────────────────
# Selects semantic domain: d* = argmax(W_r1 Β· x / Ο„)
self.tier1_router = nn.Linear(d_model, n_domains, bias=False)
self.gumbel_tau = nn.Parameter(
torch.tensor(gumbel_tau_init), requires_grad=False # annealed externally
)
# ── Tier-2 Routers (one per domain) ──────────────────────────────────
self.tier2_routers = nn.ModuleList([
nn.Linear(d_model, n_experts_per_domain, bias=False)
for _ in range(n_domains)
])
# ── Experts ───────────────────────────────────────────────────────────
self.experts = nn.ModuleList([
FeedForwardExpert(d_model, d_ff_expert)
for _ in range(self.n_experts_total)
])
# ── Dropout for expert regularization ────────────────────────────────
self.expert_dropout_layer = nn.Dropout(expert_dropout)
# ── Aux loss accumulator (for logging) ────────────────────────────────
self._last_aux_loss: Optional[torch.Tensor] = None
self._init_weights()
def _init_weights(self):
"""Initialize router weights to near-uniform distribution."""
nn.init.normal_(self.tier1_router.weight, std=0.02)
for router in self.tier2_routers:
nn.init.normal_(router.weight, std=0.02)
for expert in self.experts:
nn.init.normal_(expert.gate.weight, std=0.02)
nn.init.normal_(expert.up.weight, std=0.02)
nn.init.normal_(expert.down.weight, std=0.02 / math.sqrt(2))
def update_gumbel_tau(self, progress: float):
"""
Anneal Gumbel-softmax temperature from init to final.
progress: float in [0, 1] representing training progress.
"""
tau_init = 1.0
tau_final = 0.1
tau = tau_init * (tau_final / tau_init) ** progress
self.gumbel_tau.data.fill_(tau)
def _compute_aux_loss(
self,
tier1_probs: torch.Tensor, # (N, n_domains)
domain_idx: torch.Tensor, # (N,)
x_flat: torch.Tensor, # (N, d_model)
) -> torch.Tensor:
"""
Hierarchical load-balance auxiliary loss.
Loss = Ξ± * [Tier1_balance + Tier2_balance]
Tier1_balance = n_domains * Ξ£_d [f_d * P_d]
Tier2_balance = Ξ£_d {n_exp * Ξ£_e [f_{d,e} * P_{d,e}]} (over assigned tokens)
"""
N = x_flat.shape[0]
# Tier-1 balance
# f_d: fraction of tokens routed to domain d
one_hot_d = F.one_hot(domain_idx, self.n_domains).float() # (N, n_domains)
f1 = one_hot_d.mean(0) # (n_domains,)
P1 = tier1_probs.mean(0) # (n_domains,)
loss_t1 = self.n_domains * (f1 * P1).sum()
# Tier-2 balance (over tokens in each domain)
loss_t2 = torch.zeros(1, device=x_flat.device)
for d in range(self.n_domains):
mask = (domain_idx == d)
n_d = mask.sum()
if n_d < 2:
continue
x_d = x_flat[mask] # (n_d, d_model)
t2_logits = self.tier2_routers[d](x_d) # (n_d, n_exp)
t2_probs = F.softmax(t2_logits, dim=-1) # (n_d, n_exp)
# Top-k indices for f computation
topk_idx = t2_logits.topk(self.top_k, dim=-1).indices # (n_d, top_k)
one_hot_e = F.one_hot(topk_idx, self.n_experts_per_domain).float() # (n_d, k, n_exp)
f2 = one_hot_e.sum(1).mean(0) / self.top_k # (n_exp,)
P2 = t2_probs.mean(0) # (n_exp,)
loss_t2 = loss_t2 + self.n_experts_per_domain * (f2 * P2).sum()
return self.aux_loss_alpha * (loss_t1 + loss_t2)
def forward(
self,
x: torch.Tensor, # (B, T, d_model)
return_aux_loss: bool = True,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
"""
Forward pass with hierarchical routing (vectorized for maximum GPU throughput).
Returns (output, aux_loss) where aux_loss is for load balancing.
"""
B, T, D = x.shape
N = B * T
residual = x
x_norm = self.norm(x)
x_flat = x_norm.reshape(N, D) # (N, D)
# ── Tier-1 Routing ────────────────────────────────────────────────────
tier1_logits = self.tier1_router(x_flat) # (N, n_domains)
tier1_probs = F.softmax(tier1_logits / float(self.gumbel_tau), dim=-1)
if self.training:
domain_idx = F.gumbel_softmax(
tier1_logits,
tau=float(self.gumbel_tau),
hard=True
).argmax(dim=-1) # (N,)
else:
domain_idx = tier1_probs.argmax(dim=-1) # (N,)
# ── Tier-2 Routing (Zero-sync batched projection + gather) ───────────
t2_weights = torch.cat([r.weight for r in self.tier2_routers], dim=0) # (32, D)
t2_all = F.linear(x_flat, t2_weights).view(N, self.n_domains, self.n_experts_per_domain)
d_expanded = domain_idx.view(N, 1, 1).expand(-1, 1, self.n_experts_per_domain)
t2_logits = torch.gather(t2_all, 1, d_expanded).squeeze(1) # (N, 8)
t2_probs = F.softmax(t2_logits, dim=-1)
w, topk_idx = t2_probs.topk(self.top_k, dim=-1)
w = (w / (w.sum(-1, keepdim=True) + 1e-9)).to(x.dtype)
global_exp_idx = topk_idx + (domain_idx.unsqueeze(-1) * self.n_experts_per_domain) # (N, 2)
# ── Vectorized Token-Expert Dispatch ──────────────────────────────────
flat_exp_ids = global_exp_idx.reshape(-1)
flat_tokens = torch.arange(N, device=x.device).unsqueeze(1).expand(-1, self.top_k).reshape(-1)
flat_weights = w.reshape(-1, 1)
sort_perm = torch.argsort(flat_exp_ids)
sorted_exp = flat_exp_ids[sort_perm]
sorted_tokens = flat_tokens[sort_perm]
sorted_weights = flat_weights[sort_perm]
counts = torch.bincount(sorted_exp, minlength=self.n_experts_total).tolist()
out_flat = torch.zeros(N, D, device=x.device, dtype=x.dtype)
offset = 0
for e in range(self.n_experts_total):
cnt = counts[e]
if cnt > 0:
t_idx = sorted_tokens[offset : offset + cnt]
w_e = sorted_weights[offset : offset + cnt]
x_e = x_flat[t_idx]
if self.training and self.expert_dropout > 0:
x_e = self.expert_dropout_layer(x_e)
y_e = self.experts[e](x_e)
out_flat.index_add_(0, t_idx, w_e * y_e)
offset += cnt
output = out_flat.reshape(B, T, D)
# ── Auxiliary loss ────────────────────────────────────────────────────
aux_loss = None
if return_aux_loss and self.training:
aux_loss = self._compute_aux_loss(tier1_probs, domain_idx, x_flat)
self._last_aux_loss = aux_loss.detach()
return residual + output, aux_loss