Dev gzf exp (#1654)

zhifu gao

2024-04-24 861147c7308b91068ffa02724fdf74ee623a909e

 funasr/models/sense_voice/rwkv_v6.py

@@ -4,11 +4,11 @@

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):
@@ -27,23 +27,40 @@

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
   
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}"])
    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):
@@ -65,8 +82,9 @@
         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)
            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
   
@@ -80,16 +98,41 @@
         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)
            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)
@@ -123,13 +166,17 @@
         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_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))
            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)
@@ -139,7 +186,9 @@
         
         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))
            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):
@@ -155,7 +204,9 @@
      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))
        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):
@@ -202,6 +253,7 @@
      
      return self.jit_func_2(x, g)


class RWKV_CMix_x060(MyModule):
   def __init__(self, args, layer_id):
      super().__init__()
@@ -245,11 +297,9 @@
      if self.layer_id == 0:
         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)
@@ -260,7 +310,6 @@
      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:
@@ -293,7 +342,6 @@
      if args.get("ln1", True):
         self.ln1 = nn.LayerNorm(args.n_embd)
      self.ln2 = nn.LayerNorm(args.n_embd)
		

      self.att = RWKV_Tmix_x060(args, layer_id)
      
@@ -353,16 +401,16 @@
   def __init__(self, args):
      super().__init__()
      self.args = args
      if not hasattr(args, 'dim_att'):
        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"]:
        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'):
        if not hasattr(args, "tiny_att_layer"):
         args.tiny_att_layer = -1
      if not hasattr(args, 'tiny_att_dim'):
        if not hasattr(args, "tiny_att_dim"):
         args.tiny_att_dim = -1
      assert args.n_embd % 32 == 0
      assert args.dim_att % 32 == 0
@@ -375,10 +423,8 @@
      self.ln_out = nn.LayerNorm(args.n_embd)
      self.head = nn.Linear(args.n_embd, args.vocab_size, bias=False)


      if args.dropout > 0:
         self.drop0 = nn.Dropout(p=args.dropout)


   def forward(self, idx):
      args = self.args