kandinsky/model/dit.py

import math

import torch
from torch import nn
import torch.nn.functional as F

from diffusers import CogVideoXDDIMScheduler

from .nn import TimeEmbeddings, TextEmbeddings, VisualEmbeddings, RoPE3D, Modulation, MultiheadSelfAttention, MultiheadSelfAttentionTP, FeedForward, OutLayer
from .utils import exist


from torch.distributed.tensor.parallel import (
    ColwiseParallel,
    PrepareModuleInput,
    PrepareModuleOutput,
    RowwiseParallel,
    SequenceParallel,
    parallelize_module,
)

from torch.distributed._tensor import Replicate, Shard

def parallelize(model, tp_mesh):
    if tp_mesh.size() > 1:

        plan = {
            "in_layer":ColwiseParallel(),
            "out_layer": RowwiseParallel(
                 output_layouts=Replicate(),
            )
        }
        parallelize_module(model.time_embeddings, tp_mesh, plan)

        plan = {
            "in_layer": ColwiseParallel(output_layouts=Replicate(),)
        }
        parallelize_module(model.text_embeddings, tp_mesh, plan)
        parallelize_module(model.visual_embeddings, tp_mesh, plan)
        
        for i, doubled_transformer_block in enumerate(model.transformer_blocks):
            for j, transformer_block in enumerate(doubled_transformer_block):
                transformer_block.self_attention = MultiheadSelfAttentionTP(transformer_block.self_attention)
                plan = {
                    #text modulation
                    "text_modulation":  PrepareModuleInput(
                        input_layouts=(None, None),
                        desired_input_layouts=(Replicate(), None),
                    ),
                    "text_modulation.out_layer": ColwiseParallel(output_layouts=Replicate(),),
                    #visual modulation
                    "visual_modulation":  PrepareModuleInput(
                        input_layouts=(None, None),
                        desired_input_layouts=(Replicate(), None), 
                    ),
                    "visual_modulation.out_layer": ColwiseParallel(output_layouts=Replicate(), use_local_output=True),

                    #self_attention_norm 
                    "self_attention_norm": SequenceParallel(sequence_dim=0, use_local_output=True), # TODO надо ли вообще это??? если у нас смешанный ввод нескольких видосом может быть

                    #self_attention
                    "self_attention.to_query":  ColwiseParallel(
                        input_layouts=Replicate(),
                    ),
                    "self_attention.to_key":  ColwiseParallel(
                        input_layouts=Replicate(),
                    ),
                    "self_attention.to_value":  ColwiseParallel(
                        input_layouts=Replicate(),
                    ),
                    
                    "self_attention.query_norm": SequenceParallel(sequence_dim=0, use_local_output=True),
                    "self_attention.key_norm": SequenceParallel(sequence_dim=0, use_local_output=True),

                    "self_attention.output_layer": RowwiseParallel( 
                        # input_layouts=(Shard(0), ),
                        output_layouts=Replicate(),
                    ),

                    #feed_forward_norm
                    "feed_forward_norm": SequenceParallel(sequence_dim=0, use_local_output=True),

                    #feed_forward
                    "feed_forward.in_layer": ColwiseParallel(),
                    "feed_forward.out_layer": RowwiseParallel(),
                }
                self_attn = transformer_block.self_attention
                self_attn.num_heads = self_attn.num_heads // tp_mesh.size()
                parallelize_module(transformer_block, tp_mesh, plan)

        plan = {
                "modulation_out":ColwiseParallel(output_layouts=Replicate(),),
                "out_layer": ColwiseParallel(output_layouts=Replicate(),),
            }
        parallelize_module(model.out_layer, tp_mesh, plan)
    
        plan={
            "time_embeddings": PrepareModuleInput(desired_input_layouts=Replicate(),),
            "text_embeddings": PrepareModuleInput(desired_input_layouts=Replicate(),),
            "visual_embeddings": PrepareModuleInput(desired_input_layouts=Replicate(),),
            "out_layer": PrepareModuleInput(
                input_layouts=(None, None, None, None),
                desired_input_layouts=(Replicate(), Replicate(), Replicate(), None)), 
        }
        parallelize_module(model, tp_mesh, {})
    return model

