KaleiNeely
commited on
Commit
•
d398e6d
1
Parent(s):
c197a93
Update modeling_rwkv6.py
Browse files- modeling_rwkv6.py +19 -97
modeling_rwkv6.py
CHANGED
@@ -37,6 +37,13 @@ from transformers.utils import (
|
|
37 |
)
|
38 |
|
39 |
from .configuration_rwkv6 import Rwkv6Config
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
40 |
|
41 |
|
42 |
logger = logging.get_logger(__name__)
|
@@ -44,102 +51,15 @@ logger = logging.get_logger(__name__)
|
|
44 |
_CHECKPOINT_FOR_DOC = "RWKV/rwkv-6-world-1b6"
|
45 |
_CONFIG_FOR_DOC = "Rwkv6Config"
|
46 |
|
47 |
-
rwkv6_cuda_kernel = None
|
48 |
-
|
49 |
-
def load_wkv6_cuda_kernel(head_size, ctx_len):
|
50 |
-
from torch.utils.cpp_extension import load as load_kernel
|
51 |
-
|
52 |
-
global rwkv6_cuda_kernel
|
53 |
-
|
54 |
-
kernel_folder = Path(__file__).parent.resolve()
|
55 |
-
cuda_kernel_files = [kernel_folder / f for f in ["wkv6_op.cpp", "wkv6_cuda.cu"]]
|
56 |
-
|
57 |
-
# Only load the kernel if it's not been loaded yet or if we changed the context length
|
58 |
-
if rwkv6_cuda_kernel is not None and rwkv6_cuda_kernel.head_size == head_size:
|
59 |
-
return
|
60 |
-
|
61 |
-
logger.info(f"Loading CUDA kernel for RWKV at head size of {head_size}.")
|
62 |
-
|
63 |
-
flags = [
|
64 |
-
"-res-usage",
|
65 |
-
# "--maxrregcount 60", # not sure, should we add this? its not in RWKV-LM
|
66 |
-
"--use_fast_math",
|
67 |
-
"-O3",
|
68 |
-
"-Xptxas -O3",
|
69 |
-
"--extra-device-vectorization",
|
70 |
-
f"-D_N_={head_size}",
|
71 |
-
f"-D_T_={ctx_len}"
|
72 |
-
]
|
73 |
-
rwkv6_cuda_kernel = load_kernel(
|
74 |
-
name=f"wkv_{head_size}_{ctx_len}",
|
75 |
-
sources=cuda_kernel_files,
|
76 |
-
verbose=(logging.get_verbosity() == logging.DEBUG),
|
77 |
-
extra_cuda_cflags=flags,
|
78 |
-
)
|
79 |
-
rwkv6_cuda_kernel.head_size = head_size
|
80 |
-
rwkv6_cuda_kernel.ctx_len = ctx_len
|
81 |
-
|
82 |
-
|
83 |
-
class Rwkv6LinearAttention(torch.autograd.Function):
|
84 |
-
@staticmethod
|
85 |
-
def forward(ctx, receptance, key, value, time_decay, time_first, state):
|
86 |
-
with torch.no_grad():
|
87 |
-
assert receptance.dtype == torch.bfloat16
|
88 |
-
assert key.dtype == torch.bfloat16
|
89 |
-
assert value.dtype == torch.bfloat16
|
90 |
-
assert time_decay.dtype == torch.bfloat16
|
91 |
-
assert time_first.dtype == torch.bfloat16
|
92 |
-
assert state.dtype == torch.float32
|
93 |
-
#assert HEAD_SIZE == C // H
|
94 |
-
Batch, SequenceLength, HiddenSize = key.shape
|
95 |
-
NumHeads, HeadSize = time_decay.shape
|
96 |
-
ctx.Batch = Batch
|
97 |
-
ctx.SequenceLength = SequenceLength
|
98 |
-
ctx.HiddenSize = HiddenSize
|
99 |
-
ctx.NumHeads = NumHeads
|
100 |
-
assert receptance.is_contiguous()
|
101 |
-
assert key.is_contiguous()
|
102 |
-
assert value.is_contiguous()
|
103 |
-
assert time_decay.is_contiguous()
|
104 |
-
assert time_first.is_contiguous()
|
105 |
-
e_time_decay = (-torch.exp(time_decay.float())).contiguous()
|
106 |
-
ctx.save_for_backward(receptance, key, value, e_time_decay, time_first)
|
107 |
-
out = torch.empty((Batch, SequenceLength, HiddenSize), device=receptance.device, dtype=torch.bfloat16, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
|
108 |
-
# FIXME - current kernel does not handle nor update state
|
109 |
-
rwkv6_cuda_kernel.forward(Batch, SequenceLength, HiddenSize, NumHeads, receptance, key, value, e_time_decay, time_first, out)
|
110 |
-
return out, state
|
111 |
-
|
112 |
-
@staticmethod
|
113 |
-
def backward(ctx, g_out, g_state):
|
114 |
-
with torch.no_grad():
|
115 |
-
assert g_out.dtype == torch.bfloat16
|
116 |
-
Batch = ctx.Batch
|
117 |
-
SequenceLength = ctx.SequenceLength
|
118 |
-
HiddenSize = ctx.HiddenSize
|
119 |
-
NumHeads = ctx.NumHeads
|
120 |
-
HeadSize = HiddenSize // NumHeads
|
121 |
-
assert g_out.is_contiguous()
|
122 |
-
receptance, key, value, e_time_decay, time_first = ctx.saved_tensors
|
123 |
-
g_receptance = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=torch.bfloat16, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
|
124 |
-
g_key = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=torch.bfloat16, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
|
125 |
-
g_value = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=torch.bfloat16, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
|
126 |
-
g_time_decay = torch.empty((B, T, C), device=gy.device, requires_grad=False, dtype=torch.bfloat16, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
|
127 |
-
g_time_first = torch.empty((B, C), device=gy.device, requires_grad=False, dtype=torch.bfloat16, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
|
128 |
-
#gs = torch.empty((B, C//H, H, H), device=gy.device, requires_grad=False, dtype=torch.float, memory_format=torch.contiguous_format)#.uniform_(-100, 100)
|
129 |
-
rwkv6_cuda_kernel.backward(B, T, C, H, receptance, key, value, e_time_decay, time_first, g_out, g_receptance, g_key, g_value, g_time_decay, g_time_first)
|
130 |
-
g_time_first = torch.sum(g_time_first, 0).view(NumHeads, HeadSize)
|
131 |
-
return (None, None, None, None, g_receptance, g_key, g_value, g_time_decay, g_time_first, None)
|
132 |
-
|
133 |
def rwkv6_linear_attention_cpu(receptance, key, value, time_decay, time_first, state):
|
134 |
-
input_dtype = receptance.dtype
|
135 |
# For CPU fallback. Will be slower and probably take more memory than the custom CUDA kernel if not executed
|
136 |
# within a torch.no_grad.
|
137 |
-
batch, seq_length,
|
138 |
num_heads, head_size = time_first.shape
|
139 |
key = key.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2).transpose(-2, -1)
|
140 |
value = value.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2)
|
141 |
receptance = receptance.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2)
|
142 |
-
time_decay = torch.exp(-torch.exp(time_decay.float())).view(batch, seq_length, num_heads, head_size).permute(0, 2, 3, 1)
|
143 |
time_first = time_first.float().reshape(-1, 1, 1).reshape(num_heads, -1, 1)
|
144 |
out = torch.zeros_like(key).reshape(batch, seq_length, num_heads, head_size)
|
145 |
|
@@ -168,24 +88,26 @@ def rwkv6_linear_attention(
|
|
168 |
# Launching the CUDA kernel for just one token will actually be slower (there is no for loop in the CPU version
|
169 |
# in this case).
|
170 |
one_token = key.size(1) == 1
|
171 |
-
if not training or
|
172 |
return rwkv6_linear_attention_cpu(
|
173 |
receptance, key, value, time_decay, time_first, state
|
174 |
)
|
175 |
else:
|
176 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
177 |
|
178 |
|
179 |
class Rwkv6SelfAttention(nn.Module):
|
180 |
def __init__(self, config, layer_id=0):
|
181 |
super().__init__()
|
182 |
self.config = config
|
183 |
-
kernel_loaded = rwkv6_cuda_kernel is not None and rwkv6_cuda_kernel.head_size == config.head_size
|
184 |
-
if is_ninja_available() and is_torch_cuda_available() and not kernel_loaded:
|
185 |
-
try:
|
186 |
-
load_wkv6_cuda_kernel(config.head_size, config.max_context_length) # FIXME - context_length is not a configured attribute
|
187 |
-
except Exception:
|
188 |
-
logger.info("Could not load the custom CUDA kernel for RWKV6 attention.")
|
189 |
self.layer_id = layer_id
|
190 |
hidden_size = config.hidden_size
|
191 |
attention_hidden_size = config.attention_hidden_size
|
|
|
37 |
)
|
38 |
|
39 |
from .configuration_rwkv6 import Rwkv6Config
|
40 |
+
try:
|
41 |
+
from fla.ops.rwkv6.recurrent_fuse import fused_recurrent_rwkv6
|
42 |
+
except ImportError:
|
43 |
+
print("Required module is not installed. Please install it using the following commands:")
|
44 |
+
print("pip install -U git+https://github.com/sustcsonglin/flash-linear-attention")
|
45 |
+
print("Additionally, ensure you have the correct version of Triton installed:")
|
46 |
+
print("pip install triton==2.2.0")
|
47 |
|
48 |
|
49 |
logger = logging.get_logger(__name__)
|
|
|
51 |
_CHECKPOINT_FOR_DOC = "RWKV/rwkv-6-world-1b6"
|
52 |
_CONFIG_FOR_DOC = "Rwkv6Config"
|
53 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
54 |
def rwkv6_linear_attention_cpu(receptance, key, value, time_decay, time_first, state):
|
|
|
55 |
# For CPU fallback. Will be slower and probably take more memory than the custom CUDA kernel if not executed
|
56 |
# within a torch.no_grad.
|
57 |
+
batch, seq_length, _ = receptance.shape
|
58 |
num_heads, head_size = time_first.shape
|
59 |
key = key.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2).transpose(-2, -1)
|
60 |
value = value.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2)
|
61 |
receptance = receptance.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2)
|
62 |
+
time_decay = torch.exp(-torch.exp(time_decay.float())).view(batch, seq_length, num_heads, head_size).permute(0, 2, 3, 1)
|
63 |
time_first = time_first.float().reshape(-1, 1, 1).reshape(num_heads, -1, 1)
|
64 |
out = torch.zeros_like(key).reshape(batch, seq_length, num_heads, head_size)
|
65 |
|
|
|
88 |
# Launching the CUDA kernel for just one token will actually be slower (there is no for loop in the CPU version
|
89 |
# in this case).
|
90 |
one_token = key.size(1) == 1
|
91 |
+
if not training or no_cuda or one_token:
|
92 |
return rwkv6_linear_attention_cpu(
|
93 |
receptance, key, value, time_decay, time_first, state
|
94 |
)
|
95 |
else:
|
96 |
+
batch, seq_length, _ = receptance.shape
|
97 |
+
num_heads, head_size = time_first.shape
|
98 |
+
key = key.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2) # B, T, H, K -> B, H, T, K
|
99 |
+
value = value.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2) # B, T, H, K - > B, H, T, V
|
100 |
+
receptance = receptance.float().view(batch, seq_length, num_heads, head_size).transpose(1, 2) # B, H, T, K
|
101 |
+
time_decay = -torch.exp(time_decay.float()).view(batch, seq_length, num_heads, head_size).permute(0, 2, 1, 3) # B, T, H, K -> B, H, T, K
|
102 |
+
time_first = time_first.float().reshape(num_heads, head_size) # H, K
|
103 |
+
out, state = fused_recurrent_rwkv6(receptance, key, value, time_decay, time_first, scale=1.0, initial_state=state, output_final_state=True)
|
104 |
+
return out.transpose(1, 2), state
|
105 |
|
106 |
|
107 |
class Rwkv6SelfAttention(nn.Module):
|
108 |
def __init__(self, config, layer_id=0):
|
109 |
super().__init__()
|
110 |
self.config = config
|
|
|
|
|
|
|
|
|
|
|
|
|
111 |
self.layer_id = layer_id
|
112 |
hidden_size = config.hidden_size
|
113 |
attention_hidden_size = config.attention_hidden_size
|