| |
| import torch, torch.nn as nn, torch.nn.functional as F |
|
|
| CHUNK = 16 |
|
|
| def wkv7_naive(r, k, v, a, b, logd, S=None): |
| B, T, H, N = r.shape |
| if S is None: S = r.new_zeros(B, H, N, N) |
| d = torch.exp(logd); outs = [] |
| for t in range(T): |
| sa = S @ a[:, t].unsqueeze(-1) |
| S = S * d[:, t].unsqueeze(-2) + sa @ b[:, t].unsqueeze(-2) + v[:, t].unsqueeze(-1) @ k[:, t].unsqueeze(-2) |
| outs.append((S @ r[:, t].unsqueeze(-1)).squeeze(-1)) |
| return torch.stack(outs, 1), S |
|
|
| def wkv7(r, k, v, a, b, logd, S=None, chunk=CHUNK): |
| B, T, H, N = r.shape |
| if S is None: S = r.new_zeros(B, H, N, N) |
| if T <= 4: return wkv7_naive(r, k, v, a, b, logd, S) |
| c = chunk; pad = (-T) % c |
| if pad: |
| r, k, v, a, b, logd = [F.pad(x, (0, 0, 0, 0, 0, pad)) for x in (r, k, v, a, b, logd)] |
| nc = (T + pad) // c |
| f = lambda x: x.reshape(B, nc, c, H, N).permute(0, 3, 1, 2, 4) |
| r, k, v, a, b, logd = map(f, (r, k, v, a, b, logd)) |
| cum = torch.cumsum(logd, 3) |
| at = torch.exp(cum - logd) * a; bt = b * torch.exp(-cum) |
| kt = k * torch.exp(-cum); rt = r * torch.exp(cum) |
| Dc = torch.exp(cum[:, :, :, -1]) |
| tri_s = torch.tril(torch.ones(c, c, dtype=r.dtype, device=r.device), -1) |
| tri = tri_s + torch.eye(c, dtype=r.dtype, device=r.device) |
| Bm = (at @ bt.transpose(-1, -2)) * tri_s |
| Km = (at @ kt.transpose(-1, -2)) * tri_s |
| Qb = (rt @ bt.transpose(-1, -2)) * tri |
| Qk = (rt @ kt.transpose(-1, -2)) * tri |
| L_m = torch.eye(c, dtype=r.dtype, device=r.device) - Bm |
| X = torch.linalg.solve_triangular(L_m, torch.cat([Km @ v, at], -1), upper=False) |
| Uloc, Wm = X[..., :N], X[..., N:] |
| Oloc = Qb @ Uloc + Qk @ v |
| P = rt + Qb @ Wm |
| Tm = torch.eye(N, dtype=r.dtype, device=r.device) + Wm.transpose(-1, -2) @ bt |
| Cm = Uloc.transpose(-1, -2) @ bt + v.transpose(-1, -2) @ kt |
| outs = [] |
| for n in range(nc): |
| outs.append(Oloc[:, :, n] + P[:, :, n] @ S.transpose(-1, -2)) |
| S = (S @ Tm[:, :, n] + Cm[:, :, n]) * Dc[:, :, n].unsqueeze(-2) |
| o = torch.stack(outs, 2).permute(0, 2, 3, 1, 4).reshape(B, nc * c, H, N) |
| return o[:, :T], S |
|
|
| class TimeMix(nn.Module): |
| def __init__(s, C, H, layer, n_layer): |
| super().__init__() |
| s.C, s.H, s.N, s.layer = C, H, C // H, layer |
| for n in ["r", "w", "k", "v", "a", "g"]: |
| s.register_parameter("x_" + n, nn.Parameter(torch.rand(1, 1, C) * 0.5)) |
| s.w0 = nn.Parameter(-6 + 5 * (torch.arange(C) / (C - 1)) ** 0.85) |
| s.w1 = nn.Parameter(torch.zeros(C, 32)); s.w2 = nn.Parameter(torch.randn(32, C) * 0.01) |
| s.a0 = nn.Parameter(torch.zeros(C)) |
| s.a1 = nn.Parameter(torch.zeros(C, 32)); s.a2 = nn.Parameter(torch.randn(32, C) * 0.01) |
| s.v0 = nn.Parameter(torch.ones(C)) |
| s.v1 = nn.Parameter(torch.zeros(C, 16)); s.v2 = nn.Parameter(torch.randn(16, C) * 0.01) |
| s.g1 = nn.Parameter(torch.zeros(C, 64)); s.g2 = nn.Parameter(torch.randn(64, C) * 0.01) |
| s.k_k = nn.Parameter(torch.full((C,), 0.85)) |
| s.k_a = nn.Parameter(torch.ones(C)) |
| s.r_k = nn.Parameter(torch.zeros(H, s.N)) |
| s.receptance = nn.Linear(C, C, bias=False) |
| s.key = nn.Linear(C, C, bias=False) |
| s.value = nn.Linear(C, C, bias=False) |
| s.output = nn.Linear(C, C, bias=False) |
| s.ln_x = nn.GroupNorm(H, C, eps=64e-5) |
|
|
| def forward(s, x, v_first, st=None): |
| B, T, C = x.shape; H, N = s.H, s.N |
| prev = x.new_zeros(B, 1, C) if st is None else st[0].unsqueeze(1) |
| xx = torch.cat([prev, x[:, :-1]], 1) - x |
| xr, xw, xk, xv, xa, xg = [x + xx * getattr(s, "x_" + n) for n in "rwkvag"] |
| r = s.receptance(xr) |
| w = -F.softplus(-(s.w0 + torch.tanh(xw @ s.w1) @ s.w2)) - 0.5 |
| k = s.key(xk); v = s.value(xv) |
| if s.layer == 0: v_first = v |
| else: v = v + (v_first - v) * torch.sigmoid(s.v0 + (xv @ s.v1) @ s.v2) |
| a = torch.sigmoid(s.a0 + (xa @ s.a1) @ s.a2) |
| g = torch.sigmoid(xg @ s.g1) @ s.g2 |
| kk = F.normalize((k * s.k_k).view(B, T, H, N), dim=-1).view(B, T, C) |
| k = k * (1 + (a - 1) * s.k_a) |
|
|
| sh = lambda t: t.view(B, T, H, N) |
| r_, k_, v_, a_, b_ = sh(r), sh(k), sh(v), sh(-kk), sh(kk * a) |
| logd = -torch.exp(sh(w)) |
| o, S = wkv7(r_, k_, v_, a_, b_, logd, None if st is None else st[1]) |
| o = s.ln_x(o.reshape(B * T, C)).view(B, T, C) |
| o = o + ((r_ * k_ * s.r_k).sum(-1, keepdim=True) * v_).view(B, T, C) |
| return s.output(o * g), v_first, (x[:, -1], S) |
|
|
| class ChannelMix(nn.Module): |
| def __init__(s, C): |
| super().__init__() |
| s.x_k = nn.Parameter(torch.rand(1, 1, C) * 0.5) |
| s.key = nn.Linear(C, C * 4, bias=False) |
| s.value = nn.Linear(C * 4, C, bias=False) |
| def forward(s, x, st=None): |
| prev = x.new_zeros(x.shape[0], 1, x.shape[2]) if st is None else st.unsqueeze(1) |
| xx = torch.cat([prev, x[:, :-1]], 1) - x |
| return s.value(torch.relu(s.key(x + xx * s.x_k)) ** 2), x[:, -1] |
|
|
| class Block(nn.Module): |
| def __init__(s, C, H, i, L): |
| super().__init__() |
| s.ln0 = nn.LayerNorm(C) if i == 0 else nn.Identity() |
| s.ln1, s.ln2 = nn.LayerNorm(C), nn.LayerNorm(C) |
| s.att, s.ffn = TimeMix(C, H, i, L), ChannelMix(C) |
| def forward(s, x, v_first, st=None): |
| x = s.ln0(x) |
| d, v_first, a_st = s.att(s.ln1(x), v_first, None if st is None else st[0]) |
| x = x + d |
| f, f_st = s.ffn(s.ln2(x), None if st is None else st[1]) |
| return x + f, v_first, (a_st, f_st) |
|
|
| class RWKV7(nn.Module): |
| |
| def __init__(s, vocab=8192, C=512, L=5, head_size=64): |
| super().__init__() |
| s.emb = nn.Embedding(vocab, C) |
| s.blocks = nn.ModuleList(Block(C, C // head_size, i, L) for i in range(L)) |
| s.ln_out = nn.LayerNorm(C) |
| s.head = nn.Linear(C, vocab, bias=False) |
| nn.init.uniform_(s.emb.weight, -1e-4, 1e-4) |
|
|
| def forward(s, idx, targets=None, state=None): |
| keep = state is not None |
| x, vf, new = s.emb(idx), None, [] |
| for i, b in enumerate(s.blocks): |
| x, vf, st = b(x, vf, state[i] if state else None) |
| new.append(st) |
| logits = s.head(s.ln_out(x)) |
| loss = None if targets is None else F.cross_entropy( |
| logits.view(-1, logits.size(-1)), targets.view(-1)) |
| return (logits, loss, new) if keep else (logits, loss) |
