From 1cdb3cc28d4d89a576cc06e5cd8eb80da1f3a3aa Mon Sep 17 00:00:00 2001
From: zhifu gao <zhifu.gzf@alibaba-inc.com>
Date: 星期五, 26 四月 2024 11:27:39 +0800
Subject: [PATCH] Dev gzf exp (#1665)
---
funasr/models/sense_voice/rwkv_v6.py | 791 ++++++++++++++++++++++++++++++--------------------------
1 files changed, 422 insertions(+), 369 deletions(-)
diff --git a/funasr/models/sense_voice/rwkv_v6.py b/funasr/models/sense_voice/rwkv_v6.py
index 6eb53fc..b91d47a 100644
--- a/funasr/models/sense_voice/rwkv_v6.py
+++ b/funasr/models/sense_voice/rwkv_v6.py
@@ -4,22 +4,22 @@
import os, math, gc, importlib
import torch
+
# torch._C._jit_set_profiling_executor(True)
# torch._C._jit_set_profiling_mode(True)
import torch.nn as nn
from torch.nn import functional as F
-
def __nop(ob):
- return ob
+ return ob
MyModule = nn.Module
MyFunction = __nop
if "RWKV_JIT_ON" in os.environ and os.environ["RWKV_JIT_ON"] == "1":
- MyModule = torch.jit.ScriptModule
- MyFunction = torch.jit.script_method
+ MyModule = torch.jit.ScriptModule
+ MyFunction = torch.jit.script_method
########################################################################################################
# CUDA Kernel
@@ -27,399 +27,452 @@
wkv6_cuda = None
-def load_rwkv_kernel(HEAD_SIZE: int=64, RWKV_CTXLEN: int=512,):
- from torch.utils.cpp_extension import load
- global wkv6_cuda
-
-
- if wkv6_cuda is not None:
- return
-
- absolute_file_path = os.path.abspath(__file__)
- cur_dir = os.path.dirname(absolute_file_path)
- wkv6_cuda = load(name="wkv6", sources=[f"{cur_dir}/cuda/wkv6_op.cpp", f"{cur_dir}/cuda/wkv6_cuda.cu"],
- verbose=True, extra_cuda_cflags=["-res-usage", "--use_fast_math", "-O3", "-Xptxas -O3",
- "--extra-device-vectorization", f"-D_N_={HEAD_SIZE}",
- f"-D_T_={RWKV_CTXLEN}"])
+
+def load_rwkv_kernel(
+ HEAD_SIZE: int = 64,
+ RWKV_CTXLEN: int = 512,
+):
+ from torch.utils.cpp_extension import load
+
+ global wkv6_cuda
+
+ if wkv6_cuda is not None:
+ return
+
+ absolute_file_path = os.path.abspath(__file__)
+ cur_dir = os.path.dirname(absolute_file_path)
+ wkv6_cuda = load(
+ name="wkv6",
+ sources=[f"{cur_dir}/cuda/wkv6_op.cpp", f"{cur_dir}/cuda/wkv6_cuda.cu"],
+ verbose=True,
+ extra_cuda_cflags=[
+ "-res-usage",
+ "--use_fast_math",
+ "-O3",
+ "-Xptxas -O3",
+ "--extra-device-vectorization",
+ f"-D_N_={HEAD_SIZE}",
+ f"-D_T_={RWKV_CTXLEN}",
+ ],
+ )
+
# dtype = torch.float
dtype = torch.bfloat16
+
+
class WKV_6(torch.autograd.Function):
- @staticmethod
- def forward(ctx, B, T, C, H, r, k, v, w, u):
- with torch.no_grad():
- # assert r.dtype == torch.bfloat16
- # assert k.dtype == torch.bfloat16
- # assert v.dtype == torch.bfloat16
- # assert w.dtype == torch.bfloat16
- # assert u.dtype == torch.bfloat16
- # assert HEAD_SIZE == C // H
- ctx.B = B
- ctx.T = T
- ctx.C = C
- ctx.H = H
- assert r.is_contiguous()
- assert k.is_contiguous()
- assert v.is_contiguous()
- assert w.is_contiguous()
- assert u.is_contiguous()
- ew = (-torch.exp(w.float())).contiguous()
- ctx.save_for_backward(r, k, v, ew, u)
- y = torch.empty((B, T, C), device=r.device, dtype=dtype,
- memory_format=torch.contiguous_format) # .uniform_(-100, 100)
- wkv6_cuda.forward(B, T, C, H, r, k, v, ew, u, y)
- return y
-
- @staticmethod
- def backward(ctx, gy):
- with torch.no_grad():
- # assert gy.dtype == torch.bfloat16
- B = ctx.B
- T = ctx.T
- C = ctx.C
- H = ctx.H
- assert gy.is_contiguous()
- r, k, v, ew, u = ctx.saved_tensors
- gr = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=dtype,
- memory_format=torch.contiguous_format) # .uniform_(-100, 100)
- gk = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=dtype,
- memory_format=torch.contiguous_format) # .uniform_(-100, 100)
- gv = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=dtype,
- memory_format=torch.contiguous_format) # .uniform_(-100, 100)
- gw = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=dtype,
- memory_format=torch.contiguous_format) # .uniform_(-100, 100)
- gu = torch.empty((B, C), device=gy.device, requires_grad=False, dtype=dtype,
- memory_format=torch.contiguous_format) # .uniform_(-100, 100)
- wkv6_cuda.backward(B, T, C, H, r, k, v, ew, u, gy, gr, gk, gv, gw, gu)
- gu = torch.sum(gu, 0).view(H, C // H)
- return (None, None, None, None, gr, gk, gv, gw, gu)
+ @staticmethod
+ def forward(ctx, B, T, C, H, r, k, v, w, u):
+ with torch.no_grad():
+ # assert r.dtype == torch.bfloat16
+ # assert k.dtype == torch.bfloat16
+ # assert v.dtype == torch.bfloat16
+ # assert w.dtype == torch.bfloat16
+ # assert u.dtype == torch.bfloat16
+ # assert HEAD_SIZE == C // H
+ ctx.B = B
+ ctx.T = T
+ ctx.C = C
+ ctx.H = H
+ assert r.is_contiguous()
+ assert k.is_contiguous()
+ assert v.is_contiguous()
+ assert w.is_contiguous()
+ assert u.is_contiguous()
+ ew = (-torch.exp(w.float())).contiguous()
+ ctx.save_for_backward(r, k, v, ew, u)
+ y = torch.empty(
+ (B, T, C), device=r.device, dtype=dtype, memory_format=torch.contiguous_format
+ ) # .uniform_(-100, 100)
+ wkv6_cuda.forward(B, T, C, H, r, k, v, ew, u, y)
+ return y
+
+ @staticmethod
+ def backward(ctx, gy):
+ with torch.no_grad():
+ # assert gy.dtype == torch.bfloat16
+ B = ctx.B
+ T = ctx.T
+ C = ctx.C
+ H = ctx.H
+ assert gy.is_contiguous()
+ r, k, v, ew, u = ctx.saved_tensors
+ gr = torch.empty(
+ (B, T, C),
+ device=gy.device,
+ requires_grad=False,
+ dtype=dtype,
+ memory_format=torch.contiguous_format,
+ ) # .uniform_(-100, 100)
+ gk = torch.empty(
+ (B, T, C),
+ device=gy.device,
+ requires_grad=False,
+ dtype=dtype,
+ memory_format=torch.contiguous_format,
+ ) # .uniform_(-100, 100)
+ gv = torch.empty(
+ (B, T, C),
+ device=gy.device,
+ requires_grad=False,
+ dtype=dtype,
+ memory_format=torch.contiguous_format,
+ ) # .uniform_(-100, 100)
+ gw = torch.empty(
+ (B, T, C),
+ device=gy.device,
+ requires_grad=False,
+ dtype=dtype,
+ memory_format=torch.contiguous_format,
+ ) # .uniform_(-100, 100)
+ gu = torch.empty(
+ (B, C),
+ device=gy.device,
+ requires_grad=False,
+ dtype=dtype,
+ memory_format=torch.contiguous_format,
+ ) # .uniform_(-100, 100)
+ wkv6_cuda.backward(B, T, C, H, r, k, v, ew, u, gy, gr, gk, gv, gw, gu)
+ gu = torch.sum(gu, 0).view(H, C // H)
+ return (None, None, None, None, gr, gk, gv, gw, gu)
def RUN_CUDA_RWKV6(B, T, C, H, r, k, v, w, u):
- return WKV_6.apply(B, T, C, H, r, k, v, w, u)
+ return WKV_6.apply(B, T, C, H, r, k, v, w, u)
class RWKV_Tmix_x060(MyModule):
- def __init__(self, args, layer_id):
- super().__init__()
- self.args = args
-
- load_rwkv_kernel(args.head_size_a, args.ctx_len)
-
- self.layer_id = layer_id
-
- self.head_size = args.head_size_a
- self.n_head = args.dim_att // self.head_size
- assert args.dim_att % self.n_head == 0
-
- with torch.no_grad():
- ratio_0_to_1 = layer_id / (args.n_layer - 1) # 0 to 1
- ratio_1_to_almost0 = 1.0 - (layer_id / args.n_layer) # 1 to ~0
- ddd = torch.ones(1, 1, args.n_embd)
- for i in range(args.n_embd):
- ddd[0, 0, i] = i / args.n_embd
-
- # fancy time_mix
- self.time_maa_x = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
- self.time_maa_w = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
- self.time_maa_k = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
- self.time_maa_v = nn.Parameter(1.0 - (torch.pow(ddd, ratio_1_to_almost0) + 0.3 * ratio_0_to_1))
- self.time_maa_r = nn.Parameter(1.0 - torch.pow(ddd, 0.5 * ratio_1_to_almost0))
- self.time_maa_g = nn.Parameter(1.0 - torch.pow(ddd, 0.5 * ratio_1_to_almost0))
-
- D_MIX_LORA = 32 # generate TIME_MIX for w,k,v,r,g
- self.time_maa_w1 = nn.Parameter(torch.zeros(args.n_embd, D_MIX_LORA * 5))
- self.time_maa_w2 = nn.Parameter(torch.zeros(5, D_MIX_LORA, args.n_embd).uniform_(-0.01, 0.01))
-
- # fancy time_decay
- decay_speed = torch.ones(args.dim_att)
- for n in range(args.dim_att):
- decay_speed[n] = -6 + 5 * (n / (args.dim_att - 1)) ** (0.7 + 1.3 * ratio_0_to_1)
- self.time_decay = nn.Parameter(decay_speed.reshape(1, 1, args.dim_att))
-
- D_DECAY_LORA = 64
- self.time_decay_w1 = nn.Parameter(torch.zeros(args.n_embd, D_DECAY_LORA))
- self.time_decay_w2 = nn.Parameter(torch.zeros(D_DECAY_LORA, args.dim_att).uniform_(-0.01, 0.01))
-
- tmp = torch.zeros(args.dim_att)
- for n in range(args.dim_att):
- zigzag = ((n + 1) % 3 - 1) * 0.1
- tmp[n] = ratio_0_to_1 * (1 - (n / (args.dim_att - 1))) + zigzag
-
- self.time_faaaa = nn.Parameter(tmp.reshape(self.n_head, self.head_size))
-
- self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
- self.receptance = nn.Linear(args.n_embd, args.dim_att, bias=False)
- self.key = nn.Linear(args.n_embd, args.dim_att, bias=False)
-
- self.value = nn.Linear(args.n_embd, args.dim_att, bias=False)
- self.output = nn.Linear(args.dim_att, args.n_embd, bias=False)
- self.gate = nn.Linear(args.n_embd, args.dim_att, bias=False)
- self.ln_x = nn.GroupNorm(self.n_head, args.dim_att, eps=(1e-5) * (args.head_size_divisor ** 2))
-
- @MyFunction
- def jit_func(self, x):
- B, T, C = x.size()
-
- xx = self.time_shift(x) - x
-
- xxx = x + xx * self.time_maa_x
- xxx = torch.tanh(xxx @ self.time_maa_w1).view(B * T, 5, -1).transpose(0, 1)
- xxx = torch.bmm(xxx, self.time_maa_w2).view(5, B, T, -1)
- mw, mk, mv, mr, mg = xxx.unbind(dim=0)
-
- xw = x + xx * (self.time_maa_w + mw)
- xk = x + xx * (self.time_maa_k + mk)
- xv = x + xx * (self.time_maa_v + mv)
- xr = x + xx * (self.time_maa_r + mr)
- xg = x + xx * (self.time_maa_g + mg)
-
- r = self.receptance(xr)
- k = self.key(xk)
- v = self.value(xv)
- g = F.silu(self.gate(xg))
-
- ww = torch.tanh(xw @ self.time_decay_w1) @ self.time_decay_w2
- w = self.time_decay + ww
-
- return r, k, v, g, w
-
- @MyFunction
- def jit_func_2(self, x, g):
- B, T, C = x.size()
- x = x.view(B * T, C)
-
- x = self.ln_x(x).view(B, T, C)
- x = self.output(x * g)
- return x
-
- def forward(self, x):
- B, T, C = x.size()
- H = self.n_head
-
- r, k, v, g, w = self.jit_func(x)
- x = RUN_CUDA_RWKV6(B, T, C, H, r, k, v, w, u=self.time_faaaa)
-
- return self.jit_func_2(x, g)
+ def __init__(self, args, layer_id):
+ super().__init__()
+ self.args = args
+
+ load_rwkv_kernel(args.head_size_a, args.ctx_len)
+
+ self.layer_id = layer_id
+
+ self.head_size = args.head_size_a
+ self.n_head = args.dim_att // self.head_size
+ assert args.dim_att % self.n_head == 0
+
+ with torch.no_grad():
+ ratio_0_to_1 = layer_id / (args.n_layer - 1) # 0 to 1
+ ratio_1_to_almost0 = 1.0 - (layer_id / args.n_layer) # 1 to ~0
+ ddd = torch.ones(1, 1, args.n_embd)
+ for i in range(args.n_embd):
+ ddd[0, 0, i] = i / args.n_embd
+
+ # fancy time_mix
+ self.time_maa_x = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
+ self.time_maa_w = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
+ self.time_maa_k = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
+ self.time_maa_v = nn.Parameter(
+ 1.0 - (torch.pow(ddd, ratio_1_to_almost0) + 0.3 * ratio_0_to_1)
+ )
+ self.time_maa_r = nn.Parameter(1.0 - torch.pow(ddd, 0.5 * ratio_1_to_almost0))
+ self.time_maa_g = nn.Parameter(1.0 - torch.pow(ddd, 0.5 * ratio_1_to_almost0))
+
+ D_MIX_LORA = 32 # generate TIME_MIX for w,k,v,r,g
+ self.time_maa_w1 = nn.Parameter(torch.zeros(args.n_embd, D_MIX_LORA * 5))
+ self.time_maa_w2 = nn.Parameter(
+ torch.zeros(5, D_MIX_LORA, args.n_embd).uniform_(-0.01, 0.01)
+ )
+
+ # fancy time_decay
+ decay_speed = torch.ones(args.dim_att)
+ for n in range(args.dim_att):
+ decay_speed[n] = -6 + 5 * (n / (args.dim_att - 1)) ** (0.7 + 1.3 * ratio_0_to_1)
+ self.time_decay = nn.Parameter(decay_speed.reshape(1, 1, args.dim_att))
+
+ D_DECAY_LORA = 64
+ self.time_decay_w1 = nn.Parameter(torch.zeros(args.n_embd, D_DECAY_LORA))
+ self.time_decay_w2 = nn.Parameter(
+ torch.zeros(D_DECAY_LORA, args.dim_att).uniform_(-0.01, 0.01)
+ )
+
+ tmp = torch.zeros(args.dim_att)
+ for n in range(args.dim_att):
+ zigzag = ((n + 1) % 3 - 1) * 0.1
+ tmp[n] = ratio_0_to_1 * (1 - (n / (args.dim_att - 1))) + zigzag
+
+ self.time_faaaa = nn.Parameter(tmp.reshape(self.n_head, self.head_size))
+
+ self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
+ self.receptance = nn.Linear(args.n_embd, args.dim_att, bias=False)
+ self.key = nn.Linear(args.n_embd, args.dim_att, bias=False)
+
+ self.value = nn.Linear(args.n_embd, args.dim_att, bias=False)
+ self.output = nn.Linear(args.dim_att, args.n_embd, bias=False)
+ self.gate = nn.Linear(args.n_embd, args.dim_att, bias=False)
+ self.ln_x = nn.GroupNorm(
+ self.n_head, args.dim_att, eps=(1e-5) * (args.head_size_divisor**2)
+ )
+
+ @MyFunction
+ def jit_func(self, x):
+ B, T, C = x.size()
+
+ xx = self.time_shift(x) - x
+
+ xxx = x + xx * self.time_maa_x
+ xxx = torch.tanh(xxx @ self.time_maa_w1).view(B * T, 5, -1).transpose(0, 1)
+ xxx = torch.bmm(xxx, self.time_maa_w2).view(5, B, T, -1)
+ mw, mk, mv, mr, mg = xxx.unbind(dim=0)
+
+ xw = x + xx * (self.time_maa_w + mw)
+ xk = x + xx * (self.time_maa_k + mk)
+ xv = x + xx * (self.time_maa_v + mv)
+ xr = x + xx * (self.time_maa_r + mr)
+ xg = x + xx * (self.time_maa_g + mg)
+
+ r = self.receptance(xr)
+ k = self.key(xk)
+ v = self.value(xv)
+ g = F.silu(self.gate(xg))
+
+ ww = torch.tanh(xw @ self.time_decay_w1) @ self.time_decay_w2
+ w = self.time_decay + ww
+
+ return r, k, v, g, w
+
+ @MyFunction
+ def jit_func_2(self, x, g):
+ B, T, C = x.size()
+ x = x.view(B * T, C)
+
+ x = self.ln_x(x).view(B, T, C)
+ x = self.output(x * g)
+ return x
+
+ def forward(self, x, **kwargs):
+ B, T, C = x.size()
+ H = self.n_head
+
+ r, k, v, g, w = self.jit_func(x)
+ x = RUN_CUDA_RWKV6(B, T, C, H, r, k, v, w, u=self.time_faaaa)
+
+ return self.jit_func_2(x, g)
+
class RWKV_CMix_x060(MyModule):
- def __init__(self, args, layer_id):
- super().__init__()
- self.args = args
- self.layer_id = layer_id
- self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
-
- with torch.no_grad(): # fancy init of time_mix
- ratio_1_to_almost0 = 1.0 - (layer_id / args.n_layer) # 1 to ~0
- ddd = torch.ones(1, 1, args.n_embd)
- for i in range(args.n_embd):
- ddd[0, 0, i] = i / args.n_embd
- self.time_maa_k = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
- self.time_maa_r = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
-
- self.key = nn.Linear(args.n_embd, args.dim_ffn, bias=False)
- self.receptance = nn.Linear(args.n_embd, args.n_embd, 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
- xk = x + xx * self.time_maa_k
- xr = x + xx * self.time_maa_r
-
- k = self.key(xk)
- k = torch.relu(k) ** 2
- kv = self.value(k)
- return torch.sigmoid(self.receptance(xr)) * kv
+ def __init__(self, args, layer_id):
+ super().__init__()
+ self.args = args
+ self.layer_id = layer_id
+ self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
+
+ with torch.no_grad(): # fancy init of time_mix
+ ratio_1_to_almost0 = 1.0 - (layer_id / args.n_layer) # 1 to ~0
+ ddd = torch.ones(1, 1, args.n_embd)
+ for i in range(args.n_embd):
+ ddd[0, 0, i] = i / args.n_embd
+ self.time_maa_k = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
+ self.time_maa_r = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
+
+ self.key = nn.Linear(args.n_embd, args.dim_ffn, bias=False)
+ self.receptance = nn.Linear(args.n_embd, args.n_embd, 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
+ xk = x + xx * self.time_maa_k
+ xr = x + xx * self.time_maa_r
+
+ k = self.key(xk)
+ k = torch.relu(k) ** 2
+ kv = self.value(k)
+ return torch.sigmoid(self.receptance(xr)) * kv
class Block(nn.Module):
- def __init__(self, args, layer_id):
- super().__init__()
- self.args = args
- self.layer_id = layer_id
-
- self.ln1 = nn.LayerNorm(args.n_embd)
- self.ln2 = nn.LayerNorm(args.n_embd)
-
- if self.layer_id == 0:
- self.ln0 = nn.LayerNorm(args.n_embd)
+ def __init__(self, args, layer_id):
+ super().__init__()
+ self.args = args
+ self.layer_id = layer_id
+ self.ln1 = nn.LayerNorm(args.n_embd)
+ self.ln2 = nn.LayerNorm(args.n_embd)
- self.att = RWKV_Tmix_x060(args, layer_id)
-
- self.ffn = RWKV_CMix_x060(args, layer_id)
+ if self.layer_id == 0:
+ self.ln0 = nn.LayerNorm(args.n_embd)
+ self.att = RWKV_Tmix_x060(args, layer_id)
- if args.dropout > 0:
- self.drop0 = nn.Dropout(p=args.dropout)
- self.drop1 = nn.Dropout(p=args.dropout)
-
- def forward(self, x, x_emb=None):
- args = self.args
- B, T, C = x.size()
- if self.layer_id == 0:
- x = self.ln0(x)
+ self.ffn = RWKV_CMix_x060(args, layer_id)
-
- if self.args.dropout == 0:
- if self.layer_id == 0 and args.pre_ffn > 0:
- x = x + self.ffnPre(self.ln1(x))
- else:
- x = x + self.att(self.ln1(x))
- x = x + self.ffn(self.ln2(x))
- else:
- if self.layer_id == 0 and args.pre_ffn > 0:
- x = self.drop0(x + self.ffnPre(self.ln1(x)))
- else:
- x = self.drop0(x + self.att(self.ln1(x)))
- x = self.drop1(x + self.ffn(self.ln2(x)))
-
- return x
+ if args.dropout > 0:
+ self.drop0 = nn.Dropout(p=args.dropout)
+ self.drop1 = nn.Dropout(p=args.dropout)
+
+ def forward(self, x, x_emb=None):
+ args = self.args
+ B, T, C = x.size()
+ if self.layer_id == 0:
+ x = self.ln0(x)
+
+ if self.args.dropout == 0:
+ if self.layer_id == 0 and args.pre_ffn > 0:
+ x = x + self.ffnPre(self.ln1(x))
+ else:
+ x = x + self.att(self.ln1(x))
+ x = x + self.ffn(self.ln2(x))
+ else:
+ if self.layer_id == 0 and args.pre_ffn > 0:
+ x = self.drop0(x + self.ffnPre(self.ln1(x)))
+ else:
+ x = self.drop0(x + self.att(self.ln1(x)))
+ x = self.drop1(x + self.ffn(self.ln2(x)))
+
+ return x
class RWKVLayer(nn.Module):
- def __init__(self, args, layer_id):
- super().__init__()
- self.args = args
- self.layer_id = layer_id
- if args.dim_ffn is None:
- args.dim_ffn = int((args.n_embd * 3.5) // 32 * 32)
- self.ln0 = None
- if self.layer_id == 0 and args.get("ln0", True):
- self.ln0 = nn.LayerNorm(args.n_embd)
-
- self.ln1 = None
- if args.get("ln1", True):
- self.ln1 = nn.LayerNorm(args.n_embd)
- self.ln2 = nn.LayerNorm(args.n_embd)
-
+ def __init__(self, args, layer_id):
+ super().__init__()
+ self.args = args
+ self.layer_id = layer_id
+ if args.dim_ffn is None:
+ args.dim_ffn = int((args.n_embd * 3.5) // 32 * 32)
+ self.ln0 = None
+ if self.layer_id == 0 and args.get("ln0", True):
+ self.ln0 = nn.LayerNorm(args.n_embd)
- self.att = RWKV_Tmix_x060(args, layer_id)
-
- self.ffn = RWKV_CMix_x060(args, layer_id)
-
- if args.dropout > 0:
- self.drop0 = nn.Dropout(p=args.dropout)
- self.drop1 = nn.Dropout(p=args.dropout)
-
- # init
- if args.get("init_rwkv", True):
- print("init_rwkv")
- nn.init.orthogonal_(self.att.receptance.weight, gain=1)
- nn.init.orthogonal_(self.att.key.weight, gain=0.1)
- nn.init.orthogonal_(self.att.value.weight, gain=1)
- nn.init.orthogonal_(self.att.gate.weight, gain=0.1)
- nn.init.zeros_(self.att.output.weight)
-
- nn.init.orthogonal_(self.ffn.key.weight, gain=1)
- nn.init.zeros_(self.ffn.value.weight)
- nn.init.zeros_(self.ffn.receptance.weight)
- scale = ((1 + layer_id) / args.get("n_layer")) ** 0.7
- nn.init.constant_(self.ln2.weight, scale)
- if self.ln0 is not None:
- nn.init.constant_(self.ln0.weight, scale)
- if self.ln1 is not None:
- nn.init.constant_(self.ln1.weight, scale)
-
- def forward(self, x, x_emb=None, mask=None, **kwargs):
-
- args = self.args
- if args.get("datatype", "bf16") == "bf16":
- x = x.bfloat16()
- B, T, C = x.size()
- if self.layer_id == 0 and self.ln0 is not None:
- x = self.ln0(x)
-
- if self.args.dropout == 0:
- if self.ln1 is None:
- x = x + self.att(x)
- else:
- x = x + self.att(self.ln1(x))
- x = x + self.ffn(self.ln2(x))
- else:
- if self.ln1 is None:
- x = self.drop0(x + self.att(x))
- else:
- x = self.drop0(x + self.att(self.ln1(x)))
- x = self.drop1(x + self.ffn(self.ln2(x)))
-
- if args.get("datatype", "bf16") == "bf16":
- x = x.to(torch.float32)
- return x
+ self.ln1 = None
+ if args.get("ln1", True):
+ self.ln1 = nn.LayerNorm(args.n_embd)
+
+ self.att = RWKV_Tmix_x060(args, layer_id)
+
+ self.ln2 = None
+ self.ffn = None
+ if args.get("use_rwkv_ffn", True):
+ self.ln2 = nn.LayerNorm(args.n_embd)
+ self.ffn = RWKV_CMix_x060(args, layer_id)
+
+ if args.dropout > 0:
+ self.drop0 = nn.Dropout(p=args.dropout)
+ self.drop1 = nn.Dropout(p=args.dropout)
+
+ # init
+ if args.get("init_rwkv", True):
+ print("init_rwkv")
+ nn.init.orthogonal_(self.att.receptance.weight, gain=1)
+ nn.init.orthogonal_(self.att.key.weight, gain=0.1)
+ nn.init.orthogonal_(self.att.value.weight, gain=1)
+ nn.init.orthogonal_(self.att.gate.weight, gain=0.1)
+ nn.init.zeros_(self.att.output.weight)
+
+ nn.init.orthogonal_(self.ffn.key.weight, gain=1)
+ nn.init.zeros_(self.ffn.value.weight)
+ nn.init.zeros_(self.ffn.receptance.weight)
+ scale = ((1 + layer_id) / args.get("n_layer")) ** 0.7
+
+ if self.ln0 is not None:
+ nn.init.constant_(self.ln0.weight, scale)
+ if self.ln1 is not None:
+ nn.init.constant_(self.ln1.weight, scale)
+ if self.ln2 is not None:
+ nn.init.constant_(self.ln2.weight, scale)
+
+ def forward(self, x, x_emb=None, mask=None, **kwargs):
+
+ args = self.args
+ if args.get("datatype", "bf16") == "bf16":
+ x = x.bfloat16()
+ B, T, C = x.size()
+ if self.layer_id == 0 and self.ln0 is not None:
+ x = self.ln0(x)
+
+ if self.args.dropout == 0:
+ if self.ln1 is None:
+ x = x + self.att(x)
+ else:
+ x = x + self.att(self.ln1(x))
+ if self.ffn is not None:
+ x = x + self.ffn(self.ln2(x))
+ else:
+ if self.ln1 is None:
+ x = self.drop0(x + self.att(x))
+ else:
+ x = self.drop0(x + self.att(self.ln1(x)))
+ if self.ffn is not None:
+ x = self.drop1(x + self.ffn(self.ln2(x)))
+
+ if args.get("datatype", "bf16") == "bf16":
+ x = x.to(torch.float32)
+ return x
class RWKV(nn.Module):
- def __init__(self, args):
- super().__init__()
- self.args = args
- if not hasattr(args, 'dim_att'):
- args.dim_att = args.n_embd
- if not hasattr(args, 'dim_ffn'):
- if '-f4' in os.environ["RWKV_MY_TESTING"]:
- args.dim_ffn = int((args.n_embd * 4) // 32 * 32)
- else:
- args.dim_ffn = int((args.n_embd * 3.5) // 32 * 32) # default = 3.5x emb size
- if not hasattr(args, 'tiny_att_layer'):
- args.tiny_att_layer = -1
- if not hasattr(args, 'tiny_att_dim'):
- args.tiny_att_dim = -1
- assert args.n_embd % 32 == 0
- assert args.dim_att % 32 == 0
- assert args.dim_ffn % 32 == 0
-
- self.emb = nn.Embedding(args.vocab_size, args.n_embd)
-
- self.blocks = nn.ModuleList([Block(args, i) for i in range(args.n_layer)])
-
- self.ln_out = nn.LayerNorm(args.n_embd)
- self.head = nn.Linear(args.n_embd, args.vocab_size, bias=False)
+ def __init__(self, args):
+ super().__init__()
+ self.args = args
+ if not hasattr(args, "dim_att"):
+ args.dim_att = args.n_embd
+ if not hasattr(args, "dim_ffn"):
+ if "-f4" in os.environ["RWKV_MY_TESTING"]:
+ args.dim_ffn = int((args.n_embd * 4) // 32 * 32)
+ else:
+ args.dim_ffn = int((args.n_embd * 3.5) // 32 * 32) # default = 3.5x emb size
+ if not hasattr(args, "tiny_att_layer"):
+ args.tiny_att_layer = -1
+ if not hasattr(args, "tiny_att_dim"):
+ args.tiny_att_dim = -1
+ assert args.n_embd % 32 == 0
+ assert args.dim_att % 32 == 0
+ assert args.dim_ffn % 32 == 0
+ self.emb = nn.Embedding(args.vocab_size, args.n_embd)
- if args.dropout > 0:
- self.drop0 = nn.Dropout(p=args.dropout)
+ self.blocks = nn.ModuleList([Block(args, i) for i in range(args.n_layer)])
+ self.ln_out = nn.LayerNorm(args.n_embd)
+ self.head = nn.Linear(args.n_embd, args.vocab_size, bias=False)
- def forward(self, idx):
- args = self.args
- B, T = idx.size()
- assert T <= args.ctx_len, "Cannot forward, model ctx_len is exhausted."
-
- x = self.emb(idx)
- x_emb = x
-
- if args.dropout > 0:
- x = self.drop0(x)
- if args.tiny_att_dim > 0:
- for block in self.blocks:
- if args.grad_cp == 1:
- x = deepspeed.checkpointing.checkpoint(block, x, x_emb)
- else:
- x = block(x, x_emb)
- else:
- for block in self.blocks:
- if args.grad_cp == 1:
- x = deepspeed.checkpointing.checkpoint(block, x)
- else:
- x = block(x)
-
- x = self.ln_out(x)
-
- if args.head_qk > 0:
- q = self.head_q(x)[:, :T, :]
- k = self.head_k(x)[:, :T, :]
- c = (q @ k.transpose(-2, -1)) * (1.0 / args.head_qk)
- c = c.masked_fill(self.copy_mask[:T, :T] == 0, 0)
-
- if "32" in os.environ["RWKV_FLOAT_MODE"]:
- c = c @ F.one_hot(idx, num_classes=args.vocab_size)
- elif os.environ["RWKV_FLOAT_MODE"] == "fp16":
- c = c @ F.one_hot(idx, num_classes=args.vocab_size).half()
- elif os.environ["RWKV_FLOAT_MODE"] == "bf16":
- c = c @ F.one_hot(idx, num_classes=args.vocab_size).bfloat16()
-
- x = self.head(x) + c
- else:
- x = self.head(x)
-
- return x
+ if args.dropout > 0:
+ self.drop0 = nn.Dropout(p=args.dropout)
+
+ def forward(self, idx):
+ args = self.args
+ B, T = idx.size()
+ assert T <= args.ctx_len, "Cannot forward, model ctx_len is exhausted."
+
+ x = self.emb(idx)
+ x_emb = x
+
+ if args.dropout > 0:
+ x = self.drop0(x)
+ if args.tiny_att_dim > 0:
+ for block in self.blocks:
+ if args.grad_cp == 1:
+ x = deepspeed.checkpointing.checkpoint(block, x, x_emb)
+ else:
+ x = block(x, x_emb)
+ else:
+ for block in self.blocks:
+ if args.grad_cp == 1:
+ x = deepspeed.checkpointing.checkpoint(block, x)
+ else:
+ x = block(x)
+
+ x = self.ln_out(x)
+
+ if args.head_qk > 0:
+ q = self.head_q(x)[:, :T, :]
+ k = self.head_k(x)[:, :T, :]
+ c = (q @ k.transpose(-2, -1)) * (1.0 / args.head_qk)
+ c = c.masked_fill(self.copy_mask[:T, :T] == 0, 0)
+
+ if "32" in os.environ["RWKV_FLOAT_MODE"]:
+ c = c @ F.one_hot(idx, num_classes=args.vocab_size)
+ elif os.environ["RWKV_FLOAT_MODE"] == "fp16":
+ c = c @ F.one_hot(idx, num_classes=args.vocab_size).half()
+ elif os.environ["RWKV_FLOAT_MODE"] == "bf16":
+ c = c @ F.one_hot(idx, num_classes=args.vocab_size).bfloat16()
+
+ x = self.head(x) + c
+ else:
+ x = self.head(x)
+
+ return x
--
Gitblit v1.9.1