235 lines
8.6 KiB
Python
235 lines
8.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
from x_transformers.x_transformers import apply_rotary_pos_emb
|
|
|
|
|
|
def _apply_rope_cached(t, cos, sin):
|
|
"""Bit-exact RoPE for ``scale==1`` / interleaved-half.
|
|
Uses precomputed ``cos``/``sin`` per seq_len to skip per-step trig and cat."""
|
|
rot_dim = cos.shape[-1]
|
|
tr = t[..., :rot_dim]
|
|
x = tr.reshape(*tr.shape[:-1], -1, 2)
|
|
rot_half = torch.stack((-x[..., 1], x[..., 0]), dim=-1).reshape_as(tr)
|
|
out = tr * cos + rot_half * sin
|
|
if rot_dim < t.shape[-1]: # partial rotary (unused by ming heads); keep general
|
|
out = torch.cat((out, t[..., rot_dim:]), dim=-1)
|
|
return out.type(t.dtype)
|
|
|
|
|
|
class RMSNorm(nn.Module):
|
|
def __init__(self, dim, eps=1e-6):
|
|
super().__init__()
|
|
self.eps = eps
|
|
self.weight = nn.Parameter(torch.ones(dim))
|
|
|
|
def forward(self, x):
|
|
if self.weight.dtype in [torch.float16, torch.bfloat16]:
|
|
x = x.to(self.weight.dtype)
|
|
x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps)
|
|
return x
|
|
|
|
|
|
class FeedForward(nn.Module):
|
|
def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate="none"):
|
|
super().__init__()
|
|
inner_dim = int(dim * mult)
|
|
dim_out = dim_out if dim_out is not None else dim
|
|
|
|
activation = nn.GELU(approximate=approximate)
|
|
project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)
|
|
self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
|
|
|
|
def forward(self, x):
|
|
return self.ff(x)
|
|
|
|
|
|
class Attention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim,
|
|
heads=8,
|
|
dim_head=64,
|
|
dropout=0.0,
|
|
qk_norm=None,
|
|
pe_attn_head=None,
|
|
attn_mask_enabled=True,
|
|
):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.heads = heads
|
|
self.inner_dim = dim_head * heads
|
|
self.dropout = dropout
|
|
self.to_q = nn.Linear(dim, self.inner_dim)
|
|
self.to_k = nn.Linear(dim, self.inner_dim)
|
|
self.to_v = nn.Linear(dim, self.inner_dim)
|
|
if qk_norm is None:
|
|
self.q_norm = None
|
|
self.k_norm = None
|
|
elif qk_norm == "rms_norm":
|
|
self.q_norm = RMSNorm(dim_head)
|
|
self.k_norm = RMSNorm(dim_head)
|
|
else:
|
|
raise ValueError(f"Unimplemented qk_norm: {qk_norm}")
|
|
|
|
self.to_out = nn.ModuleList([])
|
|
self.to_out.append(nn.Linear(self.inner_dim, dim))
|
|
self.to_out.append(nn.Dropout(dropout))
|
|
self.pe_attn_head = pe_attn_head
|
|
self.attn_mask_enabled = attn_mask_enabled
|
|
# cos/sin cache keyed by seq_len (freqs are deterministic in seq_len).
|
|
self._rope_cache: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
|
|
# TODO: materialize fused weights eagerly outside any torch.compile
|
|
# scope (e.g. post-weight-load);
|
|
self._qkv_w: torch.Tensor | None = None
|
|
self._qkv_b: torch.Tensor | None = None
|
|
|
|
def _qkv(self, x):
|
|
if self._qkv_w is None:
|
|
self._qkv_w = torch.cat([self.to_q.weight, self.to_k.weight, self.to_v.weight], dim=0)
|
|
if self.to_q.bias is not None:
|
|
self._qkv_b = torch.cat([self.to_q.bias, self.to_k.bias, self.to_v.bias], dim=0)
|
|
qkv = F.linear(x, self._qkv_w, self._qkv_b)
|
|
return qkv.split(self.inner_dim, dim=-1)
|
|
|
|
def _rope_cos_sin(self, freqs):
|
|
seq_len = freqs.shape[-2]
|
|
cached = self._rope_cache.get(seq_len)
|
|
if cached is None:
|
|
f = freqs.unsqueeze(1) if freqs.ndim == 3 else freqs # b n d -> b 1 n d for 4D q/k
|
|
cached = (f.cos(), f.sin())
|
|
self._rope_cache[seq_len] = cached
|
|
return cached
|
|
|
|
def forward(self, x, mask=None, rope=None):
|
|
batch_size = x.shape[0]
|
|
query, key, value = self._qkv(x)
|
|
|
|
inner_dim = key.shape[-1]
|
|
head_dim = inner_dim // self.heads
|
|
query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
|
key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
|
value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
|
|
|
if self.q_norm is not None:
|
|
query = self.q_norm(query)
|
|
if self.k_norm is not None:
|
|
key = self.k_norm(key)
|
|
|
|
if rope is not None:
|
|
freqs, xpos_scale = rope
|
|
scale_is_one = xpos_scale is None or (not torch.is_tensor(xpos_scale) and xpos_scale == 1.0)
|
|
if scale_is_one and self.pe_attn_head is None:
|
|
# Fast path: cached, cat-free RoPE (bit-exact for scale==1, full rotary).
|
|
cos, sin = self._rope_cos_sin(freqs)
|
|
query = _apply_rope_cached(query, cos, sin)
|
|
key = _apply_rope_cached(key, cos, sin)
|
|
else:
|
|
q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
|
|
if self.pe_attn_head is not None:
|
|
on = self.pe_attn_head
|
|
query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale)
|
|
key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale)
|
|
else:
|
|
query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
|
|
key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
|
|
|
|
if self.attn_mask_enabled and mask is not None:
|
|
valid_sample_indices = mask.any(dim=1)
|
|
final_output = torch.zeros_like(query).to(query.device)
|
|
attn_mask = mask[valid_sample_indices]
|
|
query = query[valid_sample_indices]
|
|
key = key[valid_sample_indices]
|
|
value = value[valid_sample_indices]
|
|
attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)
|
|
attn_mask = attn_mask.expand(valid_sample_indices.sum().item(), self.heads, query.shape[-2], key.shape[-2])
|
|
else:
|
|
attn_mask = None
|
|
|
|
x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
|
|
if self.attn_mask_enabled and mask is not None:
|
|
final_output[valid_sample_indices] = x
|
|
x = final_output
|
|
|
|
x = x.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim)
|
|
x = x.to(query.dtype)
|
|
x = self.to_out[0](x)
|
|
x = self.to_out[1](x)
|
|
|
|
if mask is not None:
|
|
mask = mask.unsqueeze(-1)
|
|
x = x.masked_fill(~mask, 0.0)
|
|
|
|
return x
|
|
|
|
|
|
class DiTBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
hidden_size,
|
|
num_heads,
|
|
mlp_ratio=4.0,
|
|
dropout=0.1,
|
|
qk_norm=None,
|
|
pe_attn_head=None,
|
|
attn_mask_enabled=True,
|
|
**kwargs,
|
|
):
|
|
super().__init__()
|
|
self.norm1 = RMSNorm(hidden_size)
|
|
self.attn = Attention(
|
|
dim=hidden_size,
|
|
heads=num_heads,
|
|
dim_head=hidden_size // num_heads,
|
|
dropout=dropout,
|
|
qk_norm=qk_norm,
|
|
pe_attn_head=pe_attn_head,
|
|
attn_mask_enabled=attn_mask_enabled,
|
|
)
|
|
self.norm2 = RMSNorm(hidden_size)
|
|
self.mlp = FeedForward(dim=hidden_size, mult=mlp_ratio, dropout=dropout, approximate="tanh")
|
|
|
|
def forward(self, x, mask, rope):
|
|
x = x + self.attn(self.norm1(x), mask=mask, rope=rope)
|
|
x = x + self.mlp(self.norm2(x))
|
|
return x
|
|
|
|
|
|
class FinalLayer(nn.Module):
|
|
def __init__(self, hidden_size, out_channels):
|
|
super().__init__()
|
|
self.norm_final = RMSNorm(hidden_size)
|
|
self.linear = nn.Linear(hidden_size, out_channels, bias=True)
|
|
|
|
def forward(self, x):
|
|
x = self.norm_final(x)
|
|
x = self.linear(x)
|
|
return x
|
|
|
|
|
|
class CondEmbedder(nn.Module):
|
|
def __init__(self, input_feature_size, hidden_size):
|
|
super().__init__()
|
|
self.cond_embedder = nn.Linear(input_feature_size, hidden_size)
|
|
|
|
def forward(self, llm_cond):
|
|
return self.cond_embedder(llm_cond)
|
|
|
|
|
|
def get_epss_timesteps(n, device, dtype):
|
|
dt = 1 / 32
|
|
predefined_timesteps = {
|
|
5: [0, 2, 4, 8, 16, 32],
|
|
6: [0, 2, 4, 6, 8, 16, 32],
|
|
7: [0, 2, 4, 6, 8, 16, 24, 32],
|
|
10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],
|
|
12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],
|
|
16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],
|
|
}
|
|
t = predefined_timesteps.get(n, [])
|
|
if not t:
|
|
return torch.linspace(0, 1, n + 1, device=device, dtype=dtype)
|
|
return dt * torch.tensor(t, device=device, dtype=dtype)
|