Shard ema (#15)

* 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

---------

Co-authored-by: Xuanlei Zhao <[email protected]>

Author: MaruyamaAya

opendit/utils/ckpt_utils.py

@@ -1,5 +1,7 @@
+import functools
 import json
 import logging
+import operator
 import os
 from typing import Tuple
 
@@ -11,6 +13,8 @@
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import _LRScheduler
 
+from opendit.utils.operation import model_sharding
+
 
 def load_json(file_path: str):
     with open(file_path, "r") as f:
@@ -22,6 +26,29 @@ def save_json(data, file_path: str):
         json.dump(data, f, indent=4)
 
 
+def remove_padding(tensor: torch.Tensor, original_shape: Tuple) -> torch.Tensor:
+    return tensor[: functools.reduce(operator.mul, original_shape)]
+
+
+def model_gathering(model: torch.nn.Module, model_shape_dict: dict):
+    global_rank = dist.get_rank()
+    global_size = dist.get_world_size()
+    for name, param in model.named_parameters():
+        all_params = [torch.empty_like(param.data) for _ in range(global_size)]
+        dist.all_gather(all_params, param.data, group=dist.group.WORLD)
+        if int(global_rank) == 0:
+            all_params = torch.cat(all_params)
+            param.data = remove_padding(all_params, model_shape_dict[name]).view(model_shape_dict[name])
+    dist.barrier()
+
+
+def record_model_param_shape(model: torch.nn.Module) -> dict:
+    param_shape = {}
+    for name, param in model.named_parameters():
+        param_shape[name] = param.shape
+    return param_shape
+
+
 def save(
     booster: Booster,
     model: nn.Module,
@@ -33,13 +60,19 @@ def save(
     batch_size: int,
     coordinator: DistCoordinator,
     save_dir: str,
+    shape_dict: dict,
 ):
     save_dir = os.path.join(save_dir, f"epoch{epoch}-step{step}")
     os.makedirs(os.path.join(save_dir, "model"), exist_ok=True)
 
     booster.save_model(model, os.path.join(save_dir, "model"), shard=True)
     # ema is not boosted, so we don't need to use booster.save_model
-    torch.save(ema.state_dict(), os.path.join(save_dir, "ema.pt"))
+    model_gathering(ema, shape_dict)
+    global_rank = dist.get_rank()
+    if int(global_rank) == 0:
+        torch.save(ema.state_dict(), os.path.join(save_dir, "ema.pt"))
+        model_sharding(ema)
+
     booster.save_optimizer(optimizer, os.path.join(save_dir, "optimizer"), shard=True, size_per_shard=4096)
     if lr_scheduler is not None:
         booster.save_lr_scheduler(lr_scheduler, os.path.join(save_dir, "lr_scheduler"))

opendit/utils/operation.py

@@ -71,6 +71,20 @@ def backward(ctx, *grad_output):
         return (return_grad, None, None, None)
 
 
+
+def model_sharding(model: torch.nn.Module):
+    global_rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    for name, param in model.named_parameters():
+        padding_size = (world_size - param.numel() % world_size) % world_size
+        if padding_size > 0:
+            padding_param = torch.nn.functional.pad(param.data.view(-1), [0, padding_size])
+        else:
+            padding_param = param.data.view(-1)
+        splited_params = padding_param.split(padding_param.numel() // world_size)
+        splited_params = splited_params[global_rank]
+        param.data = splited_params
+
 def all_to_all_comm(input_, process_group=None, scatter_dim=2, gather_dim=1):
     return _AllToAll.apply(input_, process_group, scatter_dim, gather_dim)
 

opendit/utils/train_utils.py

@@ -0,0 +1,62 @@
+from collections import OrderedDict
+
+import torch
+import torch.distributed as dist
+
+
+def get_model_numel(model: torch.nn.Module) -> int:
+    return sum(p.numel() for p in model.parameters())
+
+
+def format_numel_str(numel: int) -> str:
+    B = 1024**3
+    M = 1024**2
+    K = 1024
+    if numel >= B:
+        return f"{numel / B:.2f} B"
+    elif numel >= M:
+        return f"{numel / M:.2f} M"
+    elif numel >= K:
+        return f"{numel / K:.2f} K"
+    else:
+        return f"{numel}"
+
+
+def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor:
+    dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM)
+    tensor.div_(dist.get_world_size())
+    return tensor
+
+
+@torch.no_grad()
+def update_ema(
+    ema_model: torch.nn.Module, model: torch.nn.Module, optimizer=None, decay: float = 0.9999, sharded: bool = True
+) -> None:
+    """
+    Step the EMA model towards the current model.
+    """
+    ema_params = OrderedDict(ema_model.named_parameters())
+    model_params = OrderedDict(model.named_parameters())
+
+    for name, param in model_params.items():
+        if name == "pos_embed":
+            continue
+        if not sharded:
+            param_data = param.data
+            ema_params[name].mul_(decay).add_(param_data, alpha=1 - decay)
+        else:
+            if param.data.dtype != torch.float32:
+                param_id = id(param)
+                master_param = optimizer._param_store.working_to_master_param[param_id]
+                param_data = master_param.data
+            else:
+                param_data = param.data
+            ema_params[name].mul_(decay).add_(param_data, alpha=1 - decay)
+
+
+def requires_grad(model: torch.nn.Module, flag: bool = True) -> None:
+    """
+    Set requires_grad flag for all parameters in a model.
+    """
+    for p in model.parameters():
+        p.requires_grad = flag

tests/test_checkpoint.py

@@ -2,33 +2,23 @@
 import shutil
 
 import colossalai
+import pytest
 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.testing import check_state_dict_equal, clear_cache_before_run, 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)
+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()
+    model = DiT_S_2().half()
     criterion = lambda x: x.mean()
     optimizer = HybridAdam((model.parameters()), lr=0.001)
     model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion)
