File size: 3,311 Bytes
8485f65 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | # SPDX-License-Identifier: CC-BY-NC-4.0
# Copyright (c) 2026 Jakub Ćwirlej
"""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) # fixed, saved with the model
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 # 6.25% of the neurons active per token
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}')
|