From fc68b5ffe453235294a561737d8e84bb6c1689a4 Mon Sep 17 00:00:00 2001
From: zhifu gao <zhifu.gzf@alibaba-inc.com>
Date: 星期四, 25 四月 2024 21:43:47 +0800
Subject: [PATCH] Dev gzf exp (#1661)
---
funasr/models/sense_voice/decoder.py | 37 +++++++++++++++++++++++++++++++------
1 files changed, 31 insertions(+), 6 deletions(-)
diff --git a/funasr/models/sense_voice/decoder.py b/funasr/models/sense_voice/decoder.py
index 5d34ff3..133508f 100644
--- a/funasr/models/sense_voice/decoder.py
+++ b/funasr/models/sense_voice/decoder.py
@@ -156,7 +156,6 @@
return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2), qk.detach()
-from funasr.models.sense_voice.rwkv_v6 import RWKVLayer
from omegaconf import OmegaConf
@@ -168,9 +167,25 @@
rwkv_cfg = kwargs.get("rwkv_cfg", {})
args = OmegaConf.create(rwkv_cfg)
- self.attn = RWKVLayer(args=args, layer_id=layer_id)
- if args.get("datatype", "bf16") == "bf16":
- self.attn.to(torch.bfloat16)
+ if args.get("version", "v4") == "v4":
+ from funasr.models.sense_voice.rwkv_v4 import RWKVLayer
+ from funasr.models.sense_voice.rwkv_v4 import RWKV_TimeMix as RWKV_Tmix
+ elif args.get("version", "v5") == "v5":
+ from funasr.models.sense_voice.rwkv_v5 import RWKVLayer
+ from funasr.models.sense_voice.rwkv_v5 import RWKV_Tmix_x052 as RWKV_Tmix
+ else:
+ from funasr.models.sense_voice.rwkv_v6 import RWKVLayer
+ from funasr.models.sense_voice.rwkv_v6 import RWKV_Tmix_x060 as RWKV_Tmix
+ # self.att = RWKVLayer(args=args, layer_id=layer_id)
+ self.att = RWKV_Tmix(args, layer_id=layer_id)
+
+ 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)
self.ln0 = None
if layer_id == 0 and not args.get("ln0", True):
@@ -180,6 +195,7 @@
layer_id = 0
scale = ((1 + layer_id) / args.get("n_layer")) ** 0.7
nn.init.constant_(self.ln0.weight, scale)
+
self.layer_id = layer_id
self.args = args
@@ -191,6 +207,11 @@
print("init_rwkv")
scale = ((1 + layer_id) / args.get("n_layer")) ** 0.7
nn.init.constant_(self.ln1.weight, scale)
+
+ if args.get("datatype", "bf16") == "bf16":
+ self.att.to(torch.bfloat16)
+ # if self.ln1 is not None:
+ # self.ln1.to(torch.bfloat16)
self.cross_attn = MultiHeadAttention(n_state, n_head) if cross_attention else None
self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None
@@ -213,10 +234,14 @@
if self.layer_id == 0 and self.ln0 is not None:
x = self.ln0(x)
+ if self.args.get("datatype", "bf16") == "bf16":
+ x = x.bfloat16()
if self.ln1 is None:
- x = x + self.attn(x, mask=mask, kv_cache=kv_cache, is_pad_mask=is_pad_mask)[0]
+ x = x + self.att(x, mask=mask, kv_cache=kv_cache, is_pad_mask=is_pad_mask)[0]
else:
- x = x + self.attn(self.ln1(x), mask=mask, kv_cache=kv_cache, is_pad_mask=is_pad_mask)[0]
+ x = x + self.att(self.ln1(x), mask=mask, kv_cache=kv_cache, is_pad_mask=is_pad_mask)[0]
+ if self.args.get("datatype", "bf16") == "bf16":
+ x = x.to(torch.float32)
if self.cross_attn:
x = (
--
Gitblit v1.9.1