resume img pipeline and add tests (#22)

* init

* Update dependencies and installation instructions

* rename video train

* update script

* Update train.sh to train_video.sh and add train_img.sh

* Added matplotlib and flash_attn to requirements.txt

* Remove data-path argument from train_img.sh

* update args

* update test

* Add test for new model implementation

* Add weight initialization function and update model comparison in test_model()

* add ckpt

* update clip test

* Commented out unnecessary flags in train_img.sh script

* fix bugs

* Add transformers library to requirements.txt

* update

Author: oahzxl

README.md

@@ -1,25 +1,46 @@
-## Usage
+# OpenDiT
+### Install ColossalAI
+```
+git clone https://github.com/hpcaitech/ColossalAI.git
+cd ColossalAI
+git checkout adae123df3badfb15d044bd416f0cf29f250bc86
+pip install -e .
+```
 
+### Install OpenDiT
 ```
 # Prerequisite
 cd OpenDiT
 pip install -e .
-
-# Train
-bash preprocess.sh
-bash train.sh
-
-# Infer
+```
+## Image Pipeline
+```
+# train
+bash train_img.sh
+# inference
 bash sample.sh
 ```
-
-## Install ColossalAI
+## Video Pipeline
 ```
-git clone https://github.com/hpcaitech/ColossalAI.git
-cd ColossalAI
-git checkout adae123df3badfb15d044bd416f0cf29f250bc86
-pip install -e .
+# train
+bash preprocess.sh
+bash train_video.sh
 ```
+## Install kernels to speed up
+```
+# triton for modulate kernel
+pip install triton
+
+# flash attention
+pip install flash-attn
+
+# apex layernorm
+git clone https://github.com/NVIDIA/apex.git
+cd apex
+git checkout 741bdf50825a97664db08574981962d66436d16a
+pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./ --global-option="--cuda_ext" --global-option="--cpp_ext"
+```
+
 
 ## Scalable Diffusion Models with Transformers (DiT)<br><sub>Official PyTorch Implementation</sub>
 

opendit/models/dit.py

@@ -23,12 +23,6 @@
 from opendit.models.clip import TextEmbedder
 from opendit.utils.operation import all_to_all_comm, gather_forward_split_backward
 
-ULYSSES = False
-FLASH_ATTN = False
-SP_SIZE = 2
-LAYERNORM_KERNEL = False
-MODULATE_KERNEL = False
-
 
 def get_layernorm(hidden_size: torch.Tensor, eps: float, affine: bool, use_kernel: bool):
     if use_kernel:
@@ -45,16 +39,17 @@ def get_layernorm(hidden_size: torch.Tensor, eps: float, affine: bool, use_kerne
 def modulate(norm_func, x, shift, scale, use_kernel=False):
     # Suppose x is (N, T, D), shift is (N, D), scale is (N, D)
     dtype = x.dtype
-    x = norm_func(x.to(torch.float32))
+    x, shift, scale = x.to(torch.float32), shift.to(torch.float32), scale.to(torch.float32)
+    x = norm_func(x)
     if use_kernel:
         try:
             from opendit.kernels.fused_modulate import fused_modulate
 
-            x = fused_modulate(x, scale.to(torch.float32), shift.to(torch.float32))
+            x = fused_modulate(x, scale, shift)
         except ImportError:
             raise RuntimeError("FusedModulate kernel not available. Please install triton.")
     else:
-        x = x * (scale.to(torch.float32).unsqueeze(1) + 1) + shift.to(torch.float32).unsqueeze(1)
+        x = x * (scale.unsqueeze(1) + 1) + shift.unsqueeze(1)
     x = x.to(dtype)
 
     return x
@@ -156,8 +151,8 @@ def __init__(
         attn_drop: float = 0.0,
         proj_drop: float = 0.0,
         norm_layer: nn.Module = nn.LayerNorm,
-        use_flash_attn: bool = False,
-        enable_sequence_parallelism: bool = False,
+        enable_flashattn: bool = False,
+        sequence_parallel_size: int = 1,
     ) -> None:
         super().__init__()
         assert dim % num_heads == 0, "dim should be divisible by num_heads"
@@ -172,16 +167,20 @@ def __init__(
         self.attn_drop = nn.Dropout(attn_drop)
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
-        self.use_flash_attn = use_flash_attn
-        self.enable_sequence_parallelism = enable_sequence_parallelism
+        self.enable_flashattn = enable_flashattn
+        # TODO: support sequence_parallel_size > 2
+        assert sequence_parallel_size in [1, 2], "sequence_parallel_size is only supported for 1 or 2"
+        self.sequence_parallel_size = sequence_parallel_size
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         B, N, C = x.shape
         qkv = self.qkv(x)  # (B, N, C), N here is N_total // SP_SIZE
         # Todo: Change num_heads in somewhere else for a better code style
-        num_heads = self.num_heads if not self.enable_sequence_parallelism else self.num_heads // SP_SIZE
+        num_heads = (
+            self.num_heads if self.sequence_parallel_size == 1 else self.num_heads // self.sequence_parallel_size
+        )
 
-        if self.enable_sequence_parallelism:
+        if self.sequence_parallel_size > 1:
             q, k, v = qkv.split(self.head_dim * self.num_heads, dim=-1)
             # q = q.reshape(1, -1, self.head_dim * self.num_heads)
             # k = k.reshape(1, -1, self.head_dim * self.num_heads)
@@ -191,9 +190,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             k = all_to_all_comm(k, None)
             v = all_to_all_comm(v, None)
 
-            q = q.reshape(B, N * SP_SIZE, num_heads, self.head_dim).permute(0, 2, 1, 3).contiguous()
-            k = k.reshape(B, N * SP_SIZE, num_heads, self.head_dim).permute(0, 2, 1, 3).contiguous()
-            v = v.reshape(B, N * SP_SIZE, num_heads, self.head_dim).permute(0, 2, 1, 3).contiguous()
+            q = q.reshape(B, N * self.sequence_parallel_size, num_heads, self.head_dim).permute(0, 2, 1, 3).contiguous()
+            k = k.reshape(B, N * self.sequence_parallel_size, num_heads, self.head_dim).permute(0, 2, 1, 3).contiguous()
+            v = v.reshape(B, N * self.sequence_parallel_size, num_heads, self.head_dim).permute(0, 2, 1, 3).contiguous()
 
         else:
             # Todo: chunked flash attention
@@ -204,7 +203,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             #         .permute(2, 3, 0, 1, 4)
             #         .reshape(3, B * num_heads, 1, N, self.head_dim)
             #     )
-            if self.use_flash_attn:
+            if self.enable_flashattn:
                 # [3, B, num_heads, N, head_dim] => [B, N, num_heads, head_dim] * 3
                 qkv = qkv.reshape(B, N, 3, num_heads, self.head_dim).permute(2, 0, 1, 3, 4)
             else:
@@ -213,7 +212,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             q, k, v = qkv.unbind(0)
         q, k = self.q_norm(q), self.k_norm(k)
 
-        if self.use_flash_attn:
+        if self.enable_flashattn:
             from flash_attn import flash_attn_func
 
             # Todo: chunked flash attention
@@ -258,10 +257,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x = attn @ v
 
         x_output_shape = (
-            (B, N, C) if not self.enable_sequence_parallelism else (B, N * SP_SIZE, num_heads * self.head_dim)
+            (B, N, C)
+            if self.sequence_parallel_size == 1
+            else (B, N * self.sequence_parallel_size, num_heads * self.head_dim)
         )
         x = x.transpose(1, 2).reshape(x_output_shape)
-        if self.enable_sequence_parallelism:
+        if self.sequence_parallel_size > 1:
             # Todo: Use all_to_all_single for x
             # x = x.reshape(1, -1, num_heads * self.head_dim)
             x = all_to_all_comm(x, None, scatter_dim=1, gather_dim=2)
@@ -281,33 +282,37 @@ def __init__(
         hidden_size,
         num_heads,
         mlp_ratio=4.0,
-        flash_attn=False,
-        sequence_parallel=False,
-        layernorm_kernel=False,
-        modulate_kernel=False,
+        enable_flashattn=False,
+        sequence_parallel_size=False,
+        enable_layernorm_kernel=False,
+        enable_modulate_kernel=False,
         **block_kwargs,
     ):
         super().__init__()
-        self.modulate_kernel = modulate_kernel
-        self.norm1 = get_layernorm(hidden_size, eps=1e-6, affine=False, use_kernel=layernorm_kernel)
+        self.enable_modulate_kernel = enable_modulate_kernel
+        self.norm1 = get_layernorm(hidden_size, eps=1e-6, affine=False, use_kernel=enable_layernorm_kernel)
         self.attn = DistAttention(
             hidden_size,
             num_heads=num_heads,
             qkv_bias=True,
-            use_flash_attn=flash_attn,
-            enable_sequence_parallelism=sequence_parallel,
+            enable_flashattn=enable_flashattn,
+            sequence_parallel_size=sequence_parallel_size,
             **block_kwargs,
         )
-        self.norm2 = get_layernorm(hidden_size, eps=1e-6, affine=False, use_kernel=layernorm_kernel)
+        self.norm2 = get_layernorm(hidden_size, eps=1e-6, affine=False, use_kernel=enable_layernorm_kernel)
         mlp_hidden_dim = int(hidden_size * mlp_ratio)
         approx_gelu = lambda: nn.GELU(approximate="tanh")
         self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
         self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
 
     def forward(self, x, c):
         shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
-        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1, x, shift_msa, scale_msa, self.modulate_kernel))
-        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2, x, shift_mlp, scale_mlp, self.modulate_kernel))
+        x = x + gate_msa.unsqueeze(1) * self.attn(
+            modulate(self.norm1, x, shift_msa, scale_msa, self.enable_modulate_kernel)
+        )
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(
+            modulate(self.norm2, x, shift_mlp, scale_mlp, self.enable_modulate_kernel)
+        )
         return x
 
 
