akshatvishuandHongsheng Liu
2fa4c7a4a9
Signed-off-by: akshatvishu <akshatnayak197@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
194 lines
6.5 KiB
Python
194 lines
6.5 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
|
|
|
|
|
|
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
|
|
|
|
def forward(self, x, mask=None, rope=None):
|
|
batch_size = x.shape[0]
|
|
query = self.to_q(x)
|
|
key = self.to_k(x)
|
|
value = self.to_v(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
|
|
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)
|