Changed the implementation to fuse the transformers (#85)
This commit is contained in:
darknight054
@@ -1,139 +1,14 @@
|
||||
from typing import NamedTuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from torch.nn import functional as F
|
||||
from typing_extensions import Self
|
||||
|
||||
from pocket_tts.modules.layer_scale import LayerScale
|
||||
from pocket_tts.modules.rope import RotaryEmbedding
|
||||
from pocket_tts.modules.stateful_module import StatefulModule
|
||||
from pocket_tts.modules.transformer import StreamingMultiheadAttention
|
||||
from pocket_tts.utils.config import FlowLMTransformerConfig
|
||||
|
||||
|
||||
class KVCacheResult(NamedTuple):
|
||||
keys: torch.Tensor
|
||||
values: torch.Tensor
|
||||
positions: torch.Tensor
|
||||
|
||||
@staticmethod
|
||||
def from_kv(keys: torch.Tensor, values: torch.Tensor) -> "KVCacheResult":
|
||||
B, H, T, D = keys.shape
|
||||
assert tuple(values.shape[:-1]) == (B, H, T)
|
||||
positions = torch.arange(T, device=keys.device, dtype=torch.long)
|
||||
return KVCacheResult(keys, values, positions.expand(B, -1))
|
||||
|
||||
|
||||
def complete(
|
||||
cache: torch.Tensor, end_offset: torch.Tensor, k: torch.Tensor, v: torch.Tensor
|
||||
) -> KVCacheResult:
|
||||
capacity = cache.shape[3]
|
||||
assert k.shape[:-1] == v.shape[:-1], (k.shape, v.shape)
|
||||
B, H, T, D = k.shape
|
||||
assert T > 0
|
||||
indexes = torch.arange(T, device=end_offset.device, dtype=end_offset.dtype)
|
||||
indexes = indexes + end_offset.view(-1, 1)
|
||||
indexes = indexes % capacity
|
||||
# indexes is [B, T]
|
||||
# k is [B, H, T, D]
|
||||
# cache is [B, H, T', D]
|
||||
this_indexes = indexes.view(B, 1, T, 1)
|
||||
this_indexes = this_indexes.expand(-1, H, T, D)
|
||||
cache[0].scatter_(2, this_indexes, k)
|
||||
cache[1].scatter_(2, this_indexes, v)
|
||||
|
||||
keys = cache[0]
|
||||
values = cache[1]
|
||||
|
||||
indexes = torch.arange(capacity, device=end_offset.device, dtype=torch.long)
|
||||
|
||||
# end_index correspond to the actual index where the last value was written.
|
||||
last_offset = end_offset.view(-1, 1) + T - 1
|
||||
end_index = last_offset % capacity
|
||||
delta = indexes - end_index
|
||||
|
||||
positions = torch.where(delta <= 0, last_offset + delta, last_offset + delta - capacity)
|
||||
end_offset[:] = end_offset + T
|
||||
invalid = indexes >= end_offset.view(-1, 1)
|
||||
positions = torch.where(invalid, torch.full_like(positions, -1), positions)
|
||||
|
||||
return KVCacheResult(keys, values, positions)
|
||||
|
||||
|
||||
class MimiStreamingMultiheadAttention(StatefulModule):
|
||||
def __init__(self, embed_dim: int, num_heads: int, context: int, rope: RotaryEmbedding):
|
||||
super().__init__()
|
||||
|
||||
self.embed_dim = embed_dim
|
||||
self.context = context
|
||||
self.rope = rope
|
||||
self.num_heads = num_heads
|
||||
out_dim = 3 * embed_dim
|
||||
|
||||
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
|
||||
self.in_proj = nn.Linear(embed_dim, out_dim, bias=False)
|
||||
|
||||
def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
|
||||
dim_per_head = self.embed_dim // self.num_heads
|
||||
|
||||
state = {}
|
||||
device = self.in_proj.weight.device
|
||||
state["offset"] = torch.zeros(batch_size, dtype=torch.long, device=device)
|
||||
state["cache"] = torch.zeros(
|
||||
(2, batch_size, self.num_heads, sequence_length, dim_per_head), device=device
|
||||
)
|
||||
state["end_offset"] = torch.zeros(batch_size, dtype=torch.long, device=device)
|
||||
return state
|
||||
|
||||
def increment_step(self, state, increment: int = 1):
|
||||
state["offset"] += increment
|
||||
|
||||
def _complete_kv(self, k, v, model_state: dict | None) -> KVCacheResult:
|
||||
if model_state is None:
|
||||
return KVCacheResult.from_kv(k, v)
|
||||
else:
|
||||
layer_state = self.get_state(model_state)
|
||||
return complete(layer_state["cache"], layer_state["end_offset"], k, v)
|
||||
|
||||
def forward(self, query: torch.Tensor, model_state: dict | None) -> torch.Tensor:
|
||||
B, T = query.shape[:2]
|
||||
|
||||
if model_state is None:
|
||||
offset = torch.zeros(B, device=query.device, dtype=torch.long)
|
||||
else:
|
||||
offset = self.get_state(model_state)["offset"]
|
||||
|
||||
projected = self.in_proj(query)
|
||||
|
||||
q, k, v = rearrange(projected, "b t (p h d) -> p b h t d", p=3, h=self.num_heads)
|
||||
|
||||
# Permute from [b, h, t, d] to [b, t, h, d] for rope
|
||||
q = q.permute(0, 2, 1, 3)
|
||||
k = k.permute(0, 2, 1, 3)
|
||||
q, k = self.rope(q, k, offset)
|
||||
# Permute back from [b, t, h, d] to [b, h, t, d]
|
||||
q = q.permute(0, 2, 1, 3)
|
||||
k = k.permute(0, 2, 1, 3)
|
||||
|
||||
k, v, pos_k = self._complete_kv(k, v, model_state)
|
||||
pos_k = pos_k[:, None]
|
||||
pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(
|
||||
-1, 1
|
||||
)
|
||||
delta = pos_q - pos_k
|
||||
attn_bias = (pos_k >= 0) & (delta >= 0)
|
||||
attn_bias = attn_bias & (delta < self.context)
|
||||
attn_bias = attn_bias[:, None]
|
||||
|
||||
x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
|
||||
|
||||
x = rearrange(x, "b h t d -> b t (h d)")
|
||||
x = self.out_proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class StreamingTransformerLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -143,21 +18,11 @@ class StreamingTransformerLayer(nn.Module):
|
||||
context: int | None,
|
||||
rope: RotaryEmbedding,
|
||||
layer_scale: float | None = None,
|
||||
attention_kind: str = "mimi",
|
||||
):
|
||||
super().__init__()
|
||||
# Redefine self_attn to our streaming multi-head attention
|
||||
if attention_kind == "mimi":
|
||||
# TODO: we should actually use StreamingMultiheadAttention here and add context window
|
||||
# support. And we should then delete MimiStreamingMultiheadAttention.
|
||||
# The implementation is really close.
|
||||
self.self_attn = MimiStreamingMultiheadAttention(
|
||||
context=context, rope=rope, embed_dim=d_model, num_heads=num_heads
|
||||
)
|
||||
else:
|
||||
self.self_attn = StreamingMultiheadAttention(
|
||||
rope=rope, embed_dim=d_model, num_heads=num_heads
|
||||
)
|
||||
self.self_attn = StreamingMultiheadAttention(
|
||||
rope=rope, embed_dim=d_model, num_heads=num_heads, context=context
|
||||
)
|
||||
self.norm1 = nn.LayerNorm(d_model, eps=1e-5)
|
||||
self.norm2 = nn.LayerNorm(d_model, eps=1e-5)
|
||||
|
||||
@@ -199,7 +64,6 @@ class StreamingTransformer(nn.Module):
|
||||
dim_feedforward: int | list[int] = 2048,
|
||||
context: int | None = None,
|
||||
max_period: float = 10_000.0,
|
||||
kind: str = "mimi",
|
||||
):
|
||||
super().__init__()
|
||||
assert d_model % num_heads == 0
|
||||
@@ -217,7 +81,6 @@ class StreamingTransformer(nn.Module):
|
||||
context=context,
|
||||
rope=self.rope,
|
||||
layer_scale=layer_scale,
|
||||
attention_kind=kind,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -230,7 +93,6 @@ class StreamingTransformer(nn.Module):
|
||||
num_layers=config.num_layers,
|
||||
dim_feedforward=dim_feedforward,
|
||||
max_period=float(config.max_period),
|
||||
kind="flow_lm",
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, model_state: dict | None):
|
||||
|
||||
Reference in New Issue
Block a user