From 4f18296e40f616421f3ece53e6c53d1f15eb0e79 Mon Sep 17 00:00:00 2001 From: Colin Date: Mon, 3 Mar 2025 15:47:21 +0800 Subject: [PATCH] Format rwkv/RWKV-v7/rwkv_v7_demo.py --- rwkv/RWKV-v7/rwkv_v7_demo.py | 160 ++++++++++++++++++++--------------- 1 file changed, 91 insertions(+), 69 deletions(-) diff --git a/rwkv/RWKV-v7/rwkv_v7_demo.py b/rwkv/RWKV-v7/rwkv_v7_demo.py index cf49d4b..5ab5b25 100644 --- a/rwkv/RWKV-v7/rwkv_v7_demo.py +++ b/rwkv/RWKV-v7/rwkv_v7_demo.py @@ -5,7 +5,9 @@ import torch, types, os, gc, math, json import numpy as np import torch.nn as nn +from torch.nn import Module from torch.nn import functional as F + np.set_printoptions(precision=4, suppress=True, linewidth=200) torch.backends.cudnn.benchmark = True torch.backends.cudnn.allow_tf32 = True @@ -14,9 +16,9 @@ torch.backends.cuda.matmul.allow_tf32 = True # torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = True torch._C._jit_set_autocast_mode(False) -''' +""" This will load RWKV-7 "Goose" x070 and inference in GPT-mode (slower than RNN-mode for autoregressive generation) -''' +""" args = types.SimpleNamespace() @@ -35,35 +37,33 @@ D_GATE_LORA = 128 args.vocab_size = 65536 # DTYPE = torch.bfloat16 -DTYPE = torch.half # better +DTYPE = torch.half # better -args.head_size_a = 64 # don't change +args.head_size_a = 64 # don't change HEAD_SIZE = args.head_size_a -USE_CUDA_KERNEL = True # False => UNOPTIMIZED, VERY SLOW - -MyModule = torch.jit.ScriptModule -MyFunction = torch.jit.script_method -MyStatic = torch.jit.script +USE_CUDA_KERNEL = True # False => UNOPTIMIZED, VERY SLOW ######################################################################################################## # RWKV Tokenizer (slow version) ######################################################################################################## -class RWKV_TOKENIZER(): + +class RWKV_TOKENIZER: table: list[list[list[bytes]]] good: list[set[int]] wlen: list[int] + def __init__(self, file_name): self.idx2token = {} - sorted = [] # must be already sorted + sorted = [] # must be already sorted lines = open(file_name, "r", encoding="utf-8").readlines() for l in lines: - idx = int(l[:l.index(' ')]) - x = eval(l[l.index(' '):l.rindex(' ')]) + idx = int(l[: l.index(" ")]) + x = eval(l[l.index(" ") : l.rindex(" ")]) x = x.encode("utf-8") if isinstance(x, str) else x assert isinstance(x, bytes) - assert len(x) == int(l[l.rindex(' '):]) + assert len(x) == int(l[l.rindex(" ") :]) sorted += [x] self.idx2token[idx] = x @@ -76,7 +76,7 @@ class RWKV_TOKENIZER(): self.good = [set() for i in range(256)] self.wlen = [0 for i in range(256)] - for i in reversed(range(len(sorted))): # reverse order - match longer tokens first + for i in reversed(range(len(sorted))): # reverse order - match longer tokens first s = sorted[i] if len(s) >= 2: s0 = int(s[0]) @@ -107,25 +107,26 @@ class RWKV_TOKENIZER(): return tokens def decodeBytes(self, tokens): - return b''.join(map(lambda i: self.idx2token[i], tokens)) + return b"".join(map(lambda i: self.idx2token[i], tokens)) def encode(self, src: str): return self.encodeBytes(src.encode("utf-8")) def decode(self, tokens): - return self.decodeBytes(tokens).decode('utf-8') + return self.decodeBytes(tokens).decode("utf-8") def printTokens(self, tokens): for i in tokens: s = self.idx2token[i] try: - s = s.decode('utf-8') + s = s.decode("utf-8") except: pass - print(f'{repr(s)}{i}', end=' ') + print(f"{repr(s)}{i}", end=" ") # print(repr(s), i) print() + tokenizer = RWKV_TOKENIZER("rwkv_vocab_v20230424.txt") ######################################################################################################## @@ -136,8 +137,21 @@ if USE_CUDA_KERNEL: from torch.utils.cpp_extension import load - load(name="wkv7", sources=["cuda/wkv7_op.cpp", f"cuda/wkv7.cu"], is_python_module=False, - verbose=True, extra_cuda_cflags=["-res-usage", "--use_fast_math", "-O3", "-Xptxas -O3", "--extra-device-vectorization", f"-D_N_={HEAD_SIZE}"]) + load( + name="wkv7", + sources=["cuda/wkv7_op.cpp", f"cuda/wkv7.cu"], + is_python_module=False, + verbose=True, + extra_cuda_cflags=[ + "-res-usage", + "--use_fast_math", + "-O3", + "-Xptxas -O3", + "--extra-device-vectorization", + f"-D_N_={HEAD_SIZE}", + ], + ) + class WKV_7(torch.autograd.Function): @staticmethod def forward(ctx, r, w, k, v, a, b): @@ -186,7 +200,7 @@ else: vv = v[:, t, :].view(B, H, N, 1) aa = a[:, t, :].view(B, H, N, 1) bb = b[:, t, :].view(B, H, 1, N) - state = state * w[: , t, :, None, :] + state @ aa @ bb + vv @ kk + state = state * w[:, t, :, None, :] + state @ aa @ bb + vv @ kk out[:, t, :] = (state @ rr).view(B, H, N) # another method using einsum @@ -202,11 +216,13 @@ else: return out.view(B, T, C).to(dtype=DTYPE) + ######################################################################################################## # RWKV TimeMix ######################################################################################################## -class RWKV_Tmix_x070(MyModule): + +class RWKV_Tmix_x070(Module): def __init__(self, args, layer_id): super().__init__() self.args = args @@ -220,40 +236,39 @@ class RWKV_Tmix_x070(MyModule): N = self.head_size C = args.n_embd - self.x_r = nn.Parameter(torch.empty(1,1,C)) - self.x_w = nn.Parameter(torch.empty(1,1,C)) - self.x_k = nn.Parameter(torch.empty(1,1,C)) - self.x_v = nn.Parameter(torch.empty(1,1,C)) - self.x_a = nn.Parameter(torch.empty(1,1,C)) - self.x_g = nn.Parameter(torch.empty(1,1,C)) + self.x_r = nn.Parameter(torch.empty(1, 1, C)) + self.x_w = nn.Parameter(torch.empty(1, 1, C)) + self.x_k = nn.Parameter(torch.empty(1, 1, C)) + self.x_v = nn.Parameter(torch.empty(1, 1, C)) + self.x_a = nn.Parameter(torch.empty(1, 1, C)) + self.x_g = nn.Parameter(torch.empty(1, 1, C)) - self.w0 = nn.Parameter(torch.empty(1,1,C)) + self.w0 = nn.Parameter(torch.empty(1, 1, C)) self.w1 = nn.Parameter(torch.empty(C, D_DECAY_LORA)) self.w2 = nn.Parameter(torch.empty(D_DECAY_LORA, C)) - self.a0 = nn.Parameter(torch.empty(1,1,C)) + self.a0 = nn.Parameter(torch.empty(1, 1, C)) self.a1 = nn.Parameter(torch.empty(C, D_AAA_LORA)) self.a2 = nn.Parameter(torch.empty(D_AAA_LORA, C)) - self.v0 = nn.Parameter(torch.empty(1,1,C)) + self.v0 = nn.Parameter(torch.empty(1, 1, C)) self.v1 = nn.Parameter(torch.empty(C, D_MV_LORA)) self.v2 = nn.Parameter(torch.empty(D_MV_LORA, C)) self.g1 = nn.Parameter(torch.empty(C, D_GATE_LORA)) self.g2 = nn.Parameter(torch.empty(D_GATE_LORA, C)) - self.k_k = nn.Parameter(torch.empty(1,1,C)) - self.k_a = nn.Parameter(torch.empty(1,1,C)) - self.r_k = nn.Parameter(torch.empty(H,N)) + self.k_k = nn.Parameter(torch.empty(1, 1, C)) + self.k_a = nn.Parameter(torch.empty(1, 1, C)) + self.r_k = nn.Parameter(torch.empty(H, N)) self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) self.receptance = nn.Linear(C, C, bias=False) self.key = nn.Linear(C, C, bias=False) self.value = nn.Linear(C, C, bias=False) self.output = nn.Linear(C, C, bias=False) - self.ln_x = nn.GroupNorm(H, C, eps=64e-5) # !!! notice eps value !!! + self.ln_x = nn.GroupNorm(H, C, eps=64e-5) # !!! notice eps value !!! - @MyFunction def forward(self, x, v_first): B, T, C = x.size() H = self.n_head @@ -267,32 +282,36 @@ class RWKV_Tmix_x070(MyModule): xg = x + xx * self.x_g r = self.receptance(xr) - w = -F.softplus(-(self.w0 + torch.tanh(xw @ self.w1) @ self.w2)) - 0.5 # soft-clamp to (-inf, -0.5) + w = -F.softplus(-(self.w0 + torch.tanh(xw @ self.w1) @ self.w2)) - 0.5 # soft-clamp to (-inf, -0.5) k = self.key(xk) v = self.value(xv) if self.layer_id == 0: - v_first = v # store the v of the first layer + v_first = v # store the v of the first layer else: - v = v + (v_first - v) * torch.sigmoid(self.v0 + (xv @ self.v1) @ self.v2) # add value residual - a = torch.sigmoid(self.a0 + (xa @ self.a1) @ self.a2) # a is "in-context learning rate" + v = v + (v_first - v) * torch.sigmoid(self.v0 + (xv @ self.v1) @ self.v2) # add value residual + a = torch.sigmoid(self.a0 + (xa @ self.a1) @ self.a2) # a is "in-context learning rate" g = torch.sigmoid(xg @ self.g1) @ self.g2 kk = k * self.k_k - kk = F.normalize(kk.view(B,T,H,-1), dim=-1, p=2.0).view(B,T,C) - k = k * (1 + (a-1) * self.k_a) + kk = F.normalize(kk.view(B, T, H, -1), dim=-1, p=2.0).view(B, T, C) + k = k * (1 + (a - 1) * self.k_a) - x = RWKV7_OP(r, w, k, v, -kk, kk*a) + x = RWKV7_OP(r, w, k, v, -kk, kk * a) x = self.ln_x(x.view(B * T, C)).view(B, T, C) - - x = x + ((r.view(B,T,H,-1)*k.view(B,T,H,-1)*self.r_k).sum(dim=-1, keepdim=True) * v.view(B,T,H,-1)).view(B,T,C) + + x = x + ( + (r.view(B, T, H, -1) * k.view(B, T, H, -1) * self.r_k).sum(dim=-1, keepdim=True) * v.view(B, T, H, -1) + ).view(B, T, C) x = self.output(x * g) return x, v_first - + + ######################################################################################################## # RWKV ChannelMix ######################################################################################################## -class RWKV_CMix_x070(MyModule): + +class RWKV_CMix_x070(Module): def __init__(self, args, layer_id): super().__init__() self.args = args @@ -305,32 +324,32 @@ class RWKV_CMix_x070(MyModule): self.key = nn.Linear(args.n_embd, args.dim_ffn, bias=False) self.value = nn.Linear(args.dim_ffn, args.n_embd, bias=False) - @MyFunction def forward(self, x): xx = self.time_shift(x) - x - + k = x + xx * self.x_k k = torch.relu(self.key(k)) ** 2 return self.value(k) + ######################################################################################################## # RWKV Block ######################################################################################################## -class Block(MyModule): + +class Block(Module): def __init__(self, args, layer_id): super().__init__() self.args = args self.layer_id = layer_id - self.ln0 = nn.LayerNorm(args.n_embd) # only used in block 0, should be fused with emb + self.ln0 = nn.LayerNorm(args.n_embd) # only used in block 0, should be fused with emb self.ln1 = nn.LayerNorm(args.n_embd) self.ln2 = nn.LayerNorm(args.n_embd) self.att = RWKV_Tmix_x070(args, layer_id) self.ffn = RWKV_CMix_x070(args, layer_id) - - @MyFunction + def forward(self, x, v_first): if self.layer_id == 0: @@ -342,10 +361,12 @@ class Block(MyModule): return x, v_first + ######################################################################################################## # RWKV Model ######################################################################################################## + class RWKV(nn.Module): def __init__(self, args): super().__init__() @@ -371,6 +392,7 @@ class RWKV(nn.Module): return x + ######################################################################################################## # RWKV Inference ######################################################################################################## @@ -380,38 +402,38 @@ model_params = torch.load(MODEL_PATH, map_location="cpu") with torch.no_grad(): model = RWKV(args).to(dtype=DTYPE).cuda() - model.load_state_dict(model_params, strict=False) # we will ignore blocks.0.att.v0/v1/v2 + model.load_state_dict(model_params, strict=False) # we will ignore blocks.0.att.v0/v1/v2 ######################################################################################################## prompt = "中国的首都是在" input = tokenizer.encode(prompt) - print(f'\nInput:\n{input}') + print(f"\nInput:\n{input}") - out = model.forward(torch.tensor(input).reshape(1,-1).cuda()) - print(f'\nOutput:\n{out}') + out = model.forward(torch.tensor(input).reshape(1, -1).cuda()) + print(f"\nOutput:\n{out}") - # logits of the last token => prediction for the next token + # logits of the last token => prediction for the next token out = out[0, -1] - - probs = F.softmax(out.float(), dim=-1) # compute softmax in float (more accurate) - print(f'\n{prompt}') + probs = F.softmax(out.float(), dim=-1) # compute softmax in float (more accurate) - _, indices = torch.topk(probs, 10) # print top-10 possibilities + print(f"\n{prompt}") + + _, indices = torch.topk(probs, 10) # print top-10 possibilities for i in range(len(indices)): token_id = indices[i].item() token = tokenizer.decode([token_id]) token_prob = probs[token_id].item() - print(token, f'[probability {token_prob:.2%}]') + print(token, f"[probability {token_prob:.2%}]") ######################################################################################################## with open(f"misc/lambada_test.jsonl", "r", encoding="utf-8") as f: todo = [json.loads(line) for line in f] - todo = [[doc['text'].rsplit(' ', 1)[0], " " + doc['text'].rsplit(' ', 1)[1]] for doc in todo] + todo = [[doc["text"].rsplit(" ", 1)[0], " " + doc["text"].rsplit(" ", 1)[1]] for doc in todo] - print('\nCheck LAMBADA...') + print("\nCheck LAMBADA...") xsum = 0 xcnt = 0 xacc = 0 @@ -421,9 +443,9 @@ with torch.no_grad(): logits = 0 correct = True - out = model.forward(torch.tensor(src+dst).reshape(1,-1).cuda()) + out = model.forward(torch.tensor(src + dst).reshape(1, -1).cuda()) for i in range(len(dst)): - ooo = out[0,len(src)-1+i].float() + ooo = out[0, len(src) - 1 + i].float() probs = F.softmax(ooo, dim=-1) logits += math.log(probs[dst[i]]) if torch.argmax(probs).item() != dst[i]: @@ -433,4 +455,4 @@ with torch.no_grad(): xsum += logits xacc += 1 if correct else 0 if xcnt % 100 == 0 or xcnt == len(todo): - print(xcnt, 'ppl', round(math.exp(-xsum / xcnt), 2), 'acc', round(xacc/xcnt*100, 2)) + print(xcnt, "ppl", round(math.exp(-xsum / xcnt), 2), "acc", round(xacc / xcnt * 100, 2))