| |
|
|
| |
|
|
| from __future__ import annotations |
|
|
| from typing import TYPE_CHECKING |
|
|
| import ttnn |
|
|
| from ..layers.feedforward import ParallelFeedForward |
| from ..layers.linear import ColParallelLinear, prepare_chunked_linear_output |
| from ..layers.module import Module |
| from ..layers.normalization import DistributedLayerNorm |
| from ..utils.substate import rename_substate |
| from .attention import Attention |
|
|
| if TYPE_CHECKING: |
| import torch |
|
|
| from ..parallel.config import DiTParallelConfig |
| from ..parallel.manager import CCLManager |
| from ..utils.padding import PaddingConfig |
|
|
|
|
| |
| class TransformerBlock(Module): |
| def __init__( |
| self, |
| *, |
| dim: int, |
| modulation_dim: int | None = None, |
| num_heads: int, |
| head_dim: int, |
| context_pre_only: bool, |
| add_attention_to_output: bool = True, |
| context_head_scaling: bool = False, |
| ff_activation_fn: str = "gelu", |
| mesh_device: ttnn.MeshDevice, |
| ccl_manager: CCLManager | None, |
| parallel_config: DiTParallelConfig, |
| padding_config: PaddingConfig | None, |
| attention_k_chunk_size: int = 512, |
| attention_q_chunk_size: int = 128, |
| is_fsdp: bool = False, |
| ) -> None: |
| super().__init__() |
|
|
| if modulation_dim is None: |
| modulation_dim = dim |
|
|
| self.dim = dim |
| self.num_heads = num_heads |
| self.head_dim = head_dim |
| self.add_attention_to_output = add_attention_to_output |
|
|
| self.context_pre_only = context_pre_only |
| self.mesh_device = mesh_device |
| self.ccl_manager = ccl_manager |
| self.parallel_config = parallel_config |
|
|
| |
| fsdp_mesh_axis = parallel_config.sequence_parallel.mesh_axis if is_fsdp else None |
|
|
| self.norm1_linear = ColParallelLinear( |
| modulation_dim, |
| 6 * dim, |
| bias=True, |
| mesh_device=mesh_device, |
| mesh_axis=parallel_config.tensor_parallel.mesh_axis, |
| fsdp_mesh_axis=fsdp_mesh_axis, |
| ccl_manager=ccl_manager, |
| ) |
|
|
| self.norm1_norm = DistributedLayerNorm( |
| dim, |
| norm_eps=1e-6, |
| norm_elementwise_affine=False, |
| bias=False, |
| mesh_axis=parallel_config.tensor_parallel.mesh_axis, |
| mesh_device=mesh_device, |
| ccl_manager=ccl_manager, |
| ) |
|
|
| context_norm_dim = 6 * dim if not context_pre_only else 2 * dim |
| self.norm1_context_linear = ColParallelLinear( |
| modulation_dim, |
| context_norm_dim, |
| bias=True, |
| mesh_device=mesh_device, |
| mesh_axis=parallel_config.tensor_parallel.mesh_axis, |
| fsdp_mesh_axis=fsdp_mesh_axis, |
| ccl_manager=ccl_manager, |
| ) |
| self.norm1_context_norm = DistributedLayerNorm( |
| dim, |
| norm_eps=1e-6, |
| norm_elementwise_affine=False, |
| bias=False, |
| mesh_axis=parallel_config.tensor_parallel.mesh_axis, |
| mesh_device=mesh_device, |
| ccl_manager=ccl_manager, |
| ) |
|
|
| self.attn = Attention( |
| query_dim=dim, |
| head_dim=head_dim, |
| heads=num_heads, |
| out_dim=dim, |
| added_kv_proj_dim=dim, |
| context_pre_only=context_pre_only, |
| context_head_scaling=context_head_scaling, |
| eps=1e-6, |
| mesh_device=mesh_device, |
| ccl_manager=ccl_manager, |
| parallel_config=parallel_config, |
| padding_config=padding_config, |
| k_chunk_size=attention_k_chunk_size, |
| q_chunk_size=attention_q_chunk_size, |
| is_fsdp=is_fsdp, |
| ) |
|
|
| self.norm2 = DistributedLayerNorm( |
| dim, |
| norm_eps=1e-6, |
| norm_elementwise_affine=False, |
| bias=False, |
| mesh_axis=parallel_config.tensor_parallel.mesh_axis, |
| mesh_device=mesh_device, |
| ccl_manager=ccl_manager, |
| ) |
|
|
| self.ff = ParallelFeedForward( |
| dim=dim, |
| dim_out=dim, |
| activation_fn=ff_activation_fn, |
| mesh_device=mesh_device, |
| mesh_axis=parallel_config.tensor_parallel.mesh_axis, |
| fsdp_mesh_axis=fsdp_mesh_axis, |
| ccl_manager=ccl_manager, |
| ) |
|
|
| self.norm2_context = None |
| self.ff_context = None |
|
|
| if not context_pre_only: |
| self.norm2_context = DistributedLayerNorm( |
| dim, |
| norm_eps=1e-6, |
| norm_elementwise_affine=False, |
| bias=False, |
| mesh_axis=parallel_config.tensor_parallel.mesh_axis, |
| mesh_device=mesh_device, |
| ccl_manager=ccl_manager, |
| ) |
| self.ff_context = ParallelFeedForward( |
| dim=dim, |
| dim_out=dim, |
| activation_fn=ff_activation_fn, |
| mesh_device=mesh_device, |
| mesh_axis=parallel_config.tensor_parallel.mesh_axis, |
| fsdp_mesh_axis=fsdp_mesh_axis, |
| ccl_manager=ccl_manager, |
| ) |
|
|
| device_grid = self.mesh_device.compute_with_storage_grid_size() |
| self.core_grid = ttnn.CoreGrid(x=device_grid.x, y=device_grid.y) |
|
|
| def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None: |
| rename_substate(state, "norm1.linear", "norm1_linear") |
| rename_substate(state, "norm1.norm", "norm1_norm") |
| rename_substate(state, "norm1_context.linear", "norm1_context_linear") |
| rename_substate(state, "norm1_context.norm", "norm1_context_norm") |
| rename_substate(state, "ff.net.0.proj", "ff.ff1") |
| rename_substate(state, "ff.net.2", "ff.ff2") |
| rename_substate(state, "ff_context.net.0.proj", "ff_context.ff1") |
| rename_substate(state, "ff_context.net.2", "ff_context.ff2") |
|
|
| prepare_chunked_linear_output( |
| state, |
| prefix="norm1_linear", |
| device_count=self.parallel_config.tensor_parallel.factor, |
| chunks=6, |
| ) |
| prepare_chunked_linear_output( |
| state, |
| prefix="norm1_context_linear", |
| device_count=self.parallel_config.tensor_parallel.factor, |
| chunks=2 if self.context_pre_only else 6, |
| ) |
|
|
| def forward( |
| self, |
| spatial: ttnn.Tensor, |
| prompt: ttnn.Tensor, |
| time_embed: ttnn.Tensor, |
| spatial_sequence_length: int, |
| *, |
| spatial_rope: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, |
| prompt_rope: tuple[ttnn.Tensor, ttnn.Tensor] | None = None, |
| skip_time_embed_activation_fn: bool = False, |
| ) -> tuple[ttnn.Tensor, ttnn.Tensor | None]: |
| """Run the model forward. |
| |
| Args: |
| spatial: Tensor with shape [batch_size, spatial_sequence_length / sp_factor, query_dim / tp_factor]. |
| prompt: Tensor with shape [batch_size, prompt_sequence_length, query_dim / tp_factor] (sequence is not sharded!). |
| time_embed: Tensor with shape [batch_size, 1, query_dim]. |
| spatial_rope: Tuple of two tensors with shape [spatial_sequence_length / sp_factor, head_dim]. |
| prompt_rope: Tuple of two tensors with shape [prompt_sequence_length, head_dim] (sequence is not sharded!). |
| """ |
| assert len(spatial.shape) == 3 |
| assert len(prompt.shape) == 3 |
|
|
| tp_axis = self.parallel_config.tensor_parallel.mesh_axis |
|
|
| if not skip_time_embed_activation_fn: |
| time_embed = ttnn.silu(time_embed, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
|
|
| spatial_time = self.norm1_linear(time_embed) |
| prompt_time = self.norm1_context_linear(time_embed) |
|
|
| ( |
| spatial_shift_attn, |
| spatial_scale_attn, |
| spatial_gate_attn, |
| spatial_shift_ff, |
| spatial_scale_ff, |
| spatial_gate_ff, |
| ) = _chunk_time3d(spatial_time, 6) |
|
|
| spatial_normed = ttnn.squeeze( |
| self.norm1_norm( |
| ttnn.unsqueeze(spatial, 0), |
| dynamic_weight=(1 + spatial_scale_attn), |
| dynamic_bias=spatial_shift_attn, |
| ), |
| 0, |
| ) |
|
|
| if self.context_pre_only: |
| prompt_scale_attn, prompt_shift_attn = _chunk_time3d(prompt_time, 2) |
| prompt_gate_attn = None |
| prompt_shift_ff = None |
| prompt_scale_ff = None |
| prompt_gate_ff = None |
| else: |
| ( |
| prompt_shift_attn, |
| prompt_scale_attn, |
| prompt_gate_attn, |
| prompt_shift_ff, |
| prompt_scale_ff, |
| prompt_gate_ff, |
| ) = _chunk_time3d(prompt_time, 6) |
|
|
| prompt_normed = ttnn.squeeze( |
| self.norm1_context_norm( |
| ttnn.unsqueeze(prompt, 0), |
| dynamic_weight=(1 + prompt_scale_attn), |
| dynamic_bias=prompt_shift_attn, |
| ), |
| 0, |
| ) |
|
|
| |
| spatial_normed = self.ccl_manager.all_gather_persistent_buffer( |
| spatial_normed, dim=2, mesh_axis=tp_axis, use_hyperparams=True |
| ) |
| prompt_normed = self.ccl_manager.all_gather_persistent_buffer( |
| prompt_normed, dim=2, mesh_axis=tp_axis, use_hyperparams=True |
| ) |
|
|
| spatial_attn, prompt_attn = self.attn.forward( |
| spatial=spatial_normed, |
| prompt=prompt_normed, |
| spatial_rope=spatial_rope, |
| prompt_rope=prompt_rope, |
| spatial_sequence_length=spatial_sequence_length, |
| ) |
| spatial_attn = spatial_attn * spatial_gate_attn |
| prompt_attn = prompt_attn * prompt_gate_attn if prompt_gate_attn is not None else None |
|
|
| spatial_plus_attn = spatial + spatial_attn |
| if self.add_attention_to_output: |
| spatial = spatial_plus_attn |
|
|
| spatial_normed = ttnn.squeeze( |
| self.norm2( |
| ttnn.unsqueeze(spatial_plus_attn, 0), |
| dynamic_weight=(1 + spatial_scale_ff), |
| dynamic_bias=spatial_shift_ff, |
| ), |
| 0, |
| ) |
|
|
| spatial_normed = self.ccl_manager.all_gather_persistent_buffer( |
| spatial_normed, dim=2, mesh_axis=tp_axis, use_hyperparams=True |
| ) |
|
|
| spatial_ff = ttnn.squeeze(self.ff(ttnn.unsqueeze(spatial_normed, 0)), 0) |
| spatial_ff = spatial_ff * spatial_gate_ff |
|
|
| spatial = spatial + spatial_ff |
|
|
| if self.context_pre_only: |
| return spatial, None |
|
|
| prompt_plus_attn = prompt + prompt_attn |
| if self.add_attention_to_output: |
| prompt = prompt_plus_attn |
|
|
| prompt_normed = ttnn.squeeze( |
| self.norm2_context( |
| ttnn.unsqueeze(prompt_plus_attn, 0), |
| dynamic_weight=(1 + prompt_scale_ff), |
| dynamic_bias=prompt_shift_ff, |
| ), |
| 0, |
| ) |
|
|
| prompt_normed = self.ccl_manager.all_gather_persistent_buffer( |
| prompt_normed, dim=2, mesh_axis=tp_axis, use_hyperparams=True |
| ) |
|
|
| prompt_ff = ttnn.squeeze(self.ff_context(ttnn.unsqueeze(prompt_normed, 0)), 0) |
| prompt_ff = prompt_ff * prompt_gate_ff |
|
|
| prompt = prompt + prompt_ff |
|
|
| return spatial, prompt |
|
|
|
|
| def _chunk_time3d(t: ttnn.Tensor, count: int) -> list[ttnn.Tensor]: |
| size = t.shape[-1] // count |
| return [t[:, :, i * size : (i + 1) * size] for i in range(count)] |
|
|