@@ -347,17 +352,18 @@ def __init__(
         class_dropout_prob=0.1,
         num_classes=1000,
         learn_sigma=True,
-        flash_attn=FLASH_ATTN,
-        sequence_parallel=ULYSSES,
-        layernorm_kernel=LAYERNORM_KERNEL,
-        modulate_kernel=MODULATE_KERNEL,
+        enable_flashattn=False,
+        enable_layernorm_kernel=False,
+        enable_modulate_kernel=False,
+        sequence_parallel_size=1,
     ):
         super().__init__()
         self.learn_sigma = learn_sigma
         self.in_channels = in_channels
         self.out_channels = in_channels * 2 if learn_sigma else in_channels
         self.patch_size = patch_size
         self.num_heads = num_heads
+        self.sequence_parallel_size = sequence_parallel_size
 
         self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True)
         self.t_embedder = TimestepEmbedder(hidden_size)
@@ -378,10 +384,10 @@ def __init__(
                     hidden_size,
                     num_heads,
                     mlp_ratio=mlp_ratio,
-                    flash_attn=flash_attn,
-                    sequence_parallel=sequence_parallel,
-                    modulate_kernel=modulate_kernel,
-                    layernorm_kernel=layernorm_kernel,
+                    enable_flashattn=enable_flashattn,
+                    sequence_parallel_size=sequence_parallel_size,
+                    enable_modulate_kernel=enable_modulate_kernel,
+                    enable_layernorm_kernel=enable_layernorm_kernel,
                 )
                 for _ in range(depth)
             ]
