| |
| |
|
|
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import math |
|
|
| from einops import rearrange |
| from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func |
| from huggingface_hub import PyTorchModelHubMixin |
|
|
| from . import rotary |
| from .fused_add_dropout_scale import ( |
| bias_dropout_add_scale_fused_train, |
| bias_dropout_add_scale_fused_inference, |
| modulate_fused, |
| ) |
|
|
| from .dit import LayerNorm, TimestepEmbedder |
|
|
|
|
| |
| |
| |
|
|
|
|
| class DiTBlock(nn.Module): |
|
|
| def __init__(self, dim, n_heads, cond_dim, mlp_ratio=4, dropout=0.1): |
| super().__init__() |
| self.n_heads = n_heads |
|
|
| self.norm1 = LayerNorm(dim) |
| self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False) |
| self.attn_out = nn.Linear(dim, dim, bias=False) |
| self.dropout1 = nn.Dropout(dropout) |
|
|
| self.norm2 = LayerNorm(dim) |
| self.mlp = nn.Sequential( |
| nn.Linear(dim, mlp_ratio * dim, bias=True), |
| nn.GELU(approximate="tanh"), |
| nn.Linear(mlp_ratio * dim, dim, bias=True) |
| ) |
| self.dropout2 = nn.Dropout(dropout) |
|
|
| self.dropout = dropout |
| |
|
|
| self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim, bias=True) |
| self.adaLN_modulation.weight.data.zero_() |
| self.adaLN_modulation.bias.data.zero_() |
|
|
|
|
| def _get_bias_dropout_scale(self): |
| return ( |
| bias_dropout_add_scale_fused_train |
| if self.training |
| else bias_dropout_add_scale_fused_inference |
| ) |
|
|
|
|
| def forward(self, x, rotary_cos_sin, c, seqlens=None): |
| batch_size, seq_len = x.shape[0], x.shape[1] |
|
|
| bias_dropout_scale_fn = self._get_bias_dropout_scale() |
|
|
| shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c)[:, None].chunk(6, dim=2) |
|
|
| |
| x_skip = x |
| x = modulate_fused(self.norm1(x), shift_msa, scale_msa) |
| |
|
|
| qkv = self.attn_qkv(x) |
| qkv = rearrange(qkv, 'b s (three h d) -> b s three h d', three=3, h=self.n_heads) |
| with torch.cuda.amp.autocast(enabled=False): |
| cos, sin = rotary_cos_sin |
| qkv = rotary.apply_rotary_pos_emb_flash( |
| qkv, cos.to(qkv.dtype), sin.to(qkv.dtype) |
| ) |
| qkv = rearrange(qkv, 'b s ... -> (b s) ...') |
| if seqlens is None: |
| cu_seqlens = torch.arange( |
| 0, (batch_size + 1) * seq_len, step=seq_len, |
| dtype=torch.int32, device=qkv.device |
| ) |
| else: |
| cu_seqlens = seqlens.cumsum(-1) |
|
|
| x = flash_attn_varlen_qkvpacked_func( |
| qkv, cu_seqlens, seq_len, 0., causal=False) |
| |
| x = rearrange(x, '(b s) h d -> b s (h d)', b=batch_size) |
|
|
| x = bias_dropout_scale_fn(self.attn_out(x), None, gate_msa, x_skip, self.dropout) |
|
|
| |
| x = bias_dropout_scale_fn(self.mlp(modulate_fused(self.norm2(x), shift_mlp, scale_mlp)), None, gate_mlp, x, self.dropout) |
| return x |
|
|
|
|
|
|
| class EmbeddingLayer(nn.Module): |
| def __init__(self, dim, vocab_dim): |
| """ |
| Mode arg: 0 -> use a learned layer, 1 -> use eigenvectors, |
| 2-> add in eigenvectors, 3 -> use pretrained embedding matrix |
| """ |
| super().__init__() |
| |
| self.embedding = nn.Linear(vocab_dim, dim) |
| torch.nn.init.kaiming_uniform_(self.embedding.weight, a=math.sqrt(5)) |
|
|
| def forward(self, x): |
| |
| return self.embedding(x) |
|
|
|
|
| class DiTFinalLayer(nn.Module): |
| def __init__(self, hidden_size, out_channels, cond_dim): |
| super().__init__() |
| self.norm_final = LayerNorm(hidden_size) |
| self.linear = nn.Linear(hidden_size, out_channels) |
| self.linear.weight.data.zero_() |
| self.linear.bias.data.zero_() |
|
|
| self.adaLN_modulation = nn.Linear(cond_dim, 2 * hidden_size, bias=True) |
| self.adaLN_modulation.weight.data.zero_() |
| self.adaLN_modulation.bias.data.zero_() |
|
|
|
|
| def forward(self, x, c): |
| shift, scale = self.adaLN_modulation(c)[:, None].chunk(2, dim=2) |
| x = modulate_fused(self.norm_final(x), shift, scale) |
| x = self.linear(x) |
| return x |
|
|
|
|
| class DiT(nn.Module, PyTorchModelHubMixin): |
| def __init__( |
| self, |
| input_dim: int, |
| output_dim: int, |
| hidden_size: int, |
| n_heads: int, |
| cond_dim: int, |
| dropout: float, |
| n_blocks: int, |
| **kwargs |
| ): |
| super().__init__() |
|
|
| self.input_dim = input_dim |
| self.output_dim = output_dim |
|
|
| self.vocab_embed = EmbeddingLayer(hidden_size, input_dim) |
| self.sigma_map = TimestepEmbedder(cond_dim) |
| self.rotary_emb = rotary.Rotary(hidden_size // n_heads) |
|
|
| self.blocks = nn.ModuleList([ |
| DiTBlock(hidden_size, n_heads, cond_dim, dropout=dropout) |
| for _ in range(n_blocks) |
| ]) |
|
|
| self.output_layer = DiTFinalLayer(hidden_size, output_dim, cond_dim) |
| self.kwargs = kwargs |
|
|
| def _get_bias_dropout_scale(self): |
| return ( |
| bias_dropout_add_scale_fused_train |
| if self.training |
| else bias_dropout_add_scale_fused_inference |
| ) |
|
|
| def forward(self, input_vector, t): |
| x = self.vocab_embed(input_vector) |
| c = F.silu(self.sigma_map(t)) |
|
|
| rotary_cos_sin = self.rotary_emb(x) |
|
|
| with torch.cuda.amp.autocast(dtype=torch.bfloat16): |
| for i in range(len(self.blocks)): |
| x = self.blocks[i](x, rotary_cos_sin, c, seqlens=None) |
| x = self.output_layer(x, c) |
|
|
| return x |
