chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,46 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from fairseq.model_parallel.megatron.mpu import (
|
||||
ColumnParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
|
||||
from .model_parallel_init import init_method
|
||||
from .kernel.rotary import apply_rotary_emb
|
||||
from flash_attn import flash_attn_func
|
||||
|
||||
class CrossAttention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
args,
|
||||
):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.embed_dim = args.dim
|
||||
self.num_heads = args.n_attn_heads // args.model_parallel_size
|
||||
self.num_kv_heads = args.n_attn_kv_heads // args.model_parallel_size
|
||||
|
||||
self.head_dim = args.dim // args.n_attn_heads
|
||||
self.q_proj = ColumnParallelLinear(args.dim, args.dim, bias=False, gather_output=False, init_method=init_method)
|
||||
self.out_proj = RowParallelLinear(args.dim, args.dim, bias=False, input_is_parallel=True, init_method=init_method)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
key,
|
||||
value,
|
||||
rel_pos
|
||||
):
|
||||
bsz, tgt_len, _ = x.size()
|
||||
|
||||
q = self.q_proj(x)
|
||||
q = q.view(bsz, tgt_len, self.num_heads, self.head_dim)
|
||||
q = apply_rotary_emb(q, *rel_pos, interleaved=True)
|
||||
|
||||
attn = flash_attn_func(q, key, value, causal=True)
|
||||
attn = attn.view(bsz, tgt_len, self.head_dim * self.num_heads)
|
||||
|
||||
attn = self.out_proj(attn)
|
||||
return attn
|
||||
@@ -0,0 +1,33 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from fairseq.model_parallel.megatron.mpu import (
|
||||
ColumnParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
|
||||
from .kernel.swiglu import swiglu
|
||||
from .model_parallel_init import init_method
|
||||
|
||||
class FeedForwardNetwork(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
embed_dim,
|
||||
ffn_dim,
|
||||
load_checkpoint=False,
|
||||
):
|
||||
super().__init__()
|
||||
self.embed_dim = embed_dim
|
||||
self.fc1 = ColumnParallelLinear(self.embed_dim, ffn_dim, bias=False, gather_output=False, init_method=init_method)
|
||||
self.gate = ColumnParallelLinear(self.embed_dim, ffn_dim, bias=False, gather_output=False, init_method=init_method)
|
||||
self.fc2 = RowParallelLinear(ffn_dim, self.embed_dim, bias=False, input_is_parallel=True, init_method=init_method)
|
||||
|
||||
def forward(self, x):
|
||||
x_shape = x.shape
|
||||
x = x.reshape(-1, x.size(-1))
|
||||
x = self.fc2(swiglu(self.fc1(x), self.gate(x)))
|
||||
output = x.view(x_shape)
|
||||
return output
|
||||
@@ -0,0 +1,87 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from fairseq.model_parallel.megatron.mpu import (
|
||||
ColumnParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
|
||||
from .rms_norm import RMSNorm
|
||||
|
||||
from .kernel.gate_recurrent import chunk_gate_retention, recurrent_gate_retention
|
||||
from .kernel.rotary import apply_rotary_emb
|
||||
from .kernel.swiglu import swiglu
|
||||
|
||||
from .model_parallel_init import qkvg_init_method, out_init_method
|
||||
|
||||
class GateRetention(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
args,
|
||||
gate_logit_normalizer: int = 16,
|
||||
):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.embed_dim = args.dim
|
||||
self.num_heads = args.n_self_heads // args.model_parallel_size
|
||||
self.head_dim = args.dim // args.n_self_heads
|
||||
|
||||
self.q_proj = ColumnParallelLinear(args.dim, args.dim, bias=False, gather_output=False, init_method=qkvg_init_method)
|
||||
self.k_proj = ColumnParallelLinear(args.dim, args.dim, bias=False, gather_output=False, init_method=qkvg_init_method)
|
||||
self.v_proj = ColumnParallelLinear(args.dim, args.dim, bias=False, gather_output=False, init_method=qkvg_init_method)
|
||||
self.g_proj = ColumnParallelLinear(args.dim, args.dim, bias=False, gather_output=False, init_method=qkvg_init_method)
|
||||
self.gt_proj = ColumnParallelLinear(args.dim, args.n_self_heads, bias=False, gather_output=False, init_method=qkvg_init_method)
|
||||
|
||||
self.out_proj = RowParallelLinear(args.dim, args.dim, bias=False, input_is_parallel=True, init_method=out_init_method)
|
||||
|
||||
self.subln = RMSNorm(self.head_dim, elementwise_affine=False, eps=args.norm_eps)
|
||||
|
||||
self.gate_logit_normalizer = gate_logit_normalizer
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
rel_pos,
|
||||
incremental_state=None,
|
||||
is_prefilling=False,
|
||||
):
|
||||
bsz, tgt_len, _ = x.size()
|
||||
|
||||
q = self.q_proj(x)
|
||||
k = self.k_proj(x)
|
||||
v = self.v_proj(x)
|
||||
g = self.g_proj(x)
|
||||
gt = self.gt_proj(x)
|
||||
|
||||
qr = q.view(bsz, tgt_len, self.num_heads, self.head_dim)
|
||||
kr = k.view(bsz, tgt_len, self.num_heads, self.head_dim)
|
||||
v = v.view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
gt = gt.view(bsz, tgt_len, self.num_heads).transpose(1, 2)
|
||||
|
||||
qr = apply_rotary_emb(qr, *rel_pos, interleaved=True).transpose(1, 2)
|
||||
kr = apply_rotary_emb(kr, *rel_pos, interleaved=True).transpose(1, 2)
|
||||
gt = (F.logsigmoid(gt) / self.gate_logit_normalizer)
|
||||
|
||||
if incremental_state is not None and not is_prefilling:
|
||||
o = recurrent_gate_retention(qr, kr, v, gt, incremental_state)
|
||||
else:
|
||||
if incremental_state is not None:
|
||||
index_mask = incremental_state["index_mask"]
|
||||
gt_sum = gt.float().masked_fill(index_mask, 0).sum(dim=-1, keepdim=True)
|
||||
gt_mask = (gt_sum - gt.float().cumsum(dim=-1)).exp().masked_fill(index_mask, 0)
|
||||
next_hidden_state = (kr.transpose(-1, -2) * (self.head_dim ** -0.5)) @ (v * gt_mask.to(v.dtype).unsqueeze(-1))
|
||||
if "last_hidden_state" in incremental_state:
|
||||
last_hidden_state = incremental_state["last_hidden_state"]
|
||||
next_hidden_state += last_hidden_state * gt_sum.exp().unsqueeze(-1).to(v.dtype) if last_hidden_state is not None else 0
|
||||
else:
|
||||
last_hidden_state = None
|
||||
incremental_state["last_hidden_state"] = next_hidden_state
|
||||
o = chunk_gate_retention(qr, kr, v, gt, chunk_size=256, last_hidden_state=last_hidden_state)
|
||||
else:
|
||||
o = chunk_gate_retention(qr, kr, v, gt, chunk_size=256)
|
||||
|
||||
o = self.subln(o).transpose(1, 2).reshape(bsz, tgt_len, self.num_heads * self.head_dim)
|
||||
o = swiglu(g, o)
|
||||
o = self.out_proj(o)
|
||||
return o
|
||||
@@ -0,0 +1,302 @@
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
|
||||
@triton.jit
|
||||
def _fwd_recurrence(
|
||||
S, d,
|
||||
O,
|
||||
NUM_HEAD, NUM_BLOCK,
|
||||
D_MODEL_K: tl.constexpr, D_MODEL_V: tl.constexpr,
|
||||
BLOCK_MODEL_K: tl.constexpr, BLOCK_MODEL_V: tl.constexpr,
|
||||
last_kv: Optional[tl.tensor]
|
||||
):
|
||||
offset_bh = tl.program_id(0)
|
||||
offset_d = tl.program_id(1)
|
||||
offset_s = tl.program_id(2)
|
||||
|
||||
S = S + offset_bh * NUM_BLOCK * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :]
|
||||
|
||||
O = O + offset_bh * NUM_BLOCK * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :]
|
||||
|
||||
if last_kv is not None:
|
||||
last_kv = last_kv + offset_bh * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :]
|
||||
acc = tl.load(last_kv).to(tl.float32)
|
||||
else:
|
||||
acc = tl.zeros([BLOCK_MODEL_K, BLOCK_MODEL_V], dtype=tl.float32)
|
||||
|
||||
tl.store(O, acc.to(O.dtype.element_ty))
|
||||
O += D_MODEL_K * D_MODEL_V
|
||||
d = d + offset_bh * NUM_BLOCK
|
||||
for i in range(NUM_BLOCK-1):
|
||||
d_i = tl.load(d)
|
||||
S_i = tl.load(S)
|
||||
acc = acc * d_i + S_i
|
||||
tl.store(O, acc.to(O.dtype.element_ty))
|
||||
d += 1
|
||||
S += D_MODEL_K * D_MODEL_V
|
||||
O += D_MODEL_K * D_MODEL_V
|
||||
|
||||
|
||||
## NUM_SPLIT_K/V. K/V dimension split into NUM_SPLIT_K/V parts with equal size BLOCK_MODEL
|
||||
@triton.jit
|
||||
def _bwd_recurrence(
|
||||
S, d,
|
||||
DI, DG, DL, DS,
|
||||
NUM_HEAD, NUM_BLOCK,
|
||||
D_MODEL_K: tl.constexpr, D_MODEL_V: tl.constexpr,
|
||||
BLOCK_MODEL_K: tl.constexpr, BLOCK_MODEL_V: tl.constexpr,
|
||||
|
||||
):
|
||||
offset_bh = tl.program_id(0)
|
||||
offset_d = tl.program_id(1)
|
||||
offset_s = tl.program_id(2)
|
||||
|
||||
# offset_h = offset_bh % NUM_HEAD
|
||||
NUM_K = D_MODEL_K // BLOCK_MODEL_K
|
||||
NUM_V = D_MODEL_V // BLOCK_MODEL_V
|
||||
# skip the last chunk because it is never used
|
||||
S = S + offset_bh * NUM_BLOCK * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :] + (NUM_BLOCK - 2) * D_MODEL_K * D_MODEL_V
|
||||
|
||||
DI = DI + offset_bh * NUM_BLOCK * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :] + (NUM_BLOCK - 2) * D_MODEL_K * D_MODEL_V
|
||||
|
||||
# start from the last chunk
|
||||
DS = DS + offset_bh * NUM_BLOCK * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :] + (NUM_BLOCK - 1) * D_MODEL_K * D_MODEL_V
|
||||
|
||||
DG = DG + offset_bh * NUM_BLOCK * NUM_K * NUM_V + offset_d * NUM_V + offset_s + (NUM_BLOCK - 2) * NUM_K * NUM_V
|
||||
|
||||
d = d + offset_bh * NUM_BLOCK + (NUM_BLOCK - 1)
|
||||
|
||||
Dacc = tl.zeros([BLOCK_MODEL_K, BLOCK_MODEL_V], dtype=tl.float32)
|
||||
|
||||
# ignore the first chunk
|
||||
for i in range(NUM_BLOCK - 1):
|
||||
S_i = tl.load(S)
|
||||
DS_i = tl.load(DS)
|
||||
d_i = tl.load(d)
|
||||
Dacc = Dacc * d_i + DS_i
|
||||
DG_i = tl.sum(Dacc * S_i.to(tl.float32))
|
||||
|
||||
tl.store(DG, DG_i.to(DG.dtype.element_ty))
|
||||
tl.store(DI, Dacc.to(DI.dtype.element_ty))
|
||||
|
||||
S -= D_MODEL_K * D_MODEL_V
|
||||
DI -= D_MODEL_K * D_MODEL_V
|
||||
DS -= D_MODEL_K * D_MODEL_V
|
||||
DG -= NUM_K * NUM_V
|
||||
d -= 1
|
||||
|
||||
DL = DL + offset_bh * D_MODEL_K * D_MODEL_V + offset_d * D_MODEL_V * BLOCK_MODEL_K + tl.arange(0, BLOCK_MODEL_K)[:, None] * D_MODEL_V + offset_s * BLOCK_MODEL_V + tl.arange(0, BLOCK_MODEL_V)[None, :]
|
||||
DS_i = tl.load(DS)
|
||||
d_i = tl.load(d)
|
||||
Dacc = Dacc * d_i + DS_i
|
||||
tl.store(DL, Dacc.to(DL.dtype.element_ty))
|
||||
|
||||
class ChunkGateRecurrent(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, kv, cross_decay, last_kv=None):
|
||||
cross_decay = cross_decay.contiguous()
|
||||
kv = kv.contiguous()
|
||||
|
||||
B, H, N, D_k, D_v = kv.shape
|
||||
output = torch.empty_like(kv)
|
||||
BLOCK_MODEL_K = 64
|
||||
BLOCK_MODEL_V = 16
|
||||
|
||||
assert D_k % BLOCK_MODEL_K == 0
|
||||
assert D_v % BLOCK_MODEL_V == 0
|
||||
|
||||
grid = (B*H, D_k//BLOCK_MODEL_K, D_v//BLOCK_MODEL_V)
|
||||
ctx.grid = grid
|
||||
ctx.have_last_kv = last_kv is not None
|
||||
ctx.BLOCK_MODEL_K = BLOCK_MODEL_K
|
||||
ctx.BLOCK_MODEL_V = BLOCK_MODEL_V
|
||||
|
||||
_fwd_recurrence[grid](
|
||||
kv,
|
||||
cross_decay,
|
||||
output,
|
||||
D_MODEL_K=D_k, D_MODEL_V=D_v,
|
||||
NUM_BLOCK=N, NUM_HEAD=H,
|
||||
BLOCK_MODEL_K=BLOCK_MODEL_K,
|
||||
BLOCK_MODEL_V=BLOCK_MODEL_V,
|
||||
last_kv=last_kv
|
||||
)
|
||||
|
||||
ctx.save_for_backward(output, cross_decay)
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, DO):
|
||||
DO = DO.contiguous()
|
||||
|
||||
output, cross_decay = ctx.saved_tensors
|
||||
|
||||
B, H, N, D_k, D_v = output.shape
|
||||
|
||||
BLOCK_MODEL_K = 64
|
||||
BLOCK_MODEL_V = 16
|
||||
|
||||
grid = (B*H, D_k//BLOCK_MODEL_K, D_v//BLOCK_MODEL_V)
|
||||
|
||||
DI = torch.empty_like(DO)
|
||||
DG = torch.empty(B*H, N, D_k//BLOCK_MODEL_K, D_v//BLOCK_MODEL_V, device=cross_decay.device, dtype=cross_decay.dtype)
|
||||
DL = torch.empty(B, H, D_k, D_v, device=output.device, dtype=output.dtype)
|
||||
_bwd_recurrence[grid](
|
||||
output, cross_decay,
|
||||
DI, DG, DL, DO,
|
||||
NUM_HEAD=H, NUM_BLOCK = N,
|
||||
D_MODEL_K = D_k,
|
||||
D_MODEL_V = D_v,
|
||||
BLOCK_MODEL_K=BLOCK_MODEL_K,
|
||||
BLOCK_MODEL_V=BLOCK_MODEL_V,
|
||||
)
|
||||
|
||||
DI[:, :, -1] = 0
|
||||
DG[:, -1] = 0
|
||||
DG = DG.view(B, H, N, -1).sum(dim=-1)
|
||||
return DI, DG, DL if ctx.have_last_kv else None
|
||||
|
||||
def cross_chunk(q, k, v, g, last_hidden_state=None):
|
||||
kv = k.transpose(-1, -2) @ (v * (-g + g[:, :, :, -1, None]).exp()[..., None].to(v.dtype))
|
||||
cross_decay = g[:, :, :, -1].exp().to(kv.dtype)
|
||||
S = chunk_gate_recurrent(kv, cross_decay, last_hidden_state)
|
||||
cross = (q * g[..., None].exp().to(q.dtype)) @ S
|
||||
return cross
|
||||
|
||||
@torch.compile
|
||||
def inner_chunk(q, k, v, g):
|
||||
attn = q @ k.transpose(-1, -2)
|
||||
causal_mask = torch.full([q.shape[-2], q.shape[-2]], float("-inf"), device=q.device).triu(1).type_as(q)
|
||||
attn = attn * (g[..., None] - g[..., None, :] + causal_mask).exp().to(attn.dtype)
|
||||
inner = attn @ v
|
||||
return inner
|
||||
|
||||
def chunk_gate_retention(q, k, v, g, chunk_size=64, last_hidden_state=None):
|
||||
bsz, num_head, tgt_len, key_dim = q.shape
|
||||
head_dim = v.shape[-1]
|
||||
num_chunk = tgt_len // chunk_size
|
||||
q = q.view(bsz, num_head, num_chunk, chunk_size, key_dim)
|
||||
k = k.view(bsz, num_head, num_chunk, chunk_size, key_dim) * (key_dim ** -0.5)
|
||||
v = v.view(bsz, num_head, num_chunk, chunk_size, head_dim)
|
||||
g = g.view(bsz, num_head, num_chunk, chunk_size)
|
||||
g = g.float().cumsum(-1)
|
||||
cross = cross_chunk(q, k, v, g, last_hidden_state=last_hidden_state)
|
||||
inner = inner_chunk(q, k, v, g)
|
||||
o = cross + inner
|
||||
return o.view(bsz, num_head, tgt_len, head_dim)
|
||||
|
||||
# for long sequence parallelism
|
||||
def hier_chunk_gate_retention(q, k, v, g, chunk_size=64, hier_chunk_size=16384):
|
||||
bsz, num_head, tgt_len, key_dim = q.shape
|
||||
head_dim = v.shape[-1]
|
||||
num_hier_chunk = tgt_len // hier_chunk_size
|
||||
assert tgt_len == num_hier_chunk * hier_chunk_size
|
||||
|
||||
q = q.view(bsz, num_head, num_hier_chunk, hier_chunk_size, key_dim)
|
||||
k = k.view(bsz, num_head, num_hier_chunk, hier_chunk_size, key_dim)
|
||||
v = v.view(bsz, num_head, num_hier_chunk, hier_chunk_size, head_dim)
|
||||
g = g.view(bsz, num_head, num_hier_chunk, hier_chunk_size)
|
||||
hier_cross = cross_chunk(q, k * (key_dim ** -0.5), v, g.float().cumsum(-1)).view(bsz, num_head, tgt_len, head_dim)
|
||||
|
||||
qi = q.transpose(1, 2).reshape(bsz * num_hier_chunk, num_head, hier_chunk_size, key_dim)
|
||||
ki = k.transpose(1, 2).reshape(bsz * num_hier_chunk, num_head, hier_chunk_size, key_dim)
|
||||
vi = v.transpose(1, 2).reshape(bsz * num_hier_chunk, num_head, hier_chunk_size, head_dim)
|
||||
gi = g.transpose(1, 2).reshape(bsz * num_hier_chunk, num_head, hier_chunk_size)
|
||||
inner_cross = chunk_gate_retention(qi, ki, vi, gi, chunk_size)
|
||||
|
||||
inner_cross = inner_cross.view(bsz, num_hier_chunk, num_head, hier_chunk_size, head_dim).transpose(1, 2).reshape(bsz, num_head, tgt_len, head_dim)
|
||||
o = hier_cross + inner_cross
|
||||
return o
|
||||
|
||||
def recurrent_gate_retention(q, k, v, g, incremental_state):
|
||||
bsz, num_head, _, key_dim = q.shape
|
||||
k *= key_dim ** -0.5
|
||||
g = g.view(bsz, num_head, 1, 1).float().exp()
|
||||
kv = k.transpose(-1, -2) * v
|
||||
if "last_hidden_state" in incremental_state:
|
||||
prev_kv = incremental_state["last_hidden_state"]
|
||||
kv += prev_kv * g.to(prev_kv.dtype)
|
||||
|
||||
incremental_state["last_hidden_state"] = kv
|
||||
o = q @ kv
|
||||
return o
|
||||
|
||||
def parallel_gate_retention(q, k, v, g):
|
||||
k = k * (q.shape[-1] ** -0.5)
|
||||
causal_mask = torch.full([q.shape[-2], q.shape[-2]], float("-inf"), device=q.device).triu(1).type_as(q)
|
||||
g = g.float().cumsum(-1)
|
||||
mask = g[..., None] - g[..., None, :] + causal_mask
|
||||
mask = mask.exp()
|
||||
|
||||
attn = q @ k.transpose(-1, -2)
|
||||
attn = attn * mask.to(attn.dtype)
|
||||
o = attn @ v
|
||||
return o
|
||||
|
||||
def naive_kv_recurrent(kv, cross_decay, last_kv=None):
|
||||
BSZ, NUM_HEAD, NUM_BLOCK, D_MODEL_K, D_MODEL_V = kv.shape
|
||||
kv_recurrent = []
|
||||
kv_state = torch.zeros(BSZ, NUM_HEAD, D_MODEL_K, D_MODEL_V, dtype=kv.dtype, device="cuda") if last_kv is None else last_kv
|
||||
# accumulate kv by loop
|
||||
for i in range(NUM_BLOCK):
|
||||
kv_recurrent.append(kv_state)
|
||||
kv_state = kv_state * cross_decay[:, :, i, None, None] + kv[:, :, i]
|
||||
|
||||
kv_recurrent = torch.stack(kv_recurrent, dim=2)
|
||||
return kv_recurrent
|
||||
|
||||
chunk_gate_recurrent = ChunkGateRecurrent.apply
|
||||
|
||||
def main():
|
||||
BSZ = 4
|
||||
NUM_HEAD = 4
|
||||
NUM_BLOCK = 16
|
||||
D_MODEL_K = 256
|
||||
D_MODEL_V = 432
|
||||
dtype = torch.float16
|
||||
kv = torch.randn(BSZ, NUM_HEAD, NUM_BLOCK, D_MODEL_K, D_MODEL_V, dtype=dtype, device="cuda")
|
||||
last_kv = torch.randn(BSZ, NUM_HEAD, D_MODEL_K, D_MODEL_V, dtype=dtype, device="cuda")
|
||||
kv_triton = kv.clone().detach()
|
||||
last_kv_triton = last_kv.clone().detach()
|
||||
cross_decay = torch.randn(BSZ, NUM_HEAD, NUM_BLOCK, dtype=dtype, device="cuda")
|
||||
cross_decay = torch.sigmoid(cross_decay)
|
||||
cross_decay_triton = cross_decay.clone().detach()
|
||||
grad_weight = torch.randn(BSZ, NUM_HEAD, NUM_BLOCK, D_MODEL_K, D_MODEL_V, dtype=dtype, device="cuda")
|
||||
kv.requires_grad = True
|
||||
kv_triton.requires_grad = True
|
||||
last_kv.requires_grad = True
|
||||
last_kv_triton.requires_grad = True
|
||||
cross_decay.requires_grad = True
|
||||
cross_decay_triton.requires_grad = True
|
||||
|
||||
start = time.time()
|
||||
kv_recurrent = naive_kv_recurrent(kv, cross_decay, last_kv)
|
||||
kv_recurrent.mul(grad_weight).sum().backward()
|
||||
print("naive time:", time.time() - start)
|
||||
|
||||
start = time.time()
|
||||
kv_recurrent_triton = chunk_gate_recurrent(kv_triton, cross_decay_triton, last_kv_triton)
|
||||
kv_recurrent_triton.mul(grad_weight).sum().backward()
|
||||
print("triton time:", time.time() - start)
|
||||
|
||||
print(torch.allclose(kv_recurrent, kv_recurrent_triton, atol=1e-3))
|
||||
print((kv_recurrent - kv_recurrent_triton).abs().max(), (kv_recurrent - kv_recurrent_triton).abs().mean())
|
||||
|
||||
print(torch.allclose(kv.grad, kv_triton.grad, atol=1e-3))
|
||||
print((kv.grad - kv_triton.grad).abs().max(), (kv.grad - kv_triton.grad).abs().mean())
|
||||
|
||||
print(torch.allclose(last_kv.grad, last_kv_triton.grad, atol=1e-3))
|
||||
print((last_kv.grad - last_kv_triton.grad).abs().max(), (last_kv.grad - last_kv_triton.grad).abs().mean())
|
||||
|
||||
print(torch.allclose(cross_decay.grad, cross_decay_triton.grad, atol=1e-3))
|
||||
print((cross_decay.grad - cross_decay_triton.grad).abs().max(), (cross_decay.grad - cross_decay_triton.grad).abs().mean())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,332 @@
|
||||
# Copyright (c) 2023, Tri Dao.
|
||||
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
|
||||
# @triton.autotune(
|
||||
# configs=[
|
||||
# triton.Config({"BLOCK_M": 2}),
|
||||
# triton.Config({"BLOCK_M": 4}),
|
||||
# triton.Config({"BLOCK_M": 8}),
|
||||
# triton.Config({"BLOCK_M": 16}),
|
||||
# ],
|
||||
# key=["CACHE_KEY_SEQLEN", "BLOCK_K", "INTERLEAVED"],
|
||||
# )
|
||||
@triton.jit
|
||||
def rotary_kernel(
|
||||
OUT, # Pointers to matrices
|
||||
X,
|
||||
COS,
|
||||
SIN,
|
||||
CU_SEQLENS,
|
||||
SEQLEN_OFFSETS, # this could be int or a pointer
|
||||
# Matrix dimensions
|
||||
seqlen,
|
||||
nheads,
|
||||
rotary_dim,
|
||||
seqlen_ro,
|
||||
CACHE_KEY_SEQLEN,
|
||||
# strides
|
||||
stride_out_batch,
|
||||
stride_out_seqlen,
|
||||
stride_out_nheads,
|
||||
stride_out_headdim,
|
||||
stride_x_batch,
|
||||
stride_x_seqlen,
|
||||
stride_x_nheads,
|
||||
stride_x_headdim,
|
||||
# Meta-parameters
|
||||
BLOCK_K: tl.constexpr,
|
||||
IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr,
|
||||
IS_VARLEN: tl.constexpr,
|
||||
INTERLEAVED: tl.constexpr,
|
||||
CONJUGATE: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
):
|
||||
pid_m = tl.program_id(axis=0)
|
||||
pid_batch = tl.program_id(axis=1)
|
||||
pid_head = tl.program_id(axis=2)
|
||||
rotary_dim_half = rotary_dim // 2
|
||||
|
||||
if not IS_VARLEN:
|
||||
X = X + pid_batch * stride_x_batch + pid_head * stride_x_nheads
|
||||
OUT = OUT + pid_batch * stride_out_batch + pid_head * stride_out_nheads
|
||||
else:
|
||||
start_idx = tl.load(CU_SEQLENS + pid_batch)
|
||||
seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx
|
||||
X = X + start_idx * stride_x_seqlen + pid_head * stride_x_nheads
|
||||
OUT = OUT + start_idx * stride_out_seqlen + pid_head * stride_out_nheads
|
||||
|
||||
if pid_m * BLOCK_M >= seqlen:
|
||||
return
|
||||
rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
if not IS_SEQLEN_OFFSETS_TENSOR:
|
||||
rm_cs = rm + SEQLEN_OFFSETS
|
||||
else:
|
||||
rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch)
|
||||
rk = tl.arange(0, BLOCK_K)
|
||||
rk_half = tl.arange(0, BLOCK_K // 2)
|
||||
|
||||
if not INTERLEAVED:
|
||||
# Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT
|
||||
X = X + (rm[:, None] * stride_x_seqlen + rk_half[None, :] * stride_x_headdim)
|
||||
COS = COS + (rm_cs[:, None] * rotary_dim_half + rk_half[None, :])
|
||||
SIN = SIN + (rm_cs[:, None] * rotary_dim_half + rk_half[None, :])
|
||||
cos = tl.load(
|
||||
COS, mask=(rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < rotary_dim_half), other=1.0
|
||||
).to(tl.float32)
|
||||
sin = tl.load(
|
||||
SIN, mask=(rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < rotary_dim_half), other=0.0
|
||||
).to(tl.float32)
|
||||
x0 = tl.load(
|
||||
X, mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half), other=0.0
|
||||
).to(tl.float32)
|
||||
x1 = tl.load(
|
||||
X + rotary_dim_half * stride_x_headdim,
|
||||
mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half),
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
if CONJUGATE:
|
||||
sin = -sin
|
||||
o0 = x0 * cos - x1 * sin
|
||||
o1 = x0 * sin + x1 * cos
|
||||
# write back result
|
||||
OUT = OUT + (rm[:, None] * stride_out_seqlen + rk_half[None, :] * stride_out_headdim)
|
||||
tl.store(OUT, o0, mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half))
|
||||
tl.store(
|
||||
OUT + rotary_dim_half * stride_out_headdim,
|
||||
o1,
|
||||
mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half),
|
||||
)
|
||||
else:
|
||||
# We don't want to load X[0, 2, 4, ...] and X[1, 3, 5, ...] separately since both are slow.
|
||||
# Instead, we load x0 = X[0, 1, 2, 3, ...] and x1 = X[1, 0, 3, 2, ...].
|
||||
# Loading x0 will be fast but x1 will be slow.
|
||||
# Then we load cos = COS[0, 0, 1, 1, ...] and sin = SIN[0, 0, 1, 1, ...].
|
||||
# Then we do the calculation and use tl.where to pick put the right outputs for the even
|
||||
# and for the odd indices.
|
||||
rk_swap = rk + ((rk + 1) % 2) * 2 - 1 # 1, 0, 3, 2, 5, 4, ...
|
||||
rk_repeat = tl.arange(0, BLOCK_K) // 2
|
||||
X0 = X + (rm[:, None] * stride_x_seqlen + rk[None, :] * stride_x_headdim)
|
||||
X1 = X + (rm[:, None] * stride_x_seqlen + rk_swap[None, :] * stride_x_headdim)
|
||||
COS = COS + (rm_cs[:, None] * rotary_dim_half + rk_repeat[None, :])
|
||||
SIN = SIN + (rm_cs[:, None] * rotary_dim_half + rk_repeat[None, :])
|
||||
cos = tl.load(
|
||||
COS,
|
||||
mask=(rm_cs[:, None] < seqlen_ro) & (rk_repeat[None, :] < rotary_dim_half),
|
||||
other=1.0,
|
||||
).to(tl.float32)
|
||||
sin = tl.load(
|
||||
SIN,
|
||||
mask=(rm_cs[:, None] < seqlen_ro) & (rk_repeat[None, :] < rotary_dim_half),
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
x0 = tl.load(X0, mask=(rm[:, None] < seqlen) & (rk[None, :] < rotary_dim), other=0.0).to(
|
||||
tl.float32
|
||||
)
|
||||
x1 = tl.load(
|
||||
X1, mask=(rm[:, None] < seqlen) & (rk_swap[None, :] < rotary_dim), other=0.0
|
||||
).to(tl.float32)
|
||||
if CONJUGATE:
|
||||
sin = -sin
|
||||
x0_cos = x0 * cos
|
||||
x1_sin = x1 * sin
|
||||
out = tl.where(rk[None, :] % 2 == 0, x0_cos - x1_sin, x0_cos + x1_sin)
|
||||
OUT = OUT + (rm[:, None] * stride_out_seqlen + rk[None, :] * stride_out_headdim)
|
||||
tl.store(OUT, out, mask=(rm[:, None] < seqlen) & (rk[None, :] < rotary_dim))
|
||||
|
||||
|
||||
def apply_rotary(
|
||||
x: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
seqlen_offsets: Union[int, torch.Tensor] = 0,
|
||||
cu_seqlens: Optional[torch.Tensor] = None,
|
||||
max_seqlen: Optional[int] = None,
|
||||
interleaved=False,
|
||||
inplace=False,
|
||||
conjugate=False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Arguments:
|
||||
x: (batch, seqlen, nheads, headdim) if cu_seqlens is None
|
||||
else (total_seqlen, nheads, headdim).
|
||||
cos: (seqlen_ro, rotary_dim / 2)
|
||||
sin: (seqlen_ro, rotary_dim / 2)
|
||||
seqlen_offsets: integer or integer tensor of size (batch,)
|
||||
cu_seqlens: (batch + 1,) or None
|
||||
max_seqlen: int
|
||||
Returns:
|
||||
y: (batch, seqlen, nheads, headdim)
|
||||
"""
|
||||
is_varlen = cu_seqlens is not None
|
||||
if not is_varlen:
|
||||
batch, seqlen, nheads, headdim = x.shape
|
||||
else:
|
||||
assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed"
|
||||
total_seqlen, nheads, headdim = x.shape
|
||||
batch_p_1 = cu_seqlens.shape[0]
|
||||
batch = batch_p_1 - 1
|
||||
seqlen = max_seqlen
|
||||
seqlen_ro, rotary_dim = cos.shape
|
||||
assert sin.shape == cos.shape
|
||||
rotary_dim *= 2
|
||||
assert rotary_dim <= headdim, "rotary_dim must be <= headdim"
|
||||
assert headdim <= 256, "Only support headdim <= 256"
|
||||
assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen"
|
||||
|
||||
assert (
|
||||
cos.dtype == sin.dtype
|
||||
), f"cos and sin must have the same dtype, got {cos.dtype} and {sin.dtype}"
|
||||
assert (
|
||||
x.dtype == cos.dtype
|
||||
), f"Input and cos/sin must have the same dtype, got {x.dtype} and {cos.dtype}"
|
||||
|
||||
cos, sin = cos.contiguous(), sin.contiguous()
|
||||
if isinstance(seqlen_offsets, torch.Tensor):
|
||||
assert seqlen_offsets.shape == (batch,)
|
||||
assert seqlen_offsets.dtype in [torch.int32, torch.int64]
|
||||
seqlen_offsets = seqlen_offsets.contiguous()
|
||||
else:
|
||||
assert seqlen_offsets + seqlen <= seqlen_ro
|
||||
|
||||
output = torch.empty_like(x) if not inplace else x
|
||||
if rotary_dim < headdim and not inplace:
|
||||
output[..., rotary_dim:].copy_(x[..., rotary_dim:])
|
||||
|
||||
BLOCK_K = (
|
||||
32
|
||||
if rotary_dim <= 32
|
||||
else (64 if rotary_dim <= 64 else (128 if rotary_dim <= 128 else 256))
|
||||
)
|
||||
grid = lambda META: (triton.cdiv(seqlen, META["BLOCK_M"]), batch, nheads) # noqa
|
||||
BLOCK_M = 4 if interleaved else (8 if rotary_dim <= 64 else 4)
|
||||
|
||||
# Need this, otherwise Triton tries to launch from cuda:0 and we get
|
||||
# ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)
|
||||
with torch.cuda.device(x.device.index):
|
||||
rotary_kernel[grid](
|
||||
output, # data ptrs
|
||||
x,
|
||||
cos,
|
||||
sin,
|
||||
cu_seqlens,
|
||||
seqlen_offsets,
|
||||
seqlen, # shapes
|
||||
nheads,
|
||||
rotary_dim,
|
||||
seqlen_ro,
|
||||
seqlen // 128, # key for triton cache (limit number of compilations)
|
||||
output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0
|
||||
output.stride(-3), # seqlen_stride or total_seqlen_stride
|
||||
output.stride(-2), # nheads_stride
|
||||
output.stride(-1), # headdim_stride
|
||||
x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0
|
||||
x.stride(-3), # seqlen stride or total_seqlen_stride
|
||||
x.stride(-2), # nheads stride
|
||||
x.stride(-1), # headdim stride
|
||||
BLOCK_K,
|
||||
isinstance(seqlen_offsets, torch.Tensor),
|
||||
is_varlen,
|
||||
interleaved,
|
||||
conjugate,
|
||||
BLOCK_M,
|
||||
)
|
||||
return output
|
||||
|
||||
class ApplyRotaryEmb(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx,
|
||||
x,
|
||||
cos,
|
||||
sin,
|
||||
interleaved=False,
|
||||
inplace=False,
|
||||
seqlen_offsets: Union[int, torch.Tensor] = 0,
|
||||
cu_seqlens: Optional[torch.Tensor] = None,
|
||||
max_seqlen: Optional[int] = None,
|
||||
):
|
||||
out = apply_rotary(
|
||||
x,
|
||||
cos,
|
||||
sin,
|
||||
seqlen_offsets=seqlen_offsets,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
interleaved=interleaved,
|
||||
inplace=inplace,
|
||||
)
|
||||
if isinstance(seqlen_offsets, int):
|
||||
# Can't save int with save_for_backward
|
||||
ctx.save_for_backward(cos, sin, cu_seqlens)
|
||||
ctx.seqlen_offsets = seqlen_offsets
|
||||
else:
|
||||
ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets)
|
||||
ctx.seqlen_offsets = None
|
||||
ctx.interleaved = interleaved
|
||||
ctx.inplace = inplace
|
||||
ctx.max_seqlen = max_seqlen
|
||||
return out if not inplace else x
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, do):
|
||||
seqlen_offsets = ctx.seqlen_offsets
|
||||
if seqlen_offsets is None:
|
||||
cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors
|
||||
else:
|
||||
cos, sin, cu_seqlens = ctx.saved_tensors
|
||||
# TD [2023-09-02]: For some reason Triton (2.0.0.post1) errors with
|
||||
# "[CUDA]: invalid device context", and cloning makes it work. Idk why. Triton 2.1.0 works.
|
||||
if not ctx.interleaved and not ctx.inplace:
|
||||
do = do.clone()
|
||||
dx = apply_rotary(
|
||||
do,
|
||||
cos,
|
||||
sin,
|
||||
seqlen_offsets=seqlen_offsets,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=ctx.max_seqlen,
|
||||
interleaved=ctx.interleaved,
|
||||
inplace=ctx.inplace,
|
||||
conjugate=True,
|
||||
)
|
||||
return dx, None, None, None, None, None, None, None
|
||||
|
||||
|
||||
def apply_rotary_emb(
|
||||
x,
|
||||
cos,
|
||||
sin,
|
||||
interleaved=False,
|
||||
inplace=False,
|
||||
seqlen_offsets: Union[int, torch.Tensor] = 0,
|
||||
cu_seqlens: Optional[torch.Tensor] = None,
|
||||
max_seqlen: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Arguments:
|
||||
x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None
|
||||
else (total_seqlen, nheads, headdim)
|
||||
cos, sin: (seqlen_rotary, rotary_dim / 2)
|
||||
interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
|
||||
of 1st half and 2nd half (GPT-NeoX style).
|
||||
inplace: if True, apply rotary embedding in-place.
|
||||
seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount.
|
||||
Most commonly used in inference when we have KV cache.
|
||||
cu_seqlens: (batch + 1,) or None
|
||||
max_seqlen: int
|
||||
Return:
|
||||
out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None
|
||||
else (total_seqlen, nheads, headdim)
|
||||
rotary_dim must be <= headdim
|
||||
Apply rotary embedding to the first rotary_dim of x.
|
||||
"""
|
||||
return ApplyRotaryEmb.apply(
|
||||
x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen
|
||||
)
|
||||
@@ -0,0 +1,32 @@
|
||||
import torch
|
||||
|
||||
|
||||
swiglu_fwd_codestring = """
|
||||
template <typename T> T swiglu_fwd(T x, T y) {
|
||||
return float(x) * float(y) / (1.0f + ::exp(-float(x)));
|
||||
}
|
||||
"""
|
||||
swiglu_bwd_codestring = """
|
||||
template <typename T> T swiglu_bwd(T x, T y, T g, T& dx, T& dy) {
|
||||
float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x)));
|
||||
dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y);
|
||||
dy = float(x) * x_sigmoid * float(g);
|
||||
}
|
||||
"""
|
||||
swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring)
|
||||
swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2)
|
||||
|
||||
|
||||
class SwiGLUFunction(torch.autograd.Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, x, y):
|
||||
ctx.save_for_backward(x, y)
|
||||
return swiglu_fwd(x, y)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dout):
|
||||
x, y = ctx.saved_tensors
|
||||
return swiglu_bwd(x, y, dout)
|
||||
|
||||
swiglu = SwiGLUFunction.apply
|
||||
@@ -0,0 +1,16 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
def init_method(tensor, **kwargs):
|
||||
nn.init.kaiming_uniform_(tensor, a=math.sqrt(5))
|
||||
|
||||
def qkvg_init_method(tensor, **kwargs):
|
||||
nn.init.xavier_uniform_(tensor, gain = 2 ** -2.5)
|
||||
|
||||
def out_init_method(tensor, **kwargs):
|
||||
nn.init.xavier_uniform_(tensor, gain = 2 ** -1)
|
||||
|
||||
def vocab_init_method(tensor, **kwargs):
|
||||
torch.nn.init.normal_(tensor, mean=0, std=tensor.shape[1] ** -0.5)
|
||||
@@ -0,0 +1,26 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine=True):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.eps = eps
|
||||
self.elementwise_affine = elementwise_affine
|
||||
if self.elementwise_affine:
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
else:
|
||||
self.register_parameter('weight', None)
|
||||
|
||||
def _norm(self, x):
|
||||
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||
|
||||
def forward(self, x):
|
||||
output = self._norm(x.float()).type_as(x)
|
||||
if self.weight is not None:
|
||||
output = output * self.weight
|
||||
return output
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return f'dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}'
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from fairseq.model_parallel.megatron.mpu import (
|
||||
ColumnParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
|
||||
from .model_parallel_init import init_method
|
||||
from .kernel.rotary import apply_rotary_emb
|
||||
|
||||
from flash_attn import flash_attn_func
|
||||
|
||||
class SlidingWindowAttention(nn.Module):
|
||||
def __init__(self, args):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.embed_dim = args.dim
|
||||
self.num_heads = args.n_self_heads // args.model_parallel_size
|
||||
self.window_size = args.sliding_window
|
||||
|
||||
self.head_dim = args.dim // args.n_self_heads
|
||||
|
||||
self.q_proj = ColumnParallelLinear(args.dim, args.dim, bias=False, gather_output=False, init_method=init_method)
|
||||
self.k_proj = ColumnParallelLinear(args.dim, args.dim, bias=False, gather_output=False, init_method=init_method)
|
||||
self.v_proj = ColumnParallelLinear(args.dim, args.dim, bias=False, gather_output=False, init_method=init_method)
|
||||
self.out_proj = RowParallelLinear(args.dim, args.dim, bias=False, input_is_parallel=True, init_method=init_method)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
rel_pos,
|
||||
start_pos=0,
|
||||
incremental_state=None,
|
||||
):
|
||||
bsz, tgt_len, embed_dim = x.size()
|
||||
src_len = tgt_len
|
||||
|
||||
q = self.q_proj(x)
|
||||
k = self.k_proj(x)
|
||||
v = self.v_proj(x)
|
||||
|
||||
q = q.view(bsz, tgt_len, self.num_heads, self.head_dim)
|
||||
k = k.view(bsz, src_len, self.num_heads, self.head_dim)
|
||||
v = v.view(bsz, src_len, self.num_heads, self.head_dim)
|
||||
|
||||
q = apply_rotary_emb(q, *rel_pos, interleaved=True)
|
||||
k = apply_rotary_emb(k, *rel_pos, interleaved=True)
|
||||
if incremental_state is not None:
|
||||
if "prev_key" not in incremental_state:
|
||||
incremental_state["prev_key"] = torch.empty(self.args.max_batch_size, self.window_size, self.num_heads, self.head_dim, device=x.device, dtype=x.dtype)
|
||||
incremental_state["prev_value"] = torch.empty(self.args.max_batch_size, self.window_size, self.num_heads, self.head_dim, device=x.device, dtype=x.dtype)
|
||||
|
||||
key = torch.cat([incremental_state["prev_key"][:bsz, :start_pos], k], dim=1)
|
||||
value = torch.cat([incremental_state["prev_value"][:bsz, :start_pos], v], dim=1)
|
||||
if key.shape[1] > self.window_size:
|
||||
incremental_state["prev_key"][:bsz] = key[:, -self.window_size:]
|
||||
incremental_state["prev_value"][:bsz] = value[:, -self.window_size:]
|
||||
else:
|
||||
incremental_state["prev_key"][:bsz, start_pos : start_pos + tgt_len] = k
|
||||
incremental_state["prev_value"][:bsz, start_pos : start_pos + tgt_len] = v
|
||||
else:
|
||||
key, value = k, v
|
||||
attn = flash_attn_func(q, key, value, causal=True, window_size=(self.window_size - 1, 0))
|
||||
attn = attn.reshape(bsz, tgt_len, self.head_dim * self.num_heads)
|
||||
|
||||
attn = self.out_proj(attn)
|
||||
return attn
|
||||
@@ -0,0 +1,251 @@
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from flash_attn import flash_attn_func
|
||||
|
||||
from fairseq.model_parallel.megatron.mpu import (
|
||||
ColumnParallelLinear,
|
||||
RowParallelLinear,
|
||||
copy_to_model_parallel_region,
|
||||
VocabParallelEmbedding
|
||||
)
|
||||
|
||||
from fairscale.nn import checkpoint_wrapper
|
||||
|
||||
from .rms_norm import RMSNorm
|
||||
from .kernel.rotary import apply_rotary_emb
|
||||
from .model_parallel_init import init_method, vocab_init_method
|
||||
|
||||
|
||||
def precompute_freqs_cis(dim: int, end: int, theta: float) -> torch.Tensor:
|
||||
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
|
||||
t = torch.arange(end, device=freqs.device) # type: ignore
|
||||
freqs = torch.outer(t, freqs).float() # type: ignore
|
||||
return freqs
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArgs:
|
||||
dim: int
|
||||
n_layers: int
|
||||
head_dim: int
|
||||
hidden_dim: int
|
||||
n_heads: int
|
||||
n_kv_heads: int
|
||||
norm_eps: float
|
||||
vocab_size: int
|
||||
|
||||
max_batch_size: int = 0
|
||||
max_seq_len: int = -1
|
||||
model_parallel_size: int = 1
|
||||
load_checkpoint: bool = False
|
||||
rope_theta: float = 10000.0
|
||||
sliding_window: Optional[int] = None
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
|
||||
self.dim = args.dim
|
||||
self.head_dim = args.head_dim
|
||||
self.hidden_dim = args.n_heads * args.head_dim
|
||||
self.key_value_dim = args.n_kv_heads * args.head_dim
|
||||
self.n_heads = args.n_heads // args.model_parallel_size
|
||||
self.n_kv_heads = args.n_kv_heads // args.model_parallel_size
|
||||
self.activate_sliding_window = args.sliding_window is not None
|
||||
self.cache_len = args.sliding_window - 1 if self.activate_sliding_window else args.max_seq_len
|
||||
|
||||
self.repeats = self.n_heads // self.n_kv_heads
|
||||
|
||||
self.scale = self.args.head_dim**-0.5
|
||||
|
||||
self.wq = ColumnParallelLinear(self.dim, self.hidden_dim, bias=False, gather_output=False, init_method=init_method)
|
||||
self.wk = ColumnParallelLinear(self.dim, self.key_value_dim, bias=False, gather_output=False, init_method=init_method)
|
||||
self.wv = ColumnParallelLinear(self.dim, self.key_value_dim, bias=False, gather_output=False, init_method=init_method)
|
||||
self.wo = RowParallelLinear(self.hidden_dim, self.dim, bias=False, input_is_parallel=True, init_method=init_method)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
rel_pos: Tuple[torch.Tensor, torch.Tensor],
|
||||
start_pos: int,
|
||||
incremental_state = None,
|
||||
) -> torch.Tensor:
|
||||
bsz, seqlen, _ = x.shape
|
||||
|
||||
xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
|
||||
xq = xq.view(bsz, seqlen, self.n_heads, self.head_dim)
|
||||
xk = xk.view(bsz, seqlen, self.n_kv_heads, self.head_dim)
|
||||
xv = xv.view(bsz, seqlen, self.n_kv_heads, self.head_dim)
|
||||
xq = apply_rotary_emb(xq, *rel_pos)
|
||||
xk = apply_rotary_emb(xk, *rel_pos)
|
||||
if incremental_state is not None:
|
||||
if "cache_k" not in incremental_state:
|
||||
incremental_state["cache_k"] = torch.zeros(
|
||||
(
|
||||
self.args.max_batch_size,
|
||||
self.cache_len,
|
||||
self.n_kv_heads,
|
||||
self.head_dim,
|
||||
)
|
||||
).to(xk)
|
||||
incremental_state["cache_v"] = torch.zeros(
|
||||
(
|
||||
self.args.max_batch_size,
|
||||
self.cache_len,
|
||||
self.n_kv_heads,
|
||||
self.head_dim,
|
||||
)
|
||||
).to(xv)
|
||||
key = torch.cat([incremental_state["cache_k"][:, :start_pos], xk], dim=1)
|
||||
value = torch.cat([incremental_state["cache_v"][:, :start_pos], xv], dim=1)
|
||||
if key.shape[1] > self.cache_len:
|
||||
incremental_state["cache_k"][:bsz] = key[:, -self.cache_len:]
|
||||
incremental_state["cache_v"][:bsz] = value[:, -self.cache_len:]
|
||||
else:
|
||||
incremental_state["cache_k"][:bsz, start_pos : start_pos + seqlen] = xk
|
||||
incremental_state["cache_v"][:bsz, start_pos : start_pos + seqlen] = xv
|
||||
else:
|
||||
key, value = xk, xv
|
||||
|
||||
output = flash_attn_func(xq, key, value, causal=True, window_size=(self.args.sliding_window - 1, 0) if self.activate_sliding_window else (-1, -1))
|
||||
|
||||
return self.wo(output.view(bsz, seqlen, self.n_heads * self.head_dim))
|
||||
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.w1 = ColumnParallelLinear(args.dim, args.hidden_dim, bias=False, gather_output=False, init_method=init_method)
|
||||
self.w2 = RowParallelLinear(args.hidden_dim, args.dim, bias=False, input_is_parallel=True, init_method=init_method)
|
||||
self.w3 = ColumnParallelLinear(args.dim, args.hidden_dim, bias=False, gather_output=False, init_method=init_method)
|
||||
|
||||
def forward(self, x) -> torch.Tensor:
|
||||
return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
self.n_heads = args.n_heads
|
||||
self.dim = args.dim
|
||||
self.attention = Attention(args)
|
||||
self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
||||
self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
||||
self.args = args
|
||||
|
||||
self.feed_forward: nn.Module
|
||||
self.feed_forward = FeedForward(args=args)
|
||||
|
||||
def forward(
|
||||
self, x: torch.Tensor, rel_pos: Tuple[torch.Tensor, torch.Tensor], start_pos: int, incremental_state = None
|
||||
) -> torch.Tensor:
|
||||
r = self.attention.forward(self.attention_norm(x), rel_pos, start_pos, incremental_state)
|
||||
h = x + r
|
||||
r = self.feed_forward.forward(self.ffn_norm(h))
|
||||
out = h + r
|
||||
return out
|
||||
|
||||
|
||||
class Transformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
args: ModelArgs,
|
||||
mp_rank: int = 0,
|
||||
checkpoint_activations: bool = False
|
||||
):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.vocab_size = args.vocab_size
|
||||
self.n_layers = args.n_layers
|
||||
self._precomputed_freqs_cis: Optional[torch.Tensor] = None
|
||||
self._window_precomputed_freqs_cis: Optional[torch.Tensor] = None
|
||||
self._global_precomputed_freqs_cis: Optional[torch.Tensor] = None
|
||||
assert self.vocab_size > 0
|
||||
self.mp_rank = mp_rank
|
||||
self.checkpoint_activations = checkpoint_activations
|
||||
self.tok_embeddings = VocabParallelEmbedding(
|
||||
args.vocab_size, args.dim, -1, init_method=vocab_init_method
|
||||
)
|
||||
self.norm = RMSNorm(args.dim, eps=args.norm_eps)
|
||||
self.output = nn.Linear(args.dim, args.vocab_size // args.model_parallel_size, bias=False)
|
||||
# Initialize all layers but slice off those not of this rank.
|
||||
layers = [TransformerBlock(args=args) for idx in range(args.n_layers)]
|
||||
if checkpoint_activations:
|
||||
layers = [checkpoint_wrapper(layer) for layer in layers]
|
||||
self.layers = nn.ModuleList(layers)
|
||||
self.n_local_layers = len(self.layers)
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
return next(self.parameters()).dtype
|
||||
|
||||
@property
|
||||
def device(self) -> torch.device:
|
||||
return next(self.parameters()).device
|
||||
|
||||
def build_rel_pos(self, x, start_pos):
|
||||
if self._precomputed_freqs_cis is None:
|
||||
theta = self.args.rope_theta
|
||||
self._precomputed_freqs_cis = precompute_freqs_cis(
|
||||
self.args.head_dim, self.args.max_seq_len, theta
|
||||
)
|
||||
if self._precomputed_freqs_cis.device != self.device:
|
||||
self._precomputed_freqs_cis = self._precomputed_freqs_cis.to(
|
||||
device=self.device
|
||||
)
|
||||
cos = torch.cos(self._precomputed_freqs_cis[start_pos:start_pos+x.size(1)])
|
||||
sin = torch.sin(self._precomputed_freqs_cis[start_pos:start_pos+x.size(1)])
|
||||
rel_pos = (cos.to(x.dtype), sin.to(x.dtype))
|
||||
return rel_pos
|
||||
|
||||
def forward_partial(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
start_pos: Optional[int] = 0,
|
||||
incremental_state = None,
|
||||
) -> torch.Tensor:
|
||||
h = self.tok_embeddings(input_ids)
|
||||
rel_pos = self.build_rel_pos(h, start_pos)
|
||||
for local_layer_id, layer in enumerate(self.layers):
|
||||
if incremental_state is not None:
|
||||
if local_layer_id not in incremental_state:
|
||||
incremental_state[local_layer_id] = {}
|
||||
h = layer(h, rel_pos, start_pos, incremental_state=incremental_state[local_layer_id] if incremental_state is not None else None)
|
||||
|
||||
return self.norm(h)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
start_pos: Optional[int] = 0,
|
||||
incremental_state = None,
|
||||
) -> torch.Tensor:
|
||||
h = self.forward_partial(input_ids, start_pos, incremental_state)
|
||||
if self.args.model_parallel_size > 1:
|
||||
h = copy_to_model_parallel_region(h)
|
||||
outs = self.output(h)
|
||||
return outs.float(), None
|
||||
|
||||
def load_state_dict(self, state_dict, strict=False, assign=False):
|
||||
state_to_load = {}
|
||||
for k, v in state_dict.items():
|
||||
if k.startswith("tok_embeddings") or k.startswith("output"):
|
||||
state_to_load[k] = v.view(self.args.model_parallel_size, self.vocab_size // self.args.model_parallel_size, self.args.dim)[self.mp_rank]
|
||||
elif "wq" in k or "wk" in k or "wv" in k or "w1" in k or "w3" in k:
|
||||
state_to_load[k] = v.view(self.args.model_parallel_size, -1, v.shape[1])[self.mp_rank]
|
||||
elif "wo" in k or "w2" in k:
|
||||
state_to_load[k] = v.view(v.shape[0], self.args.model_parallel_size, -1)[:, self.mp_rank]
|
||||
else:
|
||||
state_to_load[k] = v
|
||||
super().load_state_dict(state_to_load, strict=False, assign=assign)
|
||||
print("Loaded state dict from checkpoint.")
|
||||
@@ -0,0 +1,294 @@
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from fairscale.nn import checkpoint_wrapper
|
||||
|
||||
from fairseq.model_parallel.megatron.mpu import (
|
||||
ColumnParallelLinear,
|
||||
copy_to_model_parallel_region,
|
||||
VocabParallelEmbedding
|
||||
)
|
||||
|
||||
from .gate_retention import GateRetention
|
||||
from .sliding_window_attention import SlidingWindowAttention
|
||||
from .cross_attention import CrossAttention
|
||||
from .feedforward_network import FeedForwardNetwork, init_method
|
||||
from .rms_norm import RMSNorm
|
||||
|
||||
from .kernel.rotary import apply_rotary_emb
|
||||
from .model_parallel_init import vocab_init_method, init_method
|
||||
|
||||
|
||||
@dataclass
|
||||
class YOCOArgs:
|
||||
dim: int
|
||||
n_layers: int
|
||||
hidden_dim: int
|
||||
n_self_heads: int
|
||||
n_attn_heads: int
|
||||
n_attn_kv_heads: int
|
||||
vocab_size: int
|
||||
|
||||
max_batch_size: int = 0
|
||||
max_seq_len: int = -1
|
||||
model_parallel_size: int = 1
|
||||
load_checkpoint: bool = False
|
||||
rope_theta: float = 10000.0
|
||||
norm_eps: float = 1e-5
|
||||
sliding_window: Optional[int] = None
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
args: YOCOArgs,
|
||||
is_cross_layer=False
|
||||
):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.is_cross_layer = is_cross_layer
|
||||
|
||||
if is_cross_layer:
|
||||
self.mixer = CrossAttention(args)
|
||||
elif args.sliding_window is not None:
|
||||
self.mixer = SlidingWindowAttention(args)
|
||||
else:
|
||||
self.mixer = GateRetention(args)
|
||||
|
||||
self.mixer_layer_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
||||
|
||||
self.ffn = FeedForwardNetwork(
|
||||
args.dim,
|
||||
args.hidden_dim,
|
||||
args.load_checkpoint
|
||||
)
|
||||
|
||||
self.final_layer_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
start_pos=0,
|
||||
key=None,
|
||||
value=None,
|
||||
rel_pos=None,
|
||||
incremental_state=None,
|
||||
is_prefilling=False,
|
||||
):
|
||||
residual = x
|
||||
x = self.mixer_layer_norm(x)
|
||||
|
||||
if self.is_cross_layer:
|
||||
x = self.mixer(
|
||||
x,
|
||||
key,
|
||||
value,
|
||||
rel_pos=rel_pos,
|
||||
)
|
||||
elif self.args.sliding_window is not None:
|
||||
x = self.mixer(
|
||||
x,
|
||||
rel_pos=rel_pos,
|
||||
start_pos=start_pos,
|
||||
incremental_state=incremental_state,
|
||||
)
|
||||
else:
|
||||
x = self.mixer(
|
||||
x,
|
||||
rel_pos=rel_pos,
|
||||
incremental_state=incremental_state,
|
||||
is_prefilling=is_prefilling,)
|
||||
|
||||
x = x + residual
|
||||
residual = x
|
||||
x = self.final_layer_norm(x)
|
||||
|
||||
x = self.ffn(x)
|
||||
|
||||
x = x + residual
|
||||
return x
|
||||
|
||||
class SelfDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
args: YOCOArgs,
|
||||
checkpoint_activations: bool = False
|
||||
):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
layers = [DecoderLayer(args, is_cross_layer=False,) for idx in range(args.n_layers // 2)]
|
||||
if checkpoint_activations:
|
||||
layers = [checkpoint_wrapper(layer) for layer in layers]
|
||||
self.layers = nn.ModuleList(layers)
|
||||
self.head_dim = args.dim // args.n_self_heads
|
||||
self.block_size = 256
|
||||
self._precomputed_freqs_cis = None
|
||||
|
||||
def build_rel_pos(self, x, start_pos):
|
||||
if self._precomputed_freqs_cis is None:
|
||||
angle = 1.0 / (self.args.rope_theta ** torch.linspace(0, 1, self.head_dim // 2, dtype=torch.float, device=x.device))
|
||||
index = torch.arange(self.args.max_seq_len).to(angle)
|
||||
self._precomputed_freqs_cis = index[:, None] * angle
|
||||
|
||||
cos = torch.cos(self._precomputed_freqs_cis[start_pos:start_pos+x.size(1)])
|
||||
sin = torch.sin(self._precomputed_freqs_cis[start_pos:start_pos+x.size(1)])
|
||||
rel_pos = (cos.to(x.dtype), sin.to(x.dtype))
|
||||
return rel_pos
|
||||
|
||||
def get_index_mask(self, x, length, pad_length):
|
||||
return torch.arange(pad_length, device=x.device) >= length
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
incremental_state=None,
|
||||
is_prefilling=False,
|
||||
start_pos=0
|
||||
):
|
||||
if is_prefilling and x.size(1) % self.block_size != 0 and self.args.sliding_window is None:
|
||||
padding_len = self.block_size - x.size(1) % self.block_size
|
||||
x = F.pad(x, (0, 0, 0, padding_len), value=0)
|
||||
else:
|
||||
padding_len = 0
|
||||
|
||||
if incremental_state is not None and is_prefilling:
|
||||
index_mask = self.get_index_mask(x, x.size(1) - padding_len, x.size(1))
|
||||
|
||||
rel_pos = self.build_rel_pos(x, start_pos)
|
||||
for idx, layer in enumerate(self.layers):
|
||||
if incremental_state is not None:
|
||||
if idx not in incremental_state:
|
||||
incremental_state[idx] = {}
|
||||
if is_prefilling:
|
||||
incremental_state[idx]["index_mask"] = index_mask
|
||||
x = layer(
|
||||
x,
|
||||
start_pos=start_pos,
|
||||
rel_pos=rel_pos,
|
||||
incremental_state=incremental_state[idx] if incremental_state is not None else None,
|
||||
is_prefilling=is_prefilling,)
|
||||
|
||||
x = x[:, :x.size(1) - padding_len, :]
|
||||
return x
|
||||
|
||||
class CrossDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
args: YOCOArgs,
|
||||
checkpoint_activations: bool = False
|
||||
):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.num_heads = args.n_attn_kv_heads
|
||||
self.head_dim = args.dim // args.n_attn_heads
|
||||
self.k_proj = ColumnParallelLinear(args.dim, self.head_dim * args.n_attn_kv_heads, bias=False, gather_output=False, init_method=init_method)
|
||||
self.v_proj = ColumnParallelLinear(args.dim, self.head_dim * args.n_attn_kv_heads, bias=False, gather_output=False, init_method=init_method)
|
||||
self.kv_layer_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
||||
layers = [DecoderLayer(args, is_cross_layer=True) for idx in range(args.n_layers // 2)]
|
||||
if checkpoint_activations:
|
||||
layers = [checkpoint_wrapper(layer) for layer in layers]
|
||||
self.layers = nn.ModuleList(layers)
|
||||
self._precomputed_freqs_cis = None
|
||||
|
||||
def build_rel_pos(self, x, start_pos):
|
||||
if self._precomputed_freqs_cis is None:
|
||||
angle = 1.0 / (self.args.rope_theta ** torch.linspace(0, 1, self.head_dim // 2, dtype=torch.float, device=x.device))
|
||||
index = torch.arange(self.args.max_seq_len).to(angle)
|
||||
self._precomputed_freqs_cis = index[:, None] * angle
|
||||
|
||||
cos = torch.cos(self._precomputed_freqs_cis[start_pos:start_pos+x.size(1)])
|
||||
sin = torch.sin(self._precomputed_freqs_cis[start_pos:start_pos+x.size(1)])
|
||||
rel_pos = (cos.to(x.dtype), sin.to(x.dtype))
|
||||
return rel_pos
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
incremental_state=None,
|
||||
start_pos=0,
|
||||
skip_cross_decoder=False,
|
||||
):
|
||||
bsz, seqlen, embed_dim = x.size()
|
||||
x_norm = self.kv_layer_norm(x)
|
||||
key, value = self.k_proj(x_norm), self.v_proj(x_norm)
|
||||
key = key.view(bsz, seqlen, self.num_heads, self.head_dim)
|
||||
value = value.view(bsz, seqlen, self.num_heads, self.head_dim)
|
||||
rel_pos = self.build_rel_pos(x, start_pos)
|
||||
key = apply_rotary_emb(key, *rel_pos, interleaved=True)
|
||||
if incremental_state is not None:
|
||||
if "prev_key" not in incremental_state:
|
||||
incremental_state["prev_key"] = torch.empty(bsz, self.args.max_seq_len, self.num_heads, self.head_dim, device=x.device, dtype=x.dtype)
|
||||
incremental_state["prev_value"] = torch.empty(bsz, self.args.max_seq_len, self.num_heads, self.head_dim, device=x.device, dtype=x.dtype)
|
||||
incremental_state["prev_key"][:, start_pos : start_pos + seqlen] = key
|
||||
incremental_state["prev_value"][:, start_pos : start_pos + seqlen] = value
|
||||
key = incremental_state["prev_key"][:, : start_pos + seqlen]
|
||||
value = incremental_state["prev_value"][:, : start_pos + seqlen]
|
||||
|
||||
if skip_cross_decoder:
|
||||
return torch.zeros(bsz, 1, embed_dim, device=x.device, dtype=x.dtype)
|
||||
for layer in self.layers:
|
||||
x = layer(
|
||||
x,
|
||||
key=key,
|
||||
value=value,
|
||||
rel_pos=rel_pos)
|
||||
|
||||
return x
|
||||
|
||||
class YOCO(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
args: YOCOArgs,
|
||||
checkpoint_activations: bool = False,
|
||||
share_input_output_embed: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.embed_scale = math.sqrt(args.dim)
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
args.vocab_size, args.dim, -1, init_method=vocab_init_method
|
||||
)
|
||||
self.output_projection = nn.Linear(args.dim, args.vocab_size, bias=False)
|
||||
if share_input_output_embed:
|
||||
self.output_projection.weight = self.embed_tokens.weight
|
||||
|
||||
self.self_decoder = SelfDecoder(args, checkpoint_activations)
|
||||
self.cross_decoder = CrossDecoder(args, checkpoint_activations)
|
||||
self.layer_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
start_pos=0,
|
||||
incremental_state=None,
|
||||
is_prefilling=True,
|
||||
skip_cross_decoder=False
|
||||
):
|
||||
x = self.embed_scale * self.embed_tokens(x)
|
||||
|
||||
x = self.self_decoder(
|
||||
x,
|
||||
incremental_state=incremental_state,
|
||||
is_prefilling=is_prefilling,
|
||||
start_pos=start_pos,
|
||||
)
|
||||
|
||||
x = self.cross_decoder(
|
||||
x,
|
||||
start_pos=start_pos,
|
||||
incremental_state=incremental_state,
|
||||
skip_cross_decoder=skip_cross_decoder,
|
||||
)
|
||||
|
||||
x = self.layer_norm(x)
|
||||
x = self.output_layer(x)
|
||||
|
||||
return x, None
|
||||
|
||||
def output_layer(self, features):
|
||||
if self.args.model_parallel_size > 1:
|
||||
features = copy_to_model_parallel_region(features)
|
||||
return self.output_projection(features)
|
||||
Reference in New Issue
Block a user