Feature/zero (#279)

* add zero1 (#209)

* add zero1

* add test zero1

* update zero stage 1 develop (#212)

* Implement naive zero3 (#240)

* naive zero3 works well

* add zero3 param manager

* add TODOs in comments

* add gather full param ctx

* fix sub module streams

* add offload

* fix bugs of hook and add unit tests

* fix bugs of hook and add unit tests (#252)

* add gather full param ctx

* fix sub module streams

* add offload

* fix bugs of hook and add unit tests

* polish code and add state dict hook

* fix bug

* update unit test

* refactor reconstructed zero code

* clip_grad support zero3 and add unit test

* add unit test for Zero3ParameterManager

* [WIP] initialize the shard param class

* [WIP] Yet another sharded model implementation (#274)

* [WIP] initialize the shard param class

* [WIP] Yes another implementation of shardModel. Using a better hook method.

* torch.concat -> torch.cat

* fix test_zero_level_1.py::test_zero_level_1 unitest

* remove deepspeed implementation and refactor for the reconstructed zero module

* polish zero dp unittests
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: Frank Lee <somerlee.9@gmail.com>

colossalai/zero/sharded_optim/bookkeeping/__init__.py

+from .gradient_store import GradientStore
+from .parameter_store import ParameterStore
+from .bucket_store import BucketStore
+from .tensor_bucket import TensorBucket
+
+__all__ = ['GradientStore', 'ParameterStore', 'BucketStore', 'TensorBucket']
\ No newline at end of file

colossalai/zero/sharded_optim/bookkeeping/base_store.py

+from colossalai.core import global_context as gpc
+from colossalai.context import ParallelMode
+
+
+class BaseStore:
+
+    def __init__(self, dp_parallel_mode=ParallelMode.DATA):
+        self._world_size = gpc.get_world_size(dp_parallel_mode)
+        self._local_rank = gpc.get_local_rank(dp_parallel_mode)
+
+    @property
+    def world_size(self):
+        return self._world_size
+
+    @property
+    def local_rank(self):
+        return self._local_rank

colossalai/zero/sharded_optim/bookkeeping/bucket_store.py

+from colossalai.core import global_context as gpc
+from colossalai.context import ParallelMode
+from .base_store import BaseStore
+
+class BucketStore(BaseStore):
+
+    def __init__(self, dp_parallel_mode):
+        super().__init__(dp_parallel_mode)
+        self._grads = dict()
+        self._params = dict()
+        self._num_elements_in_bucket = dict()
+
+        self.reset()
+
+    def num_elements_in_bucket(self, reduce_rank: int = None):
+        return self._num_elements_in_bucket[reduce_rank]
+
+    def add_num_elements_in_bucket(self, num_elements, reduce_rank: int = None):
+        self._num_elements_in_bucket[reduce_rank] += num_elements
+
+    def add_grad(self, tensor, reduce_rank: int = None):
+        self._grads[reduce_rank].append(tensor)
+
+    def add_param(self, tensor, reduce_rank: int = None):
+        self._params[reduce_rank].append(tensor)
+
+    def reset(self):
+        keys = [None] + list(range(self._world_size))
+        self._grads = {rank: [] for rank in keys}
+        self._params = {rank: [] for rank in keys}
+        self._num_elements_in_bucket = {rank: 0 for rank in keys}
+    
+    def reset_by_rank(self, reduce_rank=None):
+        self._grads[reduce_rank] = []
+        self._params[reduce_rank] = []
+        self._num_elements_in_bucket[reduce_rank] = 0
+
+
+    def get_grad(self, reduce_rank: int = None):
+        return self._grads[reduce_rank]
+
+    def get_param(self, reduce_rank: int = None):
+        return self._params[reduce_rank]

colossalai/zero/sharded_optim/bookkeeping/gradient_store.py

+from typing import List
+from torch import Tensor
+from .base_store import BaseStore
+
+
+class GradientStore(BaseStore):
+
+    def __init__(self, *args):
+        super().__init__(*args)
+        # bookkeeping data structures
+        self._averaged_gradients = dict()
+
+        # for backward reduction hooks
+        self._grad_acc_objs = []
+
+    def add_accumulate_grad_object(self, obj):
+        """
+        Keep :class:`AccumulateGrad` objects. If these objects are not kept, reduction hooks may not 
+        be attached successfully.
+
+        :param obj: An object of :class:`AccumulateGrad` class
+        :type obj: :class:`AccumulateGrad`
+        """
+
+        self._grad_acc_objs.append(obj)
+
+    def get_averaged_gradients_by_group(self, group_id: int) -> List[Tensor]:
+        """
+        Return average gradients of a parameter group
+
+        :param group_id: The index of parameter group
+        :type group_id: int
+
+        :return: Return the list of averaged gradients of a parameter group. Each element is a gradient, not a parameter.
+        :rtype: List[torch.Tensor]
+        """
+
+        return self._averaged_gradients[group_id]
+
+    def add_average_gradient_by_group(self, group_id: int, tensor: Tensor) -> None:
+        """
+        Append an average gradient to the list of averaged gradients of a parameter group
+
+        :param group_id: The index of a parameter group
+        :param tensor: A :class:`torch.Tensor` object
+        :type group_id: int
+        :type tensor: torch.Tensor
+
+        """
+
+        if group_id in self._averaged_gradients:
+            self._averaged_gradients[group_id].append(tensor)
+        else:
+            self._averaged_gradients[group_id] = [tensor]
+    
+    def reset_average_gradients_by_group(self, group_id: int) -> None:
+        """
+        Reset the bookkeeping data structure for averaged gradients to an empty list
+
+        :param group_id: The index of a parameter group
+        :type group_id: int
+        """
+        
+        self._averaged_gradients[group_id] = []
+
+

colossalai/zero/sharded_optim/bookkeeping/parameter_store.py

+from .base_store import BaseStore
+from torch import Tensor
+from typing import List
+
+
+class ParameterStore(BaseStore):
+
+    def __init__(self, dp_paralle_mode):
+        super().__init__(dp_paralle_mode)
+        # param partitioning data structures
+        self._fp16_param_to_rank = dict()
+        self._rank_groupid_to_fp16_param_list = dict()
+        self._rank_group_id_to_flat_fp16_param = dict()
+
+        # param reduction data structures
+        self._is_param_reduced = dict()
+        self._reduced_param = []
+
+    def set_param_to_rank(self, tensor: Tensor, rank: int) -> None:
+        """
+        Set the mapping between parameter to rank, each parameter should be owned by a rank.
+
+        :param tensor: A :class:`torch.Tensor` object
+        :type tensor: torch.Tensor
+        :param rank: The rank of which the process is responsible for updating the parameter
+        :type rank: int
+        """
+
+        self._fp16_param_to_rank[tensor] = rank
+
+    def get_param_rank(self, tensor: Tensor) ->  int:
+        """
+        Gives the rank which the parameter belongs to
+
+        :param tensor: A :class:`torch.Tensor` object
+        :type tensor: torch.Tensor
+        """
+        return self._fp16_param_to_rank[tensor]
+
+    def belongs_to_current_rank(self, tensor) -> bool:
+        """
+        Check whether a parameter is supposed to be updated by the process of the current rank
+
+        :param tensor: A :class:`torch.Tensor` object
+        :type tensor: torch.Tensor
+
+        :return: True if the parameter should be updated by the current rank. Otherwise false.
+        :rtype: bool
+        """
+
+        tensor_rank = self._fp16_param_to_rank[tensor]
+        return tensor_rank == self._local_rank
+
+    def add_fp16_param_list_by_rank_group(self, rank, group_id,
+                                          tensor_list) -> None:
+        if rank not in self._rank_groupid_to_fp16_param_list:
+            self._rank_groupid_to_fp16_param_list[rank] = dict()
+
+        if group_id not in self._rank_groupid_to_fp16_param_list[rank]:
+            self._rank_groupid_to_fp16_param_list[rank][group_id] = []
+
+        self._rank_groupid_to_fp16_param_list[rank][group_id].extend(
+            tensor_list)
+
+    def get_fp16_params_by_rank_group(self, rank, group_id) -> List[Tensor]:
+        return self._rank_groupid_to_fp16_param_list[rank][group_id]
+
+    def add_flat_fp16_param_by_rank_group(self, rank, group_id, tensor) -> None:
+        if rank not in self._rank_group_id_to_flat_fp16_param:
+            self._rank_group_id_to_flat_fp16_param[rank] = dict()
+
+        self._rank_group_id_to_flat_fp16_param[rank][group_id] = tensor
+
+    def get_flat_fp16_param_by_rank_group(self, rank, group_id) -> Tensor:
+        return self._rank_group_id_to_flat_fp16_param[rank][group_id]
+
+    def is_param_reduced(self, tensor):
+        return self._is_param_reduced[tensor]
+
+    def set_param_reduction_state(self, tensor, state):
+        self._is_param_reduced[tensor] = state
+
+    def get_param_reduction_states(self):
+        return self._is_param_reduced
+
+    def reset_previous_reduced_params(self):
+        self._reduced_param = []
+
+    def add_previous_reduced_param(self, tensor):
+        self._reduced_param.append(tensor)
+
+    def clear_grads_of_previous_reduced_params(self):
+        if len(self._reduced_param) > 0:
+            for param in self._reduced_param:
+                param.grad = None
+            self.reset_previous_reduced_params()

colossalai/zero/sharded_optim/bookkeeping/tensor_bucket.py

+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+
+
+class TensorBucket:
+
+    def __init__(self, size):
+        self._max_size = size
+        self._current_size = 0
+        self._bucket = []
+
+    @property
+    def max_size(self):
+        return self._max_size
+
+    @property
+    def current_size(self):
+        return self._current_size
+
+    def is_full_or_oversized(self):
+        return self._current_size >= self._max_size
+
+    def is_empty(self):
+        return len(self._bucket) == 0
+
+    def add_to_bucket(self, tensor, allow_oversize=False):
+        tensor_size = tensor.numel()
+
+        if not allow_oversize and self.will_exceed_max_size(tensor_size):
+            msg = f"The param bucket max size {self._max_size} is exceeded" \
+                + f"by tensor (size {tensor_size})"
+            raise RuntimeError(msg)
+
+        self._bucket.append(tensor)
+        self._current_size += tensor_size
+
+    def will_exceed_max_size(self, tensor_size):
+        expected_size = self._current_size + tensor_size
+        return expected_size > self._max_size
+
+    def get_bucket(self):
+        return self._bucket
+
+    def empty(self):
+        self._bucket = []
+        self._size = 0
+
+    def flatten(self):
+        return _flatten_dense_tensors(self._bucket)
+
+    def unflatten_and_copy(self, flat_tensor):
+        unflattened_tensor_list = _unflatten_dense_tensors(
+            flat_tensor, self._bucket)
+        for old, new in zip(self._bucket, unflattened_tensor_list):
+            old.copy_(new)

tests/test_utils/test_zero_gradient_clippling.py

+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+import copy
+import operator as op
+from functools import partial, reduce
+from typing import List
+
+import colossalai
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import torch.nn as nn
+from colossalai.logging import disable_existing_loggers
+from colossalai.utils import checkpoint, clip_grad_norm_fp32, free_port
+from colossalai.zero.sharded_model import ShardedModel
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.nn.utils import clip_grad_norm_
+
+
+class Enumerator:
+    def __init__(self, arg_names: List[str], arg_values: List[tuple]) -> None:
+        self.arg_names = arg_names
+        self.enums = Enumerator.all_enumerate(arg_values)
+
+    def __len__(self):
+        return len(self.enums)
+
+    def __getitem__(self, idx):
+        return {name: self.enums[idx][i] for i, name in enumerate(self.arg_names)}
+
+    @staticmethod
+    def all_enumerate(args: List[tuple]):
+        num_states = reduce(op.mul, map(lambda xs: len(xs), args))
+        idxs = [0] * len(args)
+        states = []
+        for _ in range(num_states):
+            states.append(tuple(args[j][idx] for j, idx in enumerate(idxs)))
+            if len(states) == num_states:
+                break
+            i = 0
+            while idxs[i] + 1 == len(args[i]):
+                idxs[i] = 0
+                i += 1
+            idxs[i] += 1
+        return states
+
+
+def checkpoint_wrapper(module, enable=True):
+    if enable:
+        module.forward = partial(checkpoint, module.forward)
+    return module
+
+
+class Net(nn.Module):
+    def __init__(self, checkpoint=False) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(5, 5)
+        self.fc2 = nn.Linear(5, 5)
+        self.fc3 = nn.Linear(5, 1)
+        if checkpoint:
+            self.fc1 = checkpoint_wrapper(self.fc1)
+        self.layers = [
+            self.fc1,
+            self.fc2,
+            self.fc1,
+            self.fc2,
+            self.fc3
+        ]
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+def run_step(model, optimizer, x, enable_autocast=False, norm_type=2.0):
+    model.train()
+    optimizer.zero_grad()
+    with torch.cuda.amp.autocast(enabled=enable_autocast):
+        y = model(x)
+        loss = y.sum()
+    loss = loss.float()
+    loss.backward()
+    clip_grad(model, norm_type)
+    optimizer.step()
+
+
+def clip_grad(model, norm_type):
+    if isinstance(model, DDP):
+        clip_grad_norm_(model.parameters(), max_norm=1.0, norm_type=norm_type)
+    else:
+        clip_grad_norm_fp32(model.parameters(), max_norm=1.0, norm_type=norm_type)
+
+
+def allclose(tensor_a: torch.Tensor, tensor_b: torch.Tensor, loose=False) -> bool:
+    if loose:
+        return torch.allclose(tensor_a, tensor_b, atol=1e-3, rtol=1e-3)
+    return torch.allclose(tensor_a, tensor_b)
+
+
+def check_grads(model, zero_model, loose=False):
+    rank = dist.get_rank()
+    for p, zero_p in zip(model.parameters(), zero_model.parameters()):
+        zero_grad = zero_p.grad.clone().to(p.device)
+        chunks = torch.flatten(p.grad).chunk(4)
+        if rank >= len(chunks):
+            continue
+        grad = chunks[rank]
+        if zero_p.zero_shard_padding > 0:
+            zero_grad = zero_grad[:-zero_p.zero_shard_padding]
+        assert grad.dtype == zero_grad.dtype
+        assert allclose(grad, zero_grad, loose=loose)
+
+
+def check_params(model, zero_model, loose=False):
+    rank = dist.get_rank()
+    for p, zero_p in zip(model.parameters(), zero_model.parameters()):
+        zero_shard_padding = zero_p.zero_shard_padding
+        zero_p = zero_p.clone().to(p.device)
+        chunks = torch.flatten(p).chunk(4)
+        if rank >= len(chunks):
+            continue
+        p = chunks[rank]
+        if zero_shard_padding > 0:
+            zero_p = zero_p[:-zero_shard_padding]
+        assert p.dtype == zero_p.dtype
+        assert allclose(p, zero_p, loose=loose)
+
+
+def check_config(checkpoint=False, fp16=False, offload=False, norm_type=2.0):
+    model = Net(checkpoint=checkpoint).cuda()
+    zero_model = copy.deepcopy(model)
+    ddp_model = DDP(model)
+
+    offload_config = {}
+    if offload:
+        offload_config['device'] = 'cpu'
+        zero_model = zero_model.cpu()
+    zero_model = ShardedModel(zero_model, mixed_precision=fp16, offload_config=offload_config)
+
+    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=1e-3)
+    zero_optimizer = torch.optim.Adam(zero_model.parameters(), lr=1e-3)
+    for _ in range(5):
+        x = torch.rand(2, 5).cuda()
+        run_step(ddp_model, optimizer, x, enable_autocast=fp16, norm_type=norm_type)
+        run_step(zero_model, zero_optimizer, x, enable_autocast=fp16, norm_type=norm_type)
+        check_grads(ddp_model, zero_model)
+        check_params(ddp_model, zero_model)
+    for _ in range(5):
+        x = torch.rand(2, 5).cuda()
+        run_step(ddp_model, optimizer, x, enable_autocast=False, norm_type=norm_type)
+        run_step(zero_model, zero_optimizer, x, enable_autocast=False, norm_type=norm_type)
+        check_grads(ddp_model, zero_model, loose=True)
+        check_params(ddp_model, zero_model, loose=True)
+
+
+def run_dist(rank, world_size, port):
+    disable_existing_loggers()
+    colossalai.launch(config={},
+                      rank=rank,
+                      world_size=world_size,
+                      host='localhost',
+                      port=port,
+                      backend='nccl')
+
+    args = ['checkpoint', 'fp16', 'offload', 'norm_type']
+    arg_values = [(False, True), (False, True), (False, True), (1.0, 2.0, float('inf'))]
+    arg_enumerator = Enumerator(args, arg_values)
+
+    for kwargs in arg_enumerator:
+        if dist.get_rank() == 0:
+            print(kwargs)
+        check_config(**kwargs)
+    check_config()
+
+
+@ pytest.mark.dist
+def test_zero_clip_grad():
+    world_size = 4
+    run_func = partial(run_dist, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+if __name__ == '__main__':
+    test_zero_clip_grad()

tests/test_zero_data_parallel/common.py

+from functools import partial
+from operator import imod
+from colossalai.utils import checkpoint
+import torch.nn as nn
+import torch
+from colossalai.logging import disable_existing_loggers, get_dist_logger
+
+LOGGER = get_dist_logger()
+
+CONFIG = dict(
+    fp16=dict(
+        mode=None,
+    ),
+    zero=dict(
+        level=3,
+        verbose=False,
+        offload_optimizer_config=dict(
+            device='cpu',
+            pin_memory=True,
+            buffer_count=5,
+            fast_init=False
+        ),
+        offload_param_config=dict(
+            device='cpu',
+            pin_memory=True,
+            buffer_count=5,
+            buffer_size=1e8,
+            max_in_cpu=1e9
+        )
+    ),
+    parallel=dict(
+        pipeline=dict(size=1),
+        tensor=dict(size=1, mode=None)
+    )
+)
+
+def checkpoint_wrapper(module, enable=True):
+    if enable:
+        module.forward = partial(checkpoint, module.forward)
+    return module
+
+
+class Net(nn.Module):
+    def __init__(self, checkpoint=False) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(5, 5)
+        self.fc2 = nn.Linear(5, 5)
+        self.fc3 = nn.Linear(5, 1)
+        if checkpoint:
+            self.fc1 = checkpoint_wrapper(self.fc1)
+        self.layers = [
+            self.fc1,
+            self.fc2,
+            self.fc1,
+            self.fc2,
+            self.fc3
+        ]
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+def allclose(tensor_a: torch.Tensor, tensor_b: torch.Tensor, loose=False) -> bool:
+    if loose:
+        return torch.allclose(tensor_a, tensor_b, atol=1e-3, rtol=1e-3)
+    return torch.allclose(tensor_a, tensor_b)
+
+
+def check_grads(model, zero_model, loose=False):
+    for p, zero_p in zip(model.parameters(), zero_model.parameters()):
+        zero_grad = zero_p.grad.clone().to(p.device)
+        assert p.grad.dtype == zero_grad.dtype
+        assert allclose(p.grad, zero_grad, loose=loose)
+        LOGGER.info(torch.sum(p.grad-zero_grad))
+
+def check_params(model, zero_model, loose=False):
+    for p, zero_p in zip(model.parameters(), zero_model.parameters()):
+        zero_p = zero_p.clone().to(p.device)
+        assert p.dtype == zero_p.dtype
+        assert allclose(p, zero_p, loose=loose)
+

tests/test_zero_data_parallel/test_shard_model_v2.py

+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+import copy
+from functools import partial
+from operator import mod
+from pyexpat import model
+
+import colossalai
+import pytest
+import torch
+import torch.multiprocessing as mp
+from colossalai.logging import disable_existing_loggers
+from colossalai.utils import free_port
+from colossalai.zero.sharded_model import ShardedModelV2
+from colossalai.core import global_context as gpc
+from colossalai.context.parallel_mode import ParallelMode
+from tests.test_zero_data_parallel.common import Net, CONFIG, check_grads
+
+
+def run_fwd_bwd(model, x, enable_autocast=False):
+    model.train()
+    with torch.cuda.amp.autocast(enabled=enable_autocast):
+        y = model(x)
+        loss = y.sum()
+    loss = loss.float()
+    loss.backward()
+    
+
+def run_dist(rank, world_size, port):
+    colossalai.launch(config=CONFIG,
+                      rank=rank,
+                      world_size=world_size,
+                      host='localhost',
+                      port=port,
+                      backend='nccl')
+    
+    model = Net(checkpoint=True).cuda()
+    zero_model = copy.deepcopy(model)
+    zero_model = ShardedModelV2(zero_model, process_group=gpc.get_group(ParallelMode.DATA))
+
+    for _ in range(2):
+        x = torch.rand(2, 5).cuda()
+        run_fwd_bwd(zero_model, x, False)
+        run_fwd_bwd(model, x, False)
+        check_grads(model, zero_model)
+
+
+@pytest.mark.dist
+def test_shard_model_v2():
+    world_size = 2
+    run_func = partial(run_dist, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+if __name__ == '__main__':
+    test_shard_model_v2()

tests/test_zero_data_parallel/test_shard_param.py

+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+from asyncio.log import logger
+from functools import partial
+
+import colossalai
+import pytest
+import torch
+import torch.multiprocessing as mp
+from colossalai.zero.shard_param import ShardParam
+from colossalai.utils import free_port
+from colossalai.logging import get_dist_logger, disable_existing_loggers
+from tests.test_zero_data_parallel.common import Net, CONFIG
+
+def run_shard_param_check(rank, world_size, port):
+    colossalai.launch(config=CONFIG,
+                      rank=rank,
+                      world_size=world_size,
+                      host='localhost',
+                      port=port,
+                      backend='nccl')
+    
+    logger = get_dist_logger()
+    model = Net()
+
+    # add an attribute as ca_attr to hijack the access to param.data
+    for _, param in model.named_parameters():
+        numel_ref = (param.numel() + world_size - 1) // world_size
+        param.ca_attr = ShardParam(param)
+        param.ca_attr.shard()
+        param_data = param.ca_attr.payload(torch.device('cpu'))
+        logger.info(f'shard {param_data.shape} {param_data}', ranks = [1])
+        assert(numel_ref == param_data.numel())
+
+    for _, param in model.named_parameters():
+        param.ca_attr.gather()
+        param_data = param.ca_attr.payload(torch.device('cpu'))
+        logger.info(f'gather {param_data.shape} {param_data}', ranks = [1])
+    
+    disable_existing_loggers([logger])
+
+@pytest.mark.dist
+def test_run_shard_shape():
+    world_size = 2
+    run_func = partial(run_shard_param_check, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+if __name__ == '__main__':
+    test_run_shard_shape()
\ No newline at end of file

tests/test_zero_data_parallel/test_zero_dev_3.py

+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+import copy
+from functools import partial
+
+import colossalai
+import pytest
+import torch
+import torch.multiprocessing as mp
+import torch.nn as nn
+from colossalai.logging import disable_existing_loggers
+from colossalai.utils import checkpoint, free_port
+from colossalai.zero.sharded_model import ShardedModel
+from common import Net, check_grads, check_params, check_params
+
+def checkpoint_wrapper(module, enable=True):
+    if enable:
+        module.forward = partial(checkpoint, module.forward)
+    return module
+
+
+class Net(nn.Module):
+    def __init__(self, checkpoint=False) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(5, 5)
+        self.fc2 = nn.Linear(5, 5)
+        self.fc3 = nn.Linear(5, 1)
+        if checkpoint:
+            self.fc1 = checkpoint_wrapper(self.fc1)
+        self.layers = [
+            self.fc1,
+            self.fc2,
+            self.fc1,
+            self.fc2,
+            self.fc3
+        ]
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+def run_step(model, optimizer, x, enable_autocast=False):
+    model.train()
+    optimizer.zero_grad()
+    with torch.cuda.amp.autocast(enabled=enable_autocast):
+        y = model(x)
+        loss = y.sum()
+    loss = loss.float()
+    loss.backward()
+    optimizer.step()
+
+
+def decode_booleans(intval, bits):
+    res = []
+    for bit in range(bits):
+        mask = 1 << bit
+        res.append((intval & mask) == mask)
+    return res
+
+
+def check_config(checkpoint=False, fp16=False, offload=False):
+    model = Net(checkpoint=checkpoint).cuda()
+    zero_model = copy.deepcopy(model)
+
+    offload_config = {}
+    if offload:
+        offload_config['device'] = 'cpu'
+        zero_model = zero_model.cpu()
+    zero_model = ShardedModel(zero_model, mixed_precision=fp16, offload_config=offload_config)
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+    zero_optimizer = torch.optim.Adam(zero_model.parameters(), lr=1e-3)
+    for _ in range(5):
+        x = torch.rand(2, 5).cuda()
+        run_step(model, optimizer, x, enable_autocast=fp16)
+        run_step(zero_model, zero_optimizer, x, enable_autocast=fp16)
+        check_grads(model, zero_model)
+        check_params(model, zero_model)
+    for _ in range(5):
+        x = torch.rand(2, 5).cuda()
+        run_step(model, optimizer, x, enable_autocast=False)
+        run_step(zero_model, zero_optimizer, x, enable_autocast=False)
+        check_grads(model, zero_model, loose=True)
+        check_params(model, zero_model, loose=True)
+
+
+def run_dist(rank, world_size, port):
+    disable_existing_loggers()
+    colossalai.launch(config={},
+                      rank=rank,
+                      world_size=world_size,
+                      host='localhost',
+                      port=port,
+                      backend='nccl')
+
+    args = ['checkpoint', 'fp16', 'offload']
+
+    def pack_args(i):
+        booleans = decode_booleans(i, len(args))
+        return {arg: booleans[idx] for idx, arg in enumerate(args)}
+
+    for j in range(2 ** len(args)):
+        kwargs = pack_args(j)
+        print(kwargs)
+        check_config(**kwargs)
+
+
+@pytest.mark.dist
+def test_zero_level_3():
+    world_size = 1
+    run_func = partial(run_dist, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+if __name__ == '__main__':
+    test_zero_level_3()

tests/test_zero_data_parallel/test_zero_dev_3_mp4.py

+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+import copy
+from functools import partial
+
+import colossalai
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import torch.nn as nn
+from colossalai.logging import disable_existing_loggers
+from colossalai.utils import checkpoint, free_port
+from colossalai.zero.sharded_model import ShardedModel
+from torch.nn.parallel import DistributedDataParallel as DDP
+from common import Net, allclose
+
+def run_step(model, optimizer, x, enable_autocast=False):
+    model.train()
+    optimizer.zero_grad()
+    with torch.cuda.amp.autocast(enabled=enable_autocast):
+        y = model(x)
+        loss = y.sum()
+    loss = loss.float()
+    loss.backward()
+    optimizer.step()
+
+def check_grads_padding(model, zero_model, loose=False):
+    rank = dist.get_rank()
+    for p, zero_p in zip(model.parameters(), zero_model.parameters()):
+        zero_grad = zero_p.grad.clone().to(p.device)
+        chunks = torch.flatten(p.grad).chunk(4)
+        if rank >= len(chunks):
+            continue
+        grad = chunks[rank]
+        if zero_p.zero_shard_padding > 0:
+            zero_grad = zero_grad[:-zero_p.zero_shard_padding]
+        assert grad.dtype == zero_grad.dtype
+        assert allclose(grad, zero_grad, loose=loose)
+
+
+def check_params_padding(model, zero_model, loose=False):
+    rank = dist.get_rank()
+    for p, zero_p in zip(model.parameters(), zero_model.parameters()):
+        zero_shard_padding = zero_p.zero_shard_padding
+        zero_p = zero_p.clone().to(p.device)
+        chunks = torch.flatten(p).chunk(4)
+        if rank >= len(chunks):
+            continue
+        p = chunks[rank]
+        if zero_shard_padding > 0:
+            zero_p = zero_p[:-zero_shard_padding]
+        assert p.dtype == zero_p.dtype
+        assert allclose(p, zero_p, loose=loose)
+
+
+def decode_booleans(intval, bits):
+    res = []
+    for bit in range(bits):
+        mask = 1 << bit
+        res.append((intval & mask) == mask)
+    return res
+
+
+def check_config(checkpoint=False, fp16=False, offload=False):
+    model = Net(checkpoint=checkpoint).cuda()
+    zero_model = copy.deepcopy(model)
+    ddp_model = DDP(model)
+
+    offload_config = {}
+    if offload:
+        offload_config['device'] = 'cpu'
+        zero_model = zero_model.cpu()
+    zero_model = ShardedModel(zero_model, mixed_precision=fp16, offload_config=offload_config)
+
+    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=1e-3)
+    zero_optimizer = torch.optim.Adam(zero_model.parameters(), lr=1e-3)
+    for _ in range(5):
+        x = torch.rand(2, 5).cuda()
+        run_step(ddp_model, optimizer, x, enable_autocast=fp16)
+        run_step(zero_model, zero_optimizer, x, enable_autocast=fp16)
+        check_grads_padding(ddp_model, zero_model)
+        check_params_padding(ddp_model, zero_model)
+    for _ in range(5):
+        x = torch.rand(2, 5).cuda()
+        run_step(ddp_model, optimizer, x, enable_autocast=False)
+        run_step(zero_model, zero_optimizer, x, enable_autocast=False)
+        check_grads_padding(ddp_model, zero_model, loose=True)
+        check_params_padding(ddp_model, zero_model, loose=True)
+
+
+def run_dist(rank, world_size, port):
+    disable_existing_loggers()
+    colossalai.launch(config={},
+                      rank=rank,
+                      world_size=world_size,
+                      host='localhost',
+                      port=port,
+                      backend='nccl')
+
+    args = ['checkpoint', 'fp16', 'offload']
+
+    def pack_args(i):
+        booleans = decode_booleans(i, len(args))
+        return {arg: booleans[idx] for idx, arg in enumerate(args)}
+
+    for j in range(2 ** len(args)):
+        kwargs = pack_args(j)
+        if dist.get_rank() == 0:
+            print(kwargs)
+        check_config(**kwargs)
+
+
+@pytest.mark.dist
+def test_zero_level_3():
+    world_size = 4
+    run_func = partial(run_dist, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+if __name__ == '__main__':
+    test_zero_level_3()

tests/test_zero_data_parallel/test_zero_level_2.py

-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-import os
-from functools import partial
-from pathlib import Path
-
-import colossalai
-import pytest
-import torch
-import torch.multiprocessing as mp
-from colossalai.core import global_context as gpc
-from colossalai.utils import free_port, get_dataloader
-from torchvision import transforms
-from torchvision.datasets import CIFAR10
-from torchvision.models import resnet18
-
-BATCH_SIZE = 16
-IMG_SIZE = 224
-
-CONFIG = dict(
-    fp16=dict(
-        mode=None,
-    ),
-    zero=dict(
-        level=2,
-        cpu_offload=True,
-        verbose=False,
-    ),
-    parallel=dict(
-        pipeline=dict(size=1),
-        tensor=dict(size=1, mode=None)
-    )
-)
-
-
-def run_dist(rank, world_size, port):
-    colossalai.launch(config=CONFIG,
-                      rank=rank,
-                      world_size=world_size,
-                      host='localhost',
-                      port=port,
-                      backend='nccl')
-
-    # build model
-    model = resnet18(num_classes=10)
-
-    # build dataloader# build dataloaders
-    train_dataset = CIFAR10(
-        root=Path(os.environ['DATA']),
-        download=True,
-        transform=transforms.Compose(
-            [
-                transforms.Resize(size=(IMG_SIZE, IMG_SIZE)),
-                transforms.ToTensor(),
-                transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
-            ]
-        )
-    )
-    train_dataloader = get_dataloader(dataset=train_dataset,
-                                      shuffle=True,
-                                      batch_size=BATCH_SIZE,
-                                      pin_memory=True,
-                                      drop_last=True)
-
-    # build optimizer and loss
-    # optimizer = build_optimizer(global_context.config.optimizer, model)
-    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
-    criterion = torch.nn.CrossEntropyLoss()
-
-    engine, train_dataloader, *args = colossalai.initialize(model=model,
-                                                            optimizer=optimizer,
-                                                            criterion=criterion,
-                                                            train_dataloader=train_dataloader)
-
-    # train
-    model.train()
-    for idx, (data, label) in enumerate(train_dataloader):
-        engine.zero_grad()
-        data = data.cuda()
-        label = label.cuda()
-
-        output = engine(data)
-        loss = engine.criterion(output, label)
-
-        engine.backward(loss)
-        engine.step()
-        break
-
-    gpc.destroy()
-    torch.cuda.empty_cache()
-
-
-@pytest.mark.dist
-def test_zero_level_2():
-    world_size = 4
-    run_func = partial(run_dist, world_size=world_size, port=free_port())
-    mp.spawn(run_func, nprocs=world_size)
-
-
-if __name__ == '__main__':
-    test_zero_level_2()

tests/test_zero_data_parallel/test_zero_level_3.py

-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-import os
-from functools import partial
-from pathlib import Path
-
-import colossalai
-import pytest
-import torch
-import torch.multiprocessing as mp
-from colossalai.core import global_context as gpc
-from colossalai.utils import free_port, get_dataloader
-from torchvision import transforms
-from torchvision.datasets import CIFAR10
-from torchvision.models import resnet18
-
-BATCH_SIZE = 16
-IMG_SIZE = 224
-
-CONFIG = dict(
-    fp16=dict(
-        mode=None,
-    ),
-    zero=dict(
-        level=3,
-        verbose=False,
-        offload_optimizer_config=dict(
-            device='cpu',
-            pin_memory=True,
-            buffer_count=5,
-            fast_init=False
-        ),
-        offload_param_config=dict(
-            device='cpu',
-            pin_memory=True,
-            buffer_count=5,
-            buffer_size=1e8,
-            max_in_cpu=1e9
-        )
-    ),
-    parallel=dict(
-        pipeline=dict(size=1),
-        tensor=dict(size=1, mode=None)
-    )
-)
-
-
-def run_dist(rank, world_size, port):
-    colossalai.launch(config=CONFIG,
-                      rank=rank,
-                      world_size=world_size,
-                      host='localhost',
-                      port=port,
-                      backend='nccl')
-
-    # build model
-    model = resnet18(num_classes=10)
-
-    # build dataloader# build dataloaders
-    train_dataset = CIFAR10(
-        root=Path(os.environ['DATA']),
-        download=True,
-        transform=transforms.Compose(
-            [
-                transforms.Resize(size=(IMG_SIZE, IMG_SIZE)),
-                transforms.ToTensor(),
-                transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
-            ]
-        )
-    )
-    train_dataloader = get_dataloader(dataset=train_dataset,
-                                      shuffle=True,
-                                      batch_size=BATCH_SIZE,
-                                      pin_memory=True,
-                                      drop_last=True)
-
-    # build optimizer and loss
-    # optimizer = build_optimizer(global_context.config.optimizer, model)
-    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
-    criterion = torch.nn.CrossEntropyLoss()
-
-    engine, train_dataloader, *args = colossalai.initialize(model=model,
-                                                            optimizer=optimizer,
-                                                            criterion=criterion,
-                                                            train_dataloader=train_dataloader)
-
-    # train
-    model.train()
-    for idx, (data, label) in enumerate(train_dataloader):
-        engine.zero_grad()
-        data = data.cuda()
-        label = label.cuda()
-
-        output = engine(data)
-        loss = engine.criterion(output, label)
-
-        engine.backward(loss)
-        engine.step()
-        break
-
-    gpc.destroy()
-    torch.cuda.empty_cache()
-
-
-@pytest.mark.dist
-def test_zero_level_3():
-    world_size = 4
-    run_func = partial(run_dist, world_size=world_size, port=free_port())
-    mp.spawn(run_func, nprocs=world_size)
-
-
-if __name__ == '__main__':
-    test_zero_level_3()

tests/test_zero_data_parallel/test_zero_param_mgr.py

+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+import os
+from functools import partial
+from pathlib import Path
+
+import colossalai
+import pytest
+import torch
+import torch.multiprocessing as mp
+from colossalai.zero.sharded_model.param_manager import Zero3ParameterManager
+from colossalai.core import global_context as gpc
+from colossalai.context.parallel_mode import ParallelMode
+from colossalai.utils import free_port
+from common import CONFIG
+
+def run_shard_shape_check(rank, world_size, port):
+    colossalai.launch(config=CONFIG,
+                      rank=rank,
+                      world_size=world_size,
+                      host='localhost',
+                      port=port,
+                      backend='nccl')
+    
+    model = torch.nn.Linear(2, 4 * world_size)
+    gpc.init_parallel_groups()
+    Zero3ParameterManager(module=model, process_group=gpc.get_group(ParallelMode.DATA), offload_config=CONFIG.get('offload_param_config'))
+
+    assert(model.weight.numel() == 4 * 2)
+    assert(model.bias.numel() == 4)
+
+
+@pytest.mark.dist
+def test_run_shard_shape():
+    world_size = 2
+    run_func = partial(run_shard_shape_check, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+if __name__ == '__main__':
+    test_run_shard_shape()

tests/test_zero_tensor_parallel/test_vit_2d_level_2.py

@@ -88,6 +88,7 @@ def run_2d_parallel_vision_transformer_level_2(rank, world_size, port):
 
 
 @pytest.mark.dist
+@pytest.mark.skip(reason="This test should be refactored for the reconstructed zero")
 def test_2d_vit_zero_level_2():
     world_size = 8
     run_func = partial(run_2d_parallel_vision_transformer_level_2, world_size=world_size, port=free_port())

tests/test_zero_tensor_parallel/test_vit_2d_level_3.py

@@ -88,7 +88,7 @@ def run_2d_parallel_vision_transformer_level_3(rank, world_size, port):
 
 
 @pytest.mark.dist
-@pytest.mark.skip("Level 3 has unknown bug so skip this test for now")
+@pytest.mark.skip(reason="This test should be refactored for the reconstructed zero")
 def test_3d_vit_zero_level_3():
     world_size = 8
     run_func = partial(run_2d_parallel_vision_transformer_level_3, world_size=world_size, port=free_port())