update test (#23)

* shard ema among devices and modify checkpointing

* shard ema among devices and modify checkpointing

* rewrite sharding to support padding.

* rewrite sharding to support padding.

* add test for ema sharding

* add test for ema sharding

* Delete test.sh

* update ema sharding in both scripts

* update checkpoint

* update test for ema sharding

* update scripts

* fix tests

* update test

---------

Co-authored-by: Maruyama_Aya <[email protected]>
Co-authored-by: Ziming Liu <[email protected]>

Author: oahzxl

tests/test_checkpoint.py

@@ -8,17 +8,16 @@
 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, rerun_if_address_is_in_use, spawn
+from colossalai.testing import check_state_dict_equal, rerun_if_address_is_in_use, spawn
 from colossalai.zero import LowLevelZeroOptimizer
 
-from opendit.models.dit import DiT_S_2
+from opendit.models.dit import DiT
 
 
-@clear_cache_before_run()
 def run_zero_checkpoint(stage: int, shard: bool, offload: bool):
     plugin = LowLevelZeroPlugin(precision="fp16", stage=stage, max_norm=1.0, initial_scale=32, cpu_offload=offload)
     booster = Booster(plugin=plugin)
-    model = DiT_S_2().half()
+    model = DiT(depth=2, hidden_size=64, patch_size=2, num_heads=4).half()
     criterion = lambda x: x.mean()
     optimizer = HybridAdam((model.parameters()), lr=0.001)
     model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion)
@@ -40,13 +39,12 @@ def run_zero_checkpoint(stage: int, shard: bool, offload: bool):
 
     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().half()
+    new_model = DiT(depth=2, hidden_size=64, patch_size=2, num_heads=4).half()
     new_optimizer = HybridAdam((new_model.parameters()), lr=0.001)
     new_model, new_optimizer, _, _, _ = booster.boost(new_model, new_optimizer)
 
@@ -67,6 +65,7 @@ def run_zero_checkpoint(stage: int, shard: bool, offload: bool):
 
     booster.load_optimizer(new_optimizer, optimizer_ckpt_path)
     check_state_dict_equal(optimizer.optim.state_dict(), new_optimizer.optim.state_dict(), False)
+    dist.barrier()
 
     if dist.get_rank() == 0:
         shutil.rmtree(tempdir)
@@ -76,7 +75,6 @@ def run_zero_checkpoint(stage: int, shard: bool, offload: bool):
 def run_dist(rank, world_size, port, stage: int, shard: bool, offload: bool):
     colossalai.launch(config=(dict()), rank=rank, world_size=world_size, port=port, host="localhost")
     run_zero_checkpoint(stage=stage, shard=shard, offload=offload)
-    torch.cuda.empty_cache()
 
 
 @pytest.mark.parametrize("stage", [2])

tests/test_ema_sharding.py

@@ -1,3 +1,4 @@
+import os
 from copy import deepcopy
 
 import colossalai
@@ -8,22 +9,31 @@
 from colossalai.nn.optimizer import HybridAdam
 from colossalai.testing import clear_cache_before_run, rerun_if_address_is_in_use, spawn
 
-from opendit.models.dit import DiT_S_2
+from opendit.models.dit import DiT
 from opendit.utils.ckpt_utils import model_gathering, record_model_param_shape
 from opendit.utils.operation import model_sharding
 from opendit.utils.train_utils import update_ema
 
 
+def assert_params_equal(model1, model2):
+    for (name1, param1), (name2, param2) in zip(model1.named_parameters(), model2.named_parameters()):
+        assert name1 == name2
+        if name1 == "pos_embed":
+            continue
+        assert torch.allclose(param1, param2)
+
+
 @clear_cache_before_run()
 def run_ema_sharding():
     plugin = LowLevelZeroPlugin(precision="fp16", stage=2, max_norm=1.0, initial_scale=32)
     booster = Booster(plugin=plugin)
-    model = DiT_S_2().cuda().half()
+    model = DiT(depth=2, hidden_size=64, patch_size=2, num_heads=4).cuda().half()
 
     ema_sharding = deepcopy(model).eval()
     model_param_shape = record_model_param_shape(ema_sharding)
     model_sharding(ema_sharding)
     ema_no_sharding = deepcopy(model).eval()
+    ema_to_read = deepcopy(model).eval()
 
     criterion = lambda x: x.mean()
     optimizer = HybridAdam((model.parameters()), lr=0.001)
@@ -44,25 +54,21 @@ def run_ema_sharding():
     gather_ema_sharding = deepcopy(ema_sharding)
     model_gathering(gather_ema_sharding, model_param_shape)
     if dist.get_rank() == 0:
-        for (gather_ema_sharding_name, gather_ema_sharding_param), (ema_no_sharding_name, ema_no_sharding_param) in zip(
-            gather_ema_sharding.named_parameters(), ema_no_sharding.named_parameters()
-        ):
-            assert gather_ema_sharding_name == ema_no_sharding_name
-            if gather_ema_sharding_name == "pos_embed":
-                continue
-            assert torch.allclose(gather_ema_sharding_param, ema_no_sharding_param)
+        assert_params_equal(gather_ema_sharding, ema_no_sharding)
+    dist.barrier()
+
+    # should be same after read again
+    if dist.get_rank() == 0:
+        torch.save(gather_ema_sharding.state_dict(), "tmp.pth")
+        ema_to_read.load_state_dict(torch.load("tmp.pth"))
+        assert_params_equal(gather_ema_sharding, ema_to_read)
+        os.remove("tmp.pth")
     dist.barrier()
 
     # should be same after sharding again
     if dist.get_rank() == 0:
         model_sharding(gather_ema_sharding)
-        for (gather_ema_sharding_name, gather_ema_sharding_param), (ema_sharding_name, ema_sharding_param) in zip(
-            gather_ema_sharding.named_parameters(), ema_sharding.named_parameters()
-        ):
-            assert gather_ema_sharding_name == ema_sharding_name
-            if gather_ema_sharding_name == "pos_embed":
-                continue
-            assert torch.allclose(gather_ema_sharding_param, ema_sharding_param)
+        assert_params_equal(gather_ema_sharding, ema_sharding)
     dist.barrier()
 
 

tests/test_flash_attention.py

@@ -1,6 +1,8 @@
 import copy
 
 import colossalai
+import flash_attn
+import pytest
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -120,7 +122,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
-def flash_attn(seq_len, hidden_dim, head_num, batch_size):
+def _run_flash_attn(seq_len, hidden_dim, head_num, batch_size):
     seq_len = seq_len
     hidden_dim = hidden_dim
     head_num = head_num
@@ -231,14 +233,15 @@ def flash_attn(seq_len, hidden_dim, head_num, batch_size):
 @parameterize("head_num", [16])
 @parameterize("batch_size", [2])
 def run_flash_attn(seq_len, hidden_dim, head_num, batch_size):
-    flash_attn(seq_len, hidden_dim, head_num, batch_size)
+    _run_flash_attn(seq_len, hidden_dim, head_num, batch_size)
 
 
 def check_all2all_attn(rank, world_size, port):
     colossalai.launch(config={}, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
     run_flash_attn()
 
 
+@pytest.mark.skipif(flash_attn.__version__ < "2.4.1", reason="requires flashattn 2.4.1 or higher")
 @rerun_if_address_is_in_use()
 def test_flash_attn():
     spawn(check_all2all_attn, nprocs=WORKERS)

tests/test_sequence_parallel.py

@@ -13,7 +13,7 @@
 
 from opendit.utils.operation import all_to_all_comm
 
-WORKERS = 4
+WORKERS = 2
 
 
 class DistAttention(nn.Module):