394 lines
13 KiB
Python
394 lines
13 KiB
Python
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from cache_dit import ForwardPattern
|
|
from transformers import LlamaConfig, LlamaModel
|
|
from transformers.modeling_outputs import BaseModelOutputWithPast
|
|
from transformers.models.llama.modeling_llama import LlamaDecoderLayer
|
|
|
|
from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
|
|
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
|
|
|
|
|
|
class SinusoidalPosEmb(nn.Module):
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.dim = dim
|
|
|
|
def forward(self, x):
|
|
device = x.device
|
|
dtype = x.dtype
|
|
half_dim = self.dim // 2
|
|
emb = math.log(10000) / (half_dim - 1)
|
|
emb = torch.exp(torch.arange(half_dim, device=device, dtype=dtype) * -emb)
|
|
emb = x[:, None] * emb[None, :]
|
|
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
|
|
return emb
|
|
|
|
|
|
class LlamaAdaptiveRMSNorm(nn.Module):
|
|
def __init__(self, hidden_size=1024, eps=1e-6, dim_cond=1024):
|
|
super().__init__()
|
|
self.to_weight = nn.Linear(dim_cond, hidden_size)
|
|
nn.init.zeros_(self.to_weight.weight)
|
|
nn.init.ones_(self.to_weight.bias)
|
|
self.variance_epsilon = eps
|
|
self._is_hf_initialized = True
|
|
|
|
def forward(self, hidden_states, cond_embedding):
|
|
input_dtype = hidden_states.dtype
|
|
variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
|
|
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
|
|
|
weight = self.to_weight(cond_embedding)
|
|
if len(weight.shape) == 2:
|
|
weight = weight.unsqueeze(1)
|
|
|
|
return (weight * hidden_states).to(input_dtype)
|
|
|
|
|
|
class LlamaNARDecoderLayer(LlamaDecoderLayer):
|
|
def __init__(self, config: LlamaConfig, layer_idx: int):
|
|
super().__init__(config, layer_idx)
|
|
|
|
self.input_layernorm = LlamaAdaptiveRMSNorm(
|
|
config.hidden_size,
|
|
eps=config.rms_norm_eps,
|
|
dim_cond=config.hidden_size,
|
|
)
|
|
self.post_attention_layernorm = LlamaAdaptiveRMSNorm(
|
|
config.hidden_size,
|
|
eps=config.rms_norm_eps,
|
|
dim_cond=config.hidden_size,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
cond_embedding: torch.Tensor,
|
|
attention_mask: torch.Tensor | None = None,
|
|
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
output_attentions: bool | None = False,
|
|
cache_position: torch.LongTensor | None = None,
|
|
**kwargs,
|
|
) -> tuple[torch.FloatTensor, ...]:
|
|
residual = hidden_states
|
|
hidden_states = self.input_layernorm(hidden_states, cond_embedding=cond_embedding)
|
|
|
|
# RoPE is computed by the parent model and passed in as (cos, sin).
|
|
hidden_states, self_attn_weights = self.self_attn(
|
|
hidden_states=hidden_states,
|
|
position_embeddings=position_embeddings,
|
|
attention_mask=attention_mask,
|
|
cache_position=cache_position,
|
|
output_attentions=output_attentions,
|
|
**kwargs,
|
|
)
|
|
|
|
hidden_states = residual + hidden_states
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.post_attention_layernorm(
|
|
hidden_states,
|
|
cond_embedding=cond_embedding,
|
|
)
|
|
hidden_states = self.mlp(hidden_states)
|
|
hidden_states = residual + hidden_states
|
|
|
|
outputs = (hidden_states,)
|
|
|
|
if output_attentions:
|
|
outputs += (self_attn_weights,)
|
|
|
|
return outputs
|
|
|
|
|
|
class DiffLlama(LlamaModel):
|
|
# Regional torch.compile in DiffusionModelRunner targets these blocks.
|
|
_repeated_blocks = ["LlamaNARDecoderLayer"]
|
|
_layerwise_offload_blocks_attrs = ["layers"]
|
|
|
|
# cache-dit: Llama-style backbone with static condition embeddings injected
|
|
# before the decoder loop (like Stable Audio Open). Pattern_3 handles the
|
|
# cross-attention-style conditioning on the ``layers`` ModuleList.
|
|
_cache_dit_adapter_config = CacheDiTAdapterConfig(
|
|
block_forward_patterns={
|
|
"layers": ForwardPattern.Pattern_3,
|
|
},
|
|
)
|
|
|
|
# Sequence parallelism: shard mel sequence before the first decoder layer and
|
|
# gather after the output projection. Tensor parallelism still requires
|
|
# migrating attention to vLLM parallel linear layers.
|
|
_sp_plan = {
|
|
"layers.0": {
|
|
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
|
|
},
|
|
"mel_out_mlp": SequenceParallelOutput(gather_dim=1, expected_dims=3),
|
|
}
|
|
|
|
@staticmethod
|
|
def _is_transformer_block(name: str, module: nn.Module) -> bool:
|
|
del module
|
|
return name.startswith("layers.") and name.split(".")[-1].isdigit()
|
|
|
|
_hsdp_shard_conditions = [_is_transformer_block]
|
|
|
|
def __init__(
|
|
self,
|
|
mel_dim=100,
|
|
hidden_size=1024,
|
|
num_heads=16,
|
|
num_layers=16,
|
|
dropout=0.1,
|
|
ffn_dropout=0.1,
|
|
attention_dropout=0.0,
|
|
config: LlamaConfig | None = None,
|
|
):
|
|
if config is None:
|
|
config = LlamaConfig(
|
|
vocab_size=32000,
|
|
hidden_size=hidden_size,
|
|
intermediate_size=hidden_size * 4,
|
|
num_hidden_layers=num_layers,
|
|
num_attention_heads=num_heads,
|
|
num_key_value_heads=num_heads,
|
|
max_position_embeddings=4096,
|
|
rms_norm_eps=1e-6,
|
|
attention_dropout=attention_dropout,
|
|
use_cache=False,
|
|
attn_implementation="eager",
|
|
)
|
|
|
|
config.use_cache = False
|
|
|
|
super().__init__(config)
|
|
|
|
# Keep custom decoder layers on the same attention backend.
|
|
layer_config = LlamaConfig(
|
|
vocab_size=config.vocab_size,
|
|
hidden_size=hidden_size,
|
|
intermediate_size=hidden_size * 4,
|
|
num_hidden_layers=num_layers,
|
|
num_attention_heads=num_heads,
|
|
num_key_value_heads=num_heads,
|
|
max_position_embeddings=4096,
|
|
rms_norm_eps=config.rms_norm_eps,
|
|
attention_dropout=attention_dropout,
|
|
use_cache=False,
|
|
attn_implementation="eager",
|
|
bos_token_id=None,
|
|
eos_token_id=None,
|
|
)
|
|
|
|
self.layers = nn.ModuleList(
|
|
[
|
|
LlamaNARDecoderLayer(
|
|
layer_config,
|
|
layer_idx=i,
|
|
)
|
|
for i in range(num_layers)
|
|
]
|
|
)
|
|
|
|
self.norm = LlamaAdaptiveRMSNorm(hidden_size, dim_cond=hidden_size)
|
|
|
|
self.diff_step_embedding = SinusoidalPosEmb(hidden_size)
|
|
self.diff_step_mlp = nn.Sequential(
|
|
nn.Linear(hidden_size, hidden_size * 4),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_size * 4, hidden_size),
|
|
)
|
|
|
|
self.cond_mlp = nn.Sequential(
|
|
nn.Linear(hidden_size, hidden_size * 4),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_size * 4, hidden_size),
|
|
)
|
|
|
|
self.mel_mlp = nn.Sequential(
|
|
nn.Linear(mel_dim, hidden_size * 4),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_size * 4, hidden_size),
|
|
)
|
|
|
|
self.mel_out_mlp = nn.Sequential(
|
|
nn.Linear(hidden_size, hidden_size * 4),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_size * 4, mel_dim),
|
|
)
|
|
|
|
self.embed_tokens = None
|
|
|
|
self.post_init()
|
|
|
|
def _prepare_decoder_attention_mask(
|
|
self,
|
|
attention_mask,
|
|
input_shape,
|
|
inputs_embeds,
|
|
past_key_values_length=0,
|
|
):
|
|
"""
|
|
Build a 4D additive padding mask.
|
|
|
|
This model is non-autoregressive, so no causal mask is added here.
|
|
"""
|
|
|
|
combined_attention_mask = None
|
|
|
|
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None):
|
|
bsz, src_len = mask.size()
|
|
tgt_len = tgt_len if tgt_len is not None else src_len
|
|
|
|
expanded_mask = (
|
|
mask[:, None, None, :]
|
|
.expand(
|
|
bsz,
|
|
1,
|
|
tgt_len,
|
|
src_len,
|
|
)
|
|
.to(dtype)
|
|
)
|
|
|
|
inverted_mask = 1.0 - expanded_mask
|
|
|
|
return inverted_mask.masked_fill(
|
|
inverted_mask.to(torch.bool),
|
|
torch.finfo(dtype).min,
|
|
)
|
|
|
|
if attention_mask is not None:
|
|
expanded_attn_mask = _expand_mask(
|
|
attention_mask,
|
|
inputs_embeds.dtype,
|
|
tgt_len=input_shape[-1],
|
|
).to(inputs_embeds.device)
|
|
|
|
combined_attention_mask = (
|
|
expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
|
|
)
|
|
|
|
return combined_attention_mask
|
|
|
|
def forward(
|
|
self,
|
|
x,
|
|
diffusion_step,
|
|
cond,
|
|
x_mask,
|
|
input_ids: torch.LongTensor = None,
|
|
attention_mask: torch.Tensor | None = None,
|
|
position_ids: torch.LongTensor | None = None,
|
|
past_key_values: list[torch.FloatTensor] | None = None,
|
|
inputs_embeds: torch.FloatTensor | None = None,
|
|
use_cache: bool | None = None,
|
|
output_attentions: bool | None = None,
|
|
output_hidden_states: bool | None = None,
|
|
return_dict: bool | None = False,
|
|
) -> BaseModelOutputWithPast | torch.Tensor | dict:
|
|
batch_size, seq_length, _ = x.shape
|
|
|
|
cond_embedding = self.cond_mlp(cond)
|
|
x = self.mel_mlp(x)
|
|
|
|
diffusion_step = self.diff_step_embedding(diffusion_step).to(device=x.device, dtype=x.dtype)
|
|
diffusion_step = self.diff_step_mlp(diffusion_step)
|
|
|
|
x = x + cond_embedding
|
|
|
|
inputs_embeds = x
|
|
attention_mask = x_mask
|
|
|
|
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
|
output_hidden_states = (
|
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
|
)
|
|
|
|
# KV cache is disabled because this diffusion decoder processes full sequences.
|
|
|
|
if past_key_values is not None:
|
|
raise NotImplementedError("KV cache is disabled for this non-autoregressive DiffLlama implementation.")
|
|
|
|
if position_ids is None:
|
|
device = input_ids.device if input_ids is not None else inputs_embeds.device
|
|
position_ids = torch.arange(
|
|
0,
|
|
seq_length,
|
|
dtype=torch.long,
|
|
device=device,
|
|
)
|
|
position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
|
|
else:
|
|
position_ids = position_ids.view(batch_size, seq_length).long()
|
|
|
|
if attention_mask is None:
|
|
attention_mask = torch.ones(
|
|
(batch_size, seq_length),
|
|
dtype=torch.bool,
|
|
device=inputs_embeds.device,
|
|
)
|
|
|
|
attention_mask = self._prepare_decoder_attention_mask(
|
|
attention_mask,
|
|
(batch_size, seq_length),
|
|
inputs_embeds,
|
|
past_key_values_length=0,
|
|
)
|
|
|
|
hidden_states = inputs_embeds
|
|
|
|
all_hidden_states = () if output_hidden_states else None
|
|
all_self_attns = () if output_attentions else None
|
|
all_layer_hidden_states = []
|
|
|
|
cache_position = torch.arange(
|
|
0,
|
|
seq_length,
|
|
device=hidden_states.device,
|
|
dtype=torch.long,
|
|
)
|
|
|
|
# LlamaAttention expects precomputed rotary embeddings.
|
|
position_embeddings = self.rotary_emb(hidden_states, position_ids)
|
|
|
|
for idx, decoder_layer in enumerate(self.layers):
|
|
if output_hidden_states:
|
|
all_hidden_states += (hidden_states,)
|
|
|
|
if self.gradient_checkpointing and self.training:
|
|
raise NotImplementedError
|
|
|
|
layer_outputs = decoder_layer(
|
|
hidden_states,
|
|
cond_embedding=diffusion_step,
|
|
attention_mask=attention_mask,
|
|
position_embeddings=position_embeddings,
|
|
output_attentions=output_attentions,
|
|
cache_position=cache_position,
|
|
)
|
|
|
|
hidden_states = layer_outputs[0]
|
|
all_layer_hidden_states.append(hidden_states.clone())
|
|
|
|
if output_attentions:
|
|
all_self_attns += (layer_outputs[1],)
|
|
|
|
hidden_states = self.norm(hidden_states, cond_embedding=diffusion_step)
|
|
|
|
if output_hidden_states:
|
|
all_hidden_states += (hidden_states,)
|
|
|
|
hidden_states = self.mel_out_mlp(hidden_states)
|
|
|
|
if return_dict:
|
|
return {
|
|
"output": hidden_states,
|
|
"hidden_states": all_layer_hidden_states,
|
|
"attentions": all_self_attns,
|
|
}
|
|
|
|
return hidden_states
|