| |
| import math |
| import torch, torch.nn as nn, torch.nn.functional as F |
|
|
| HEAD = 128 |
| NCAT, NFP = 8, 3 |
| |
| FAST = {'norm': True, 'rope': 'stack', 'qknorm': True, 'joint_qk': False, 'addcmul': False} |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(s, d): |
| super().__init__(); s.w = nn.Parameter(torch.ones(d)) |
|
|
| def forward(s, x): |
| if FAST['norm']: |
| return F.rms_norm(x, (x.shape[-1],), s.w, 1e-6) |
| return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + 1e-6).to(x.dtype) * s.w |
|
|
|
|
| def rope_table(side): |
| |
| f = 1.0 / (10000 ** (torch.arange(0, HEAD // 4) / (HEAD // 4))) |
| ys, xs = torch.meshgrid(torch.arange(side), torch.arange(side), indexing='ij') |
| ang = torch.cat([ys.reshape(-1, 1) * f, xs.reshape(-1, 1) * f], 1) |
| return torch.cos(ang), torch.sin(ang) |
|
|
|
|
| def rope(x, cos, sin): |
| if FAST['rope'] == 'off': |
| return x |
| if FAST['rope'] == 'half': |
| |
| c, s_ = torch.cat([cos, cos], -1).to(x.dtype), torch.cat([sin, sin], -1).to(x.dtype) |
| x1, x2 = x.chunk(2, -1) |
| return x * c + torch.cat([-x2, x1], -1) * s_ |
| if FAST['rope'] == 'stack': |
| a, b = x.unflatten(-1, (-1, 2)).unbind(-1) |
| cos, sin = cos.to(x.dtype), sin.to(x.dtype) |
| return torch.stack([a * cos - b * sin, a * sin + b * cos], -1).flatten(-2) |
| a, b = x.float().unflatten(-1, (-1, 2)).unbind(-1) |
| return torch.stack([a * cos - b * sin, a * sin + b * cos], -1).flatten(-2).to(x.dtype) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(s, h): |
| super().__init__() |
| s.h, s.nh, s.mlp = h, h // HEAD, 3 * h |
| s.norm = nn.LayerNorm(h, elementwise_affine=False, eps=1e-6) |
| s.mod = nn.Linear(h, 3 * h) |
| s.lin1 = nn.Linear(h, 3 * h + 2 * s.mlp) |
| s.qn, s.kn = RMSNorm(HEAD), RMSNorm(HEAD) |
| s.lin2 = nn.Linear(h + s.mlp, h) |
| nn.init.zeros_(s.mod.weight); nn.init.zeros_(s.mod.bias) |
|
|
| def forward(s, x, c, cos, sin): |
| B, T, _ = x.shape |
| sh, sc, g = s.mod(F.silu(c)).unsqueeze(1).chunk(3, -1) |
| xm = torch.addcmul(sh, s.norm(x), 1 + sc) if FAST['addcmul'] else s.norm(x) * (1 + sc) + sh |
| qkv, a, b = s.lin1(xm).split([3 * s.h, s.mlp, s.mlp], -1) |
| qkv = qkv.view(B, T, 3, s.nh, HEAD).permute(2, 0, 3, 1, 4) |
| v = qkv[2] |
| if FAST['joint_qk']: |
| qk = qkv[:2] |
| if FAST['qknorm']: |
| qk = F.rms_norm(qk, (HEAD,), eps=1e-6) * torch.stack([s.qn.w, s.kn.w]).view(2, 1, 1, 1, HEAD).to(qk.dtype) |
| q, k = rope(qk, cos, sin).unbind(0) |
| else: |
| q, k = qkv[0], qkv[1] |
| if FAST['qknorm']: |
| q, k = s.qn(q), s.kn(k) |
| q, k = rope(q, cos, sin), rope(k, cos, sin) |
| o = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(B, T, s.h) |
| out = s.lin2(torch.cat([o, F.silu(a) * b], -1)) |
| return torch.addcmul(x, g, out) if FAST['addcmul'] else x + g * out |
|
|
|
|
| class DiT(nn.Module): |
| def __init__(s, h, depth, patch, vocab=0): |
| super().__init__() |
| s.p, s.side = patch, 32 // patch |
| s.inp = nn.Linear(4 * patch * patch, h) |
| s.t1, s.t2 = nn.Linear(256, h), nn.Linear(h, h) |
| s.cat = nn.Embedding(NCAT + 1, h) |
| s.fp = nn.Embedding(NFP + 1, h) |
| if vocab: |
| s.wemb = nn.Embedding(vocab, h, padding_idx=0) |
| s.wproj = nn.Sequential(nn.Linear(h, h), nn.SiLU(), nn.Linear(h, h)) |
| s.blocks = nn.ModuleList([Block(h) for _ in range(depth)]) |
| s.fnorm = nn.LayerNorm(h, elementwise_affine=False, eps=1e-6) |
| s.fmod = nn.Linear(h, 2 * h) |
| s.outp = nn.Linear(h, 4 * patch * patch) |
| for m in (s.fmod, s.outp): |
| nn.init.zeros_(m.weight); nn.init.zeros_(m.bias) |
| cos, sin = rope_table(s.side) |
| s.register_buffer('cos', cos, persistent=False); s.register_buffer('sin', sin, persistent=False) |
|
|
| def forward(s, x, t, cat, fp, txt=None): |
| B, p, n = x.shape[0], s.p, s.side |
| x = x.unfold(2, p, p).unfold(3, p, p).permute(0, 2, 3, 1, 4, 5).reshape(B, n * n, 4 * p * p) |
| half = torch.exp(-math.log(10000) * torch.arange(128, device=x.device) / 128) |
| te = torch.cat([torch.cos(t[:, None] * 1000 * half), torch.sin(t[:, None] * 1000 * half)], -1) |
| c = s.t2(F.silu(s.t1(te))) + s.cat(cat) + s.fp(fp) |
| if txt is not None: |
| m = (txt > 0).to(c.dtype) |
| e = (s.wemb(txt) * m[..., None]).sum(1) / m.sum(1, keepdim=True).clamp(min=1) |
| c = c + s.wproj(e) * (m.sum(1, keepdim=True) > 0) |
| x = s.inp(x) |
| for blk in s.blocks: |
| x = blk(x, c, s.cos, s.sin) |
| sh, sc = s.fmod(F.silu(c)).unsqueeze(1).chunk(2, -1) |
| x = s.outp(s.fnorm(x) * (1 + sc) + sh) |
| return x.view(B, n, n, 4, p, p).permute(0, 3, 1, 4, 2, 5).reshape(B, 4, 32, 32) |
|
|