@@ -471,16 +477,16 @@ def forward(self, x, t, y):
         c = t + y  # (N, D)
 
         # Chunk x on sequence dimension to sp group
-        if ULYSSES:
-            x = x.chunk(SP_SIZE, dim=1)[dist.get_rank()]
+        if self.sequence_parallel_size > 1:
+            x = x.chunk(self.sequence_parallel_size, dim=1)[dist.get_rank()]
 
         for block in self.blocks:
             if self.gradient_checkpointing:
                 x = torch.utils.checkpoint.checkpoint(self.create_custom_forward(block), x, c)
             else:
                 x = block(x, c)  # (N, T, D)
 
-        if ULYSSES:
+        if self.sequence_parallel_size > 1:
             x = gather_forward_split_backward(x, dim=1, process_group=None)
 
         x = self.final_layer(x, c)  # (N, T, patch_size ** 2 * out_channels)

requirements.txt

@@ -18,5 +18,11 @@ pytorch_lightning
 h5py
 gdown
 scikit-video
-flash_attn
-pyav
+pyav
+tensorboard
+timm
+matplotlib
+accelerate
+diffusers
+transformers
+flash_attn==2.0.5

tests/test_checkpoint.py

@@ -0,0 +1,100 @@
+import os
+import shutil
+
+import colossalai
+import torch
+import torch.distributed as dist
+from colossalai.booster import Booster
+from colossalai.booster.plugin import LowLevelZeroPlugin
+from colossalai.nn.optimizer import HybridAdam
+from colossalai.testing import (
+    check_state_dict_equal,
+    clear_cache_before_run,
+    parameterize,
+    rerun_if_address_is_in_use,
+    spawn,
+)
+from colossalai.zero import LowLevelZeroOptimizer
+
+from opendit.models.dit import DiT_S_2
+
+
+# stage 1 and 2 process the optimizer/mode the same way
+# only test 2 is fine
+@clear_cache_before_run()
+@parameterize("stage", [2])
+@parameterize("shard", [True, False])
+@parameterize("offload", [False, True])
+def _test_zero_checkpoint(stage: int, shard: bool, offload: bool):
+    plugin = LowLevelZeroPlugin(stage=stage, max_norm=1.0, initial_scale=32, cpu_offload=offload)
+    booster = Booster(plugin=plugin)
+    model = DiT_S_2()
+    criterion = lambda x: x.mean()
+    optimizer = HybridAdam((model.parameters()), lr=0.001)
+    model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion)
+
+    x = torch.randn(2, 4, 32, 32).cuda().requires_grad_(True)
+    y = torch.randint(0, 10, (2,)).cuda()
+    t = torch.randint(0, 10, (2,)).cuda()
+    output = model(x, y, t)
+    loss = criterion(output)
+    booster.backward(loss, optimizer)
+    optimizer.step()
+
+    tempdir = "./tempdir"
+    if dist.get_rank() == 0:
+        if os.path.exists(tempdir):
+            shutil.rmtree(tempdir)
+        os.makedirs(tempdir)
+    dist.barrier()
+
+    model_ckpt_path = f"{tempdir}/model"
+    optimizer_ckpt_path = f"{tempdir}/optimizer"
+    # lr scheduler is tested in test_torch_ddp_checkpoint_io.py and low level zero does not change it, we can skip it here
+    booster.save_model(model, model_ckpt_path, shard=shard)
+    booster.save_optimizer(optimizer, optimizer_ckpt_path, shard=shard)
+
+    dist.barrier()
+
+    new_model = DiT_S_2()
+    new_optimizer = HybridAdam((new_model.parameters()), lr=0.001)
+    new_model, new_optimizer, _, _, _ = booster.boost(new_model, new_optimizer)
+
+    booster.load_model(new_model, model_ckpt_path)
+    check_state_dict_equal(model.state_dict(), new_model.state_dict(), False)
+    # check master weight
+    assert isinstance(new_optimizer, LowLevelZeroOptimizer)
+    working_param_id_set = set(id(p) for p in new_model.parameters())
+    for p_id, master_param in new_optimizer._param_store.working_to_master_param.items():
+        assert p_id in working_param_id_set
+        working_param = new_optimizer._param_store.master_to_working_param[id(master_param)]
+        padding = new_optimizer._param_store.get_param_padding_size(working_param)
+        padded_param = torch.nn.functional.pad(working_param.data.view(-1), (0, padding))
+        working_shard = padded_param.chunk(dist.get_world_size())[dist.get_rank()]
+        assert torch.equal(
+            working_shard, master_param.data.view(-1).to(dtype=padded_param.dtype, device=padded_param.device)
+        )
+
+    booster.load_optimizer(new_optimizer, optimizer_ckpt_path)
+    check_state_dict_equal(optimizer.optim.state_dict(), new_optimizer.optim.state_dict(), False)
+    torch.cuda.empty_cache()
+
+    if dist.get_rank() == 0:
+        shutil.rmtree(tempdir)
+    dist.barrier()
+
+
+def run_dist(rank, world_size, port):
+    colossalai.launch(config=(dict()), rank=rank, world_size=world_size, port=port, host="localhost")
+    _test_zero_checkpoint()
+    torch.cuda.empty_cache()
+
+
+@rerun_if_address_is_in_use()
+@clear_cache_before_run()
+def test_zero_checkpoint():
+    spawn(run_dist, 2)
+
+
+if __name__ == "__main__":
+    test_zero_checkpoint()