@@ -56,7 +46,7 @@ def _test_zero_checkpoint(stage: int, shard: bool, offload: bool):
 
     dist.barrier()
 
-    new_model = DiT_S_2()
+    new_model = DiT_S_2().half()
     new_optimizer = HybridAdam((new_model.parameters()), lr=0.001)
     new_model, new_optimizer, _, _, _ = booster.boost(new_model, new_optimizer)
 
@@ -77,24 +67,25 @@ def _test_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)
-    torch.cuda.empty_cache()
 
     if dist.get_rank() == 0:
         shutil.rmtree(tempdir)
     dist.barrier()
 
 
-def run_dist(rank, world_size, port):
+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")
-    _test_zero_checkpoint()
+    run_zero_checkpoint(stage=stage, shard=shard, offload=offload)
     torch.cuda.empty_cache()
 
 
+@pytest.mark.parametrize("stage", [2])
+@pytest.mark.parametrize("shard", [True, False])
+@pytest.mark.parametrize("offload", [False, True])
 @rerun_if_address_is_in_use()
-@clear_cache_before_run()
-def test_zero_checkpoint():
-    spawn(run_dist, 2)
+def test_zero_checkpoint(stage, shard, offload):
+    spawn(run_dist, 2, stage=stage, shard=shard, offload=offload)
 
 
 if __name__ == "__main__":
-    test_zero_checkpoint()
+    test_zero_checkpoint(2, True, False)

tests/test_ema_sharding.py

@@ -0,0 +1,81 @@
+from copy import deepcopy
+
+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 clear_cache_before_run, rerun_if_address_is_in_use, spawn
+
+from opendit.models.dit import DiT_S_2
+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
+
+
+@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()
+
+    ema_sharding = deepcopy(model).eval()
+    model_param_shape = record_model_param_shape(ema_sharding)
+    model_sharding(ema_sharding)
+    ema_no_sharding = deepcopy(model).eval()
+
+    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()
+
+    update_ema(ema_sharding, model.module, optimizer=optimizer, sharded=True, decay=0.5)
+    update_ema(ema_no_sharding, model.module, optimizer=optimizer, sharded=False, decay=0.5)
+
+    # should be equal after update
+    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)
+    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)
+    dist.barrier()
+
+
+def run_dist(rank, world_size, port):
+    colossalai.launch(config=(dict()), rank=rank, world_size=world_size, port=port, host="localhost")
+    run_ema_sharding()
+    torch.cuda.empty_cache()
+
+
+@rerun_if_address_is_in_use()
+def test_ema_sharding():
+    spawn(run_dist, 2)
+
+
+if __name__ == "__main__":
+    test_ema_sharding()