class TransformerBlock(nn.Module):

    def __init__(self, model_dim, time_dim, ff_dim, head_dim=64):
        super().__init__()
        self.visual_modulation = Modulation(time_dim, model_dim)
        self.text_modulation = Modulation(time_dim, model_dim)

        self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=True)
        self.self_attention = MultiheadSelfAttention(model_dim, head_dim)

        self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=True)
        self.feed_forward = FeedForward(model_dim, ff_dim)

    def forward(self, visual_embed, text_embed, time_embed, rope, visual_cu_seqlens, text_cu_seqlens, num_groups, attention_type):
        visual_shape = visual_embed.shape[:-1]
        visual_self_attn_params, visual_ff_params = self.visual_modulation(time_embed, visual_cu_seqlens)
        text_self_attn_params, text_ff_params = self.text_modulation(time_embed, text_cu_seqlens)
    
        visual_shift, visual_scale, visual_gate = torch.chunk(visual_self_attn_params, 3, dim=-1)
        text_shift, text_scale, text_gate = torch.chunk(text_self_attn_params, 3, dim=-1)
        visual_out = self.self_attention_norm(visual_embed) * (visual_scale[:, None, None] + 1.) + visual_shift[:, None, None]
        text_out = self.self_attention_norm(text_embed) * (text_scale + 1.) + text_shift
        visual_out, text_out = self.self_attention(visual_out, text_out, rope, visual_cu_seqlens, text_cu_seqlens, num_groups, attention_type)
        
        visual_embed = visual_embed + visual_gate[:, None, None] * visual_out
        text_embed = text_embed + text_gate * text_out
       
        visual_shift, visual_scale, visual_gate = torch.chunk(visual_ff_params, 3, dim=-1)
        visual_out = self.feed_forward_norm(visual_embed) * (visual_scale[:, None, None] + 1.) + visual_shift[:, None, None]
        visual_embed = visual_embed + visual_gate[:, None, None] * self.feed_forward(visual_out)

        text_shift, text_scale, text_gate = torch.chunk(text_ff_params, 3, dim=-1)
        text_out = self.feed_forward_norm(text_embed) * (text_scale + 1.) + text_shift
        text_embed = text_embed + text_gate * self.feed_forward(text_out)
        return visual_embed, text_embed


class DiffusionTransformer3D(nn.Module):

    def __init__(
        self,
        in_visual_dim=4,
        in_text_dim=2048,
        time_dim=512,
        out_visual_dim=4,
        patch_size=(1, 2, 2),
        model_dim=2048,
        ff_dim=5120,
        num_blocks=8,
        axes_dims=(16, 24, 24),
    ):
        super().__init__()
        head_dim = sum(axes_dims)
        self.in_visual_dim = in_visual_dim
        self.model_dim = model_dim
        self.num_blocks = num_blocks

        self.time_embeddings = TimeEmbeddings(model_dim, time_dim)
        self.text_embeddings = TextEmbeddings(in_text_dim, model_dim)
        self.visual_embeddings = VisualEmbeddings(in_visual_dim, model_dim, patch_size)
        self.rope_embeddings = RoPE3D(axes_dims)

        self.transformer_blocks = nn.ModuleList([
            nn.ModuleList([
                TransformerBlock(model_dim, time_dim, ff_dim, head_dim),
                TransformerBlock(model_dim, time_dim, ff_dim, head_dim),
            ]) for _ in range(num_blocks)
        ])

        self.out_layer = OutLayer(model_dim, time_dim, out_visual_dim, patch_size)

    def forward(self, x, text_embed, time, visual_cu_seqlens, text_cu_seqlens, num_groups=(1, 1, 1), scale_factor=(1., 1., 1.)):
        time_embed = self.time_embeddings(time)
        text_embed = self.text_embeddings(text_embed)
        visual_embed = self.visual_embeddings(x)
        rope = self.rope_embeddings(visual_embed, visual_cu_seqlens, scale_factor)
        
        for i, (local_attention, global_attention) in enumerate(self.transformer_blocks):
            visual_embed, text_embed = local_attention(
                visual_embed, text_embed, time_embed, rope, visual_cu_seqlens, text_cu_seqlens, num_groups, 'local'
            )
            visual_embed, text_embed = global_attention(
                visual_embed, text_embed, time_embed, rope, visual_cu_seqlens, text_cu_seqlens, num_groups, 'global'
            )

        return self.out_layer(visual_embed, text_embed, time_embed, visual_cu_seqlens)


def get_dit(conf):
    dit = DiffusionTransformer3D(**conf.params)
    state_dict = torch.load(conf.checkpoint_path, weights_only=True, map_location=torch.device('cpu'))
    dit.load_state_dict(state_dict, strict=False)
    return dit