[ddp] add save/load state dict for ColoDDP (#1127)

* add save/load state dict for ColoDDP

* add unit test

* refactor unit test folder

* polish unit test

* rename unit test

tests/test_ddp/test_ddp_ignore_params.py

@@ -0,0 +1,87 @@
+import pytest
+import colossalai
+import torch
+import torch.multiprocessing as mp
+from colossalai.testing import rerun_if_address_is_in_use
+from colossalai.utils.cuda import get_current_device
+from colossalai.utils import free_port
+from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.tensor import ChunkManager
+from functools import partial
+from colossalai.nn.parallel import ColoDDP, ColoDDPV2
+from colossalai.gemini.gemini_mgr import GeminiManager
+from typing import Callable
+import torch.distributed as dist
+import os
+import random
+import numpy as np
+
+
+def set_seed(seed):
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+
+def init_ddp(module: torch.nn.Module) -> ColoDDP:
+    return ColoDDP(module)
+
+
+def init_ddpv2(module: torch.nn.Module, use_chunk: bool = False) -> ColoDDPV2:
+    chunk_size = ChunkManager.search_chunk_size(module, 64, 2) if use_chunk else None
+    chunk_manager = ChunkManager(chunk_size)
+    gemini_manager = GeminiManager('cuda', chunk_manager)
+    return ColoDDPV2(module, gemini_manager)
+
+
+class Net(torch.nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.fc1 = torch.nn.Linear(3, 3, bias=False)
+        self.fc2 = torch.nn.Linear(3, 1, bias=False)
+
+    def forward(self, x):
+        return self.fc2(self.fc1(x))
+
+
+def run_fwd_bwd(ddp_cls: ColoDDP, init_ddp_func: Callable[[torch.nn.Module], ColoDDP]):
+    with ColoInitContext(device=get_current_device()):
+        model = Net().cuda()
+    w1 = model.fc1.weight
+    w2 = model.fc2.weight
+    ddp_cls.set_params_to_ignore([w2])
+    model = init_ddp_func(model)
+    x = torch.rand(2, 3, device=get_current_device())
+    logits = model(x)
+    loss = torch.sum(logits)
+    model.backward(loss)
+    w1_grads = [torch.empty_like(w1) for _ in range(dist.get_world_size())]
+    dist.all_gather(w1_grads, w1.grad)
+    assert torch.equal(w1_grads[0], w1_grads[1])
+    w2_grads = [torch.empty_like(w2) for _ in range(dist.get_world_size())]
+    dist.all_gather(w2_grads, w2.grad)
+    assert not torch.equal(w2_grads[0], w2_grads[1])
+
+
+def run_dist(rank, world_size, port):
+    colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    set_seed(dist.get_rank())
+    run_fwd_bwd(ColoDDP, init_ddp)
+    run_fwd_bwd(ColoDDPV2, partial(init_ddpv2, use_chunk=False))
+    run_fwd_bwd(ColoDDPV2, partial(init_ddpv2, use_chunk=True))
+
+
+@pytest.mark.dist
+@pytest.mark.parametrize('world_size', [2])
+@rerun_if_address_is_in_use()
+def test_ddp_ignore_params(world_size):
+    run_func = partial(run_dist, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+if __name__ == '__main__':
+    test_ddp_ignore_params(2)

tests/test_ddp/test_ddp_state_dict.py

@@ -0,0 +1,66 @@
+import pytest
+import colossalai
+import torch
+import torch.multiprocessing as mp
+from colossalai.testing import rerun_if_address_is_in_use
+from colossalai.utils.cuda import get_current_device
+from colossalai.utils import free_port
+from colossalai.utils.model.colo_init_context import ColoInitContext
+from colossalai.tensor import ChunkManager
+from functools import partial
+from tests.components_to_test.registry import non_distributed_component_funcs
+from colossalai.nn.parallel import ColoDDPV2, ColoDDP
+from colossalai.gemini.gemini_mgr import GeminiManager
+from typing import Callable
+from collections import OrderedDict
+
+
+def check_state_dict_equal(state_dict: OrderedDict, other_state_dict: OrderedDict):
+    for (k1, t1), (k2, t2) in zip(state_dict.items(), other_state_dict.items()):
+        assert k1 == k2
+        assert torch.allclose(t1, t2, atol=1e-3, rtol=1e-3)
+
+
+def init_ddp(module: torch.nn.Module) -> ColoDDP:
+    return ColoDDP(module)
+
+
+def init_ddpv2(module: torch.nn.Module, use_chunk: bool = False, use_zero: bool = False) -> ColoDDPV2:
+    chunk_size = ChunkManager.search_chunk_size(module, 64, 4) if use_chunk else None
+    chunk_manager = ChunkManager(chunk_size, enable_distributed_storage=use_zero)
+    gemini_manager = GeminiManager('cuda', chunk_manager)
+    return ColoDDPV2(module, gemini_manager)
+
+
+def run_state_dict(ddp_init_func: Callable[[torch.nn.Module], ColoDDP]):
+    get_components_func = non_distributed_component_funcs.get_callable('nested_model')
+    model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
+    torch_model = model_builder().cuda()
+    with ColoInitContext(device=get_current_device()):
+        model = model_builder()
+    model = ddp_init_func(model)
+    torch_state_dict = torch_model.state_dict()
+    model.load_state_dict(torch_state_dict)
+    state_dict = model.state_dict()
+    check_state_dict_equal(torch_state_dict, state_dict)
+
+
+def run_dist(rank, world_size, port):
+    colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    run_state_dict(init_ddp)
+    run_state_dict(partial(init_ddpv2, use_chunk=False, use_zero=False))
+    run_state_dict(partial(init_ddpv2, use_chunk=False, use_zero=True))
+    run_state_dict(partial(init_ddpv2, use_chunk=True, use_zero=False))
+    run_state_dict(partial(init_ddpv2, use_chunk=True, use_zero=True))
+
+
+@pytest.mark.dist
+@pytest.mark.parametrize('world_size', [1, 2])
+@rerun_if_address_is_in_use()
+def test_state_dict(world_size):
+    run_func = partial(run_dist, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+if __name__ == '__main__':
+    test_state_dict(2)