| |
| |
| """Usage example: a drop-in sparse SwiGLU FFN module with gate subsampling and load balancing. |
| |
| python example.py # CUDA: Triton kernels under torch.compile; CPU: dense masked fallback |
| """ |
| import torch |
| import torch.nn as nn |
|
|
| from eugen_sparse import (grouped_topk_select, kernel_supported, make_subsample_index, |
| sparse_swiglu_ffn, update_load_balance_bias) |
|
|
|
|
| class EugenSparseFFN(nn.Module): |
| """SwiGLU FFN computing only `active_groups` groups of `group_size` neurons per token.""" |
|
|
| def __init__(self, d_model, d_ff, group_size=32, active_groups=8, samples_per_group=8, layer_seed=0): |
| super().__init__() |
| self.G, self.kg = group_size, active_groups |
| 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) |
| sub, self.spg = make_subsample_index(d_ff, group_size, samples_per_group, seed=layer_seed) |
| ng = d_ff // group_size |
| self.register_buffer('subsample_idx', sub, persistent=True) |
| self.register_buffer('lb_bias', torch.zeros(ng), persistent=True) |
| self.register_buffer('lb_freq_ema', torch.full((ng,), active_groups / ng), persistent=True) |
| self.register_buffer('lb_freq_now', torch.zeros(ng), persistent=False) |
|
|
| def forward(self, x): |
| idx = grouped_topk_select(x, self.gate.weight, self.G, self.kg, self.subsample_idx, self.spg, |
| self.lb_bias) |
| if self.training: |
| with torch.no_grad(): |
| ng = self.lb_bias.numel() |
| counts = torch.zeros(ng, device=x.device).index_add_( |
| 0, idx.reshape(-1).long(), torch.ones(idx.numel(), device=x.device)) |
| self.lb_freq_now.copy_(counts / (idx.numel() // self.kg)) |
| return sparse_swiglu_ffn(x, self.gate.weight, self.up.weight, self.down.weight, idx, self.G) |
|
|
| @torch.no_grad() |
| def update_balance(self, gamma=1e-3, bias_max=0.5, momentum=0.99): |
| """Call after optimizer.step().""" |
| ng = self.lb_bias.numel() |
| update_load_balance_bias(self.lb_bias, self.lb_freq_ema, self.lb_freq_now, self.kg / ng, |
| gamma, bias_max, momentum) |
|
|
|
|
| if __name__ == '__main__': |
| dev = 'cuda' if torch.cuda.is_available() else 'cpu' |
| d_model, d_ff, G, kg = 512, 4096, 32, 8 |
| torch.manual_seed(0) |
| ffn = EugenSparseFFN(d_model, d_ff, G, kg).to(dev) |
| print(f'device {dev}: Triton kernel used = {kernel_supported(dev == "cuda", d_model, d_ff, G, kg)}') |
| opt = torch.optim.AdamW(ffn.parameters(), lr=1e-3) |
| step_fn = torch.compile(ffn) if dev == 'cuda' else ffn |
| for step in range(3): |
| x = torch.randn(4, 256, d_model, device=dev) |
| with torch.autocast(device_type=dev, dtype=torch.bfloat16, enabled=(dev == 'cuda')): |
| y = step_fn(x) |
| loss = y.float().pow(2).mean() |
| opt.zero_grad(); loss.backward(); opt.step() |
| ffn.update_balance() |
| print(f'step {step}: loss {loss.item():.5f} | most used group share {ffn.lb_freq_now.max().item():.3f}') |
|
|
