Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests

[Zero] optimize offload mode and add sub dp_group
## Description

- Optimizer offload

  Optimizer offload stores the parameters of the optimizer (12 bytes typically) to host memory. Meanwhile, it also moves the computation of `optimizer.step()` to the host and incurs extra memory copy (copy `grad` from device to host and copy `param` from host to device). The opportunity lies in that we can overlap memory copy and host computation, which can approximately accelerate the optimizer offload by 1.5x (still much slower than non-offload).

- Sub-group in Zero

  Since Zero requires a large communication volume, users may want to combine Zero and DDP (Distributed Data Parallel) if they are equipped with enough device memory. Thus, it is helpful to add a `sub_dp_group` parameter to Zero to enable a sub DDP group for Zero. That is, if `Zero=4` and `sub_group=2`, then it will be DP2-Zero2 ([0,1] and [2,3] constitute two Zero groups respectively, and [0,1]-[2,3] together constitute a DP2 group).

## Methods



Discussion:

Resolving patch:
diff --git a/colossalai/booster/plugin/hybrid_parallel_plugin.py b/colossalai/booster/plugin/hybrid_parallel_plugin.py
--- a/colossalai/booster/plugin/hybrid_parallel_plugin.py
+++ b/colossalai/booster/plugin/hybrid_parallel_plugin.py
@@ -633,6 +633,7 @@ def __init__(
         communication_dtype: Optional[torch.dtype] = None,
         overlap_communication: bool = True,
         partition_grad: bool = False,  # stage 2 flag
+        sub_dp_size: int = 1,
         cpu_offload: bool = False,  # cpu offload
         dp_process_group: Optional[ProcessGroup] = None,  # the dp pg for comm
         tp_process_group: Optional[ProcessGroup] = None,  # if using tp
@@ -663,6 +664,7 @@ def __init__(
             communication_dtype=communication_dtype,
             overlap_communication=overlap_communication,
             partition_grad=partition_grad,
+            sub_dp_size=sub_dp_size,
             cpu_offload=cpu_offload,
             dp_process_group=dp_process_group,
             forced_dtype=forced_dtype,
@@ -964,6 +966,7 @@ def __init__(
         gradient_as_bucket_view: bool = False,
         static_graph: bool = False,
         zero_bucket_size_in_m: int = 12,
+        sub_dp_size: int = 1,
         cpu_offload: bool = False,
         communication_dtype: Optional[torch.dtype] = None,
         overlap_communication: bool = True,
@@ -1070,6 +1073,7 @@ def __init__(
             reduce_bucket_size=zero_bucket_size_in_m * 1024 * 1024,
             communication_dtype=communication_dtype,
             overlap_communication=overlap_communication,
+            sub_dp_size=sub_dp_size,
             cpu_offload=cpu_offload,
             partition_grad=(self.zero_stage == 2),
             forced_dtype=PRECISION_TORCH_TYPE[precision],
diff --git a/colossalai/zero/low_level/_utils.py b/colossalai/zero/low_level/_utils.py
--- a/colossalai/zero/low_level/_utils.py
+++ b/colossalai/zero/low_level/_utils.py
@@ -1,11 +1,36 @@
 import math
-from typing import Optional
+from typing import Iterable, Optional
 
 import torch
 import torch.distributed as dist
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 
 
+class DataPrefetcher:
+    """
+    Asynchronously prefetch data from the loader via a separate stream to improve the performance.
+    """
+
+    def __init__(self, loader: Iterable):
+        self.loader = iter(loader)
+        self.stream = torch.cuda.Stream()
+        self.preload()
+
+    def next(self):
+        self.stream.synchronize()
+        data = self.data
+        self.preload()
+        return data
+
+    def preload(self):
+        try:
+            with torch.cuda.stream(self.stream):
+                self.data = next(self.loader)
+        except StopIteration:
+            self.data = None
+            return
+
+
 def flatten(input_):
     return _flatten_dense_tensors(input_)
 
@@ -190,6 +215,7 @@ def calculate_global_norm_from_list(norm_list):
         total_norm += norm**2.0
     return math.sqrt(total_norm)
 
+
 def sync_tensor(flat_tensor, tensor_list):
     """
     Synchronize the flattened tensor and unflattened tensor list. When
diff --git a/colossalai/zero/low_level/low_level_optim.py b/colossalai/zero/low_level/low_level_optim.py
--- a/colossalai/zero/low_level/low_level_optim.py
+++ b/colossalai/zero/low_level/low_level_optim.py
@@ -1,14 +1,15 @@
 # this code is inspired by the DeepSpeed library and implemented with our own design from scratch
 import copy
+import warnings
 from contextlib import contextmanager
 from functools import partial
 from typing import Dict, Iterator, List, Optional, Tuple
 
+import numpy as np
 import torch
 import torch.distributed as dist
 import torch.nn as nn
 from torch import Tensor, inf
-from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 from torch.distributed import ProcessGroup
 from torch.optim import Optimizer
 
@@ -20,9 +21,15 @@
 )
 from colossalai.interface import OptimizerWrapper
 from colossalai.logging import get_dist_logger
-from colossalai.tensor.moe_tensor.api import is_moe_tensor
 
-from ._utils import calculate_global_norm_from_list, flatten, has_inf_or_nan, release_param_grad, sync_tensor
+from ._utils import (
+    DataPrefetcher,
+    calculate_global_norm_from_list,
+    flatten,
+    has_inf_or_nan,
+    release_param_grad,
+    sync_tensor,
+)
 from .bookkeeping import BucketStore, GradientStore, ParameterStore
 
 
@@ -74,8 +81,8 @@ def __init__(
         partition_grad: bool = False,  # stage 2 flag
         cpu_offload: bool = False,  # cpu offload
         dp_process_group: Optional[ProcessGroup] = None,  # the dp pg for comm
+        sub_dp_size: int = 1,  # further divide zero into sub-dp groups and zero groups
         forced_dtype: Optional[torch.dtype] = None,
-        moe_extra_dp_process_group: Optional[ProcessGroup] = None,
         master_weights: bool = True,  # master weights
     ):
         super(LowLevelZeroOptimizer, self).__init__(optim=optimizer)
@@ -87,24 +94,46 @@ def __init__(
         self._partition_grads = partition_grad
 
         self._cpu_offload = cpu_offload
+        if cpu_offload:
+            warnings.warn(
+                "Please enlarge OMP_NUM_THREADS to speed up the CPU computation. "
+                "Use default OMP_NUM_THREADS=1 would significantly slow down the training."
+            )
 
         # grad accumulation
         self.require_grad_sync = True
 
         # if process_group is none, will use the default one
-        self.dp_pg = dp_process_group
+        if dp_process_group is None:
+            dp_process_group = dist.group.WORLD
+        assert dist.get_world_size(group=dp_process_group) % sub_dp_size == 0
+        dp_ranks = dist.get_process_group_ranks(group=dp_process_group)
+        dp_ranks = np.array(dp_ranks).reshape(sub_dp_size, -1)
+        sub_dp_rank = dist.get_rank(group=dp_process_group) % dp_ranks.shape[1]
+        zero_rank = dist.get_rank(group=dp_process_group) // dp_ranks.shape[1]
+
+        if sub_dp_size == 1:
+            self.dp_pg = dp_process_group
+        else:
+            self.dp_pg = None
+            for i in range(dp_ranks.shape[0]):
+                group = dist.new_group(dp_ranks[i])
+                if i == zero_rank:
+                    assert self.dp_pg is None
+                    self.dp_pg = group
         self._local_rank = dist.get_rank(group=self.dp_pg)
         self._world_size = dist.get_world_size(group=self.dp_pg)
 
-        # extra dp
-        # This group is used to sync moe param, dp_world_size = moe_duplicates * extra_dp_size.
-        # Non moe param will be sync by global dp pg, moe param will be sync by extra dp pg.
-        # Moe param grad is be split as non moe param by global dp pg, and grad will be merged in step.
-        # And moe working and master param are split by extra dp pg.
-        self.moe_extra_dp_pg = moe_extra_dp_process_group
-        if self.moe_extra_dp_pg is not None:
-            self.moe_extra_dp_pg_size = dist.get_world_size(group=self.moe_extra_dp_pg)
-            self.moe_extra_dp_pg_rank = dist.get_rank(group=self.moe_extra_dp_pg)
+        self.sub_dp_pg = None
+        if sub_dp_size > 1:
+            for i in range(dp_ranks.shape[1]):
+                group = dist.new_group(dp_ranks[:, i])
+                if i == sub_dp_rank:
+                    assert self.sub_dp_pg is None
+                    self.sub_dp_pg = group
+        if self.sub_dp_pg is not None:
+            self._sub_dp_rank = dist.get_rank(group=self.sub_dp_pg)
+            self._sub_dp_world_size = dist.get_world_size(group=self.sub_dp_pg)
 
         # working and master params for mixed precision training
         self._working_param_groups = dict()
@@ -137,12 +166,6 @@ def __init__(
         self._grad_store = GradientStore(self.dp_pg, partition_grad=partition_grad)
         self._bucket_store = BucketStore(self.dp_pg)
 
-        # moe param should not be stored in working_groups
-        # because they have different parallel strategy
-        # so we need to store them separately in param_groups
-        # instead of working_groups
-        self.working_moe_params = list()
-
         # iterate over the param group in the optimizer
         # partition these param groups for data parallel training
         # and add buffers to parameter store for future access
@@ -150,11 +173,6 @@ def __init__(
             group_params = list()
             for param in param_group["params"]:
                 if param.requires_grad:
-                    if self.moe_extra_dp_pg is None:
-                        # skip moe param
-                        if is_moe_tensor(param):
-                            self.working_moe_params.append(param)
-                            continue
                     group_params.append(param)
 
             # add the working params to working_param_groups for bookkeeping
@@ -168,25 +186,6 @@ def __init__(
             # managed by this data parallel rank
             param_group["params"] = master_param_current_rank
 
-        # if there are moe params, store in addtional group in optim
-        if len(self.working_moe_params) > 0:
-            self._sync_master_param = False
-            param_group = dict()
-            # create fp32 master param
-            for key, value in self.optim.param_groups[0].items():
-                if key != "params":
-                    param_group[key] = value
-            self.master_moe_params = []
-            for param in self.working_moe_params:
-                self.master_moe_params.append(param.clone().to(torch.float32).detach())
-            # create mapping from master to working for optimizer io
-            self.moe_master_to_working_map = {}
-            for master_moe_param, working_moe_param in zip(self.master_moe_params, self.working_moe_params):
-                self.moe_master_to_working_map[id(master_moe_param)] = working_moe_param
-            # add to optim
-            param_group["params"] = self.master_moe_params
-            self.optim.param_groups.append(param_group)
-
         # initialize communication stream for
         # communication-computation overlapping
         if self._overlap_communication:
@@ -250,12 +249,8 @@ def _create_master_param_current_rank(self, param_list):
                 else:
                     padding_param = param.data.view(-1)
 
-                if self.moe_extra_dp_pg is not None and is_moe_tensor(param):
-                    splited_params = padding_param.split(padding_param.numel() // self.moe_extra_dp_pg_size)
-                    splited_params = splited_params[self.moe_extra_dp_pg_rank]
-                else:
-                    splited_params = padding_param.split(padding_param.numel() // self._world_size)
-                    splited_params = splited_params[self._local_rank]
+                splited_params = padding_param.split(padding_param.numel() // self._world_size)
+                splited_params = splited_params[self._local_rank]
 
                 # use fp32 when master_weights is True
                 if self._master_weights is True:
@@ -295,43 +290,8 @@ def _run_reduction(self):
         if self._bucket_store.num_elements_in_bucket() > 0:
             self._bucket_store.build_grad_in_bucket()
 
-            if self.moe_extra_dp_pg is None:
-                flat_grads = self._bucket_store.get_flatten_grad()
-                flat_grads /= self._world_size
-            else:
-                # record moe and non moe param
-                moe_list = []
-                for param in self._bucket_store._param_list:
-                    moe_list.append(is_moe_tensor(param))
-
-                # divide them into different groups
-                moe_grad_list = []
-                non_moe_grad_list = []
-                for grad_list in self._bucket_store._grad_in_bucket.values():
-                    non_moe_cur_grad = []
-                    moe_cur_grad = []
-                    for i in range(len(grad_list)):
-                        if moe_list[i] == True:
-                            moe_cur_grad.append(grad_list[i])
-                        else:
-                            non_moe_cur_grad.append(grad_list[i])
-                    if len(moe_cur_grad) > 0:
-                        moe_grad_list.append(moe_cur_grad)
-                    if len(non_moe_cur_grad) > 0:
-                        non_moe_grad_list.append(non_moe_cur_grad)
-
-                if len(non_moe_grad_list) > 0:
-                    non_moe_flat_grads = []
-                    for grad_list in non_moe_grad_list:
-                        non_moe_flat_grads.append(_flatten_dense_tensors(grad_list))
-                    non_moe_flat_grads = _flatten_dense_tensors(non_moe_flat_grads)
-                    non_moe_flat_grads /= self._world_size
-
-                if len(moe_grad_list) > 0:
-                    moe_flat_grads = []
-                    for grad_list in moe_grad_list:
-                        moe_flat_grads.append(_flatten_dense_tensors(grad_list))
-                    moe_flat_grads = _flatten_dense_tensors(moe_flat_grads)
+            flat_grads = self._bucket_store.get_flatten_grad()
+            flat_grads /= self._world_size
 
             # ready to add other tensors to bucket
             self._bucket_store.reset_num_elements_in_bucket()
@@ -339,13 +299,7 @@ def _run_reduction(self):
             if self._overlap_communication:
                 stream = self._comm_stream
                 # in case of the memory being reused in the default stream
-                if self.moe_extra_dp_pg is None:
-                    flat_grads.record_stream(stream)
-                else:
-                    if len(non_moe_grad_list) > 0:
-                        non_moe_flat_grads.record_stream(stream)
-                    if len(moe_grad_list) > 0:
-                        moe_flat_grads.record_stream(stream)
+                flat_grads.record_stream(stream)
                 # waiting for ops in the default stream finishing
                 stream.wait_stream(get_accelerator().current_stream())
             else:
@@ -354,84 +308,34 @@ def _run_reduction(self):
             with get_accelerator().stream(stream):
                 group_id = self._bucket_store.current_group_id
 
-                if self.moe_extra_dp_pg is None:
-                    grad_dtype = flat_grads.dtype
-                    if self._communication_dtype is not None:
-                        flat_grads = flat_grads.to(self._communication_dtype)
+                grad_dtype = flat_grads.dtype
+                if self._communication_dtype is not None:
+                    flat_grads = flat_grads.to(self._communication_dtype)
 
                 if not self._partition_grads:
-                    if self.moe_extra_dp_pg is None:
-                        dist.all_reduce(flat_grads, group=self.dp_pg)
-                        if flat_grads.dtype != grad_dtype:
-                            flat_grads = flat_grads.to(grad_dtype)
-
-                        flat_grads_per_rank = flat_grads.split(flat_grads.numel() // self._world_size)
-                        grad_in_bucket = self._bucket_store.get_grad()
-                        self._update_unpartitoned_grad(grad_in_bucket.values(), flat_grads_per_rank, group_id)
-
-                    # sync extra zero group
-                    else:
-                        # sync non moe param in global dp group
-                        if len(non_moe_grad_list) > 0:
-                            dist.all_reduce(non_moe_flat_grads, group=self.dp_pg)
-                            flat_grads_per_rank = non_moe_flat_grads.split(
-                                non_moe_flat_grads.numel() // self._world_size
-                            )
-                            self._update_unpartitoned_grad(non_moe_grad_list, flat_grads_per_rank, group_id)
-
-                        # sync moe param only in zero group
-                        if len(moe_grad_list) > 0:
-                            dist.all_reduce(moe_flat_grads, group=self.moe_extra_dp_pg)
-                            flat_grads_per_rank = moe_flat_grads.split(moe_flat_grads.numel() // self._world_size)
-                            self._update_unpartitoned_grad(moe_grad_list, flat_grads_per_rank, group_id)
+                    dist.all_reduce(flat_grads, group=self.dp_pg)
+                    if self.sub_dp_pg is not None:
+                        dist.all_reduce(flat_grads, op=dist.ReduceOp.AVG, group=self.sub_dp_pg)
+
+                    if flat_grads.dtype != grad_dtype:
+                        flat_grads = flat_grads.to(grad_dtype)
+
+                    flat_grads_per_rank = flat_grads.split(flat_grads.numel() // self._world_size)
+                    grad_in_bucket = self._bucket_store.get_grad()
+                    self._update_unpartitoned_grad(grad_in_bucket.values(), flat_grads_per_rank, group_id)
 
                 else:
-                    if self.moe_extra_dp_pg is None:
-                        flat_grads_list = list(flat_grads.split(len(flat_grads) // self._world_size))
-                        recieved_grad = torch.zeros_like(flat_grads_list[0])
-                        dist.reduce_scatter(recieved_grad, flat_grads_list, group=self.dp_pg)
-
-                        if recieved_grad.dtype != grad_dtype:
-                            recieved_grad = recieved_grad.to(grad_dtype)
-
-                        grad_in_bucket_current_rank = self._bucket_store.get_grad()[self._local_rank]
-                        self._update_partitoned_grad(grad_in_bucket_current_rank, recieved_grad, group_id, 1)
-                    else:
-                        # categorize moe and non moe param
-                        grad_in_bucket_current_rank = self._bucket_store.get_grad()[self._local_rank]
-                        moe_grad_in_bucket_current_rank = []
-                        non_moe_grad_in_bucket_current_rank = []
-                        for idx, grad in enumerate(grad_in_bucket_current_rank):
-                            if moe_list[idx] == True:
-                                moe_grad_in_bucket_current_rank.append(grad)
-                            else:
-                                non_moe_grad_in_bucket_current_rank.append(grad)
-
-                        if len(non_moe_grad_list) > 0:
-                            flat_grads_list = list(
-                                non_moe_flat_grads.split(len(non_moe_flat_grads) // self._world_size)
-                            )
-                            recieved_grad = torch.zeros_like(flat_grads_list[0])
-                            dist.reduce_scatter(recieved_grad, flat_grads_list, group=self.dp_pg)
-                            self._update_partitoned_grad(
-                                non_moe_grad_in_bucket_current_rank, recieved_grad, group_id, 1
-                            )
-
-                        if len(moe_grad_list) > 0:
-                            flat_grads_list = list(
-                                moe_flat_grads.split(len(moe_flat_grads) // self.moe_extra_dp_pg_size)
-                            )
-                            recieved_grad = torch.zeros_like(flat_grads_list[0])
-                            dist.reduce_scatter(recieved_grad, flat_grads_list, group=self.moe_extra_dp_pg)
-                            param_slice = self._world_size // self.moe_extra_dp_pg_size
-                            recieved_grad = list(recieved_grad.split(len(recieved_grad) // param_slice))
-                            for split_recieved_grad in recieved_grad:
-                                split_recieved_grad = _unflatten_dense_tensors(
-                                    split_recieved_grad, moe_grad_in_bucket_current_rank
-                                )
-                                for real_grad, grad in zip(split_recieved_grad, moe_grad_in_bucket_current_rank):
-                                    param_id = self._bucket_store.get_param_id_of_grad(grad)
-                                    self._add_grad(real_grad, param_slice, group_id, param_id)
+                    flat_grads_list = list(flat_grads.split(len(flat_grads) // self._world_size))
+                    recieved_grad = torch.zeros_like(flat_grads_list[0])
+                    dist.reduce_scatter(recieved_grad, flat_grads_list, group=self.dp_pg)
+                    if self.sub_dp_pg is not None:
+                        dist.all_reduce(recieved_grad, op=dist.ReduceOp.AVG, group=self.sub_dp_pg)
+
+                    if recieved_grad.dtype != grad_dtype:
+                        recieved_grad = recieved_grad.to(grad_dtype)
+
+                    grad_in_bucket_current_rank = self._bucket_store.get_grad()[self._local_rank]
+                    self._update_partitoned_grad(grad_in_bucket_current_rank, recieved_grad, group_id, 1)
 
                 self._bucket_store.reset()
 
@@ -572,24 +476,8 @@ def step(self, closure=None):
                 # else the splited grad should be attached to the splited param
                 grads = self._grad_store.get_partitioned_gradients_by_param_id(group_id, id(working_param))
                 if len(grads) > 0:
-                    # moe hybrid zero
-                    if self.moe_extra_dp_pg is not None and is_moe_tensor(working_param):
-                        real_working_params[group_id].append(working_param)
-                        if self._partition_grads:
-                            grad = grads
-                        else:
-                            param_slice = self._world_size // self.moe_extra_dp_pg_size
-                            grad = grads[
-                                self.moe_extra_dp_pg_rank * param_slice : (self.moe_extra_dp_pg_rank + 1) * param_slice
-                            ]
-                        grad = flatten(grad)
-                    else:
-                        real_working_params[group_id].append(working_param)
-                        grad = grads[grad_index]
-                    # no need to copy fp32 grad if master_weights is False
-                    if self._master_weights:
-                        grad = grad.to(splited_param.dtype).to(splited_param.device)
-                    splited_param.grad = grad
+                    real_working_params[group_id].append(working_param)
+                    grad = grads[grad_index]
                     grad_partition_groups.append(grad)
                     real_master_params[group_id].append(splited_param)
 
@@ -603,69 +491,74 @@ def step(self, closure=None):
             # update the params in the optimizer
             self.optim.param_groups[group_id]["params"] = real_master_params[group_id]
 
-        # update param for moe ep
-        # move grad to master param and compute norm
-        if len(self.working_moe_params) > 0:
-            moe_grads = []
-            for master_moe_param, working_moe_param in zip(self.master_moe_params, self.working_moe_params):
-                if master_moe_param.grad is not None:
-                    raise RuntimeError("Moe param should not have grad here")
-                grad = working_moe_param.grad
-                # no need to copy fp32 grad if master_weights is False
-                if self._master_weights:
-                    grad = grad.to(master_moe_param.dtype).to(master_moe_param.device)
-                master_moe_param.grad = grad
-                working_moe_param.grad = None
-                moe_grads.append(grad)
-                grad_partition_groups.append(grad)
-            norm_group = self._compute_grad_norm(gradients=moe_grads)
-            norm_groups.append(norm_group)
-            self.optim.param_groups[-1]["params"] = self.master_moe_params
-            del moe_grads
-
         # unscale and clip grads
         global_norm = calculate_global_norm_from_list(norm_list=norm_groups)
         self._unscale_and_clip_grads(grad_partition_groups, global_norm)
 
-        # update the parameters
-        self.optim.step()
-
-        # release moe grad
-        if len(self.working_moe_params) > 0:
-            for master_moe_param, working_moe_param in zip(self.master_moe_params, self.working_moe_params):
-                master_moe_param.grad = None
-                working_moe_param.data = (
-                    master_moe_param.data.to(working_moe_param.device).to(working_moe_param.dtype).detach()
-                )
-
-        # release the grad
-        grad_partition_groups = []
-        for group_id in range(self.num_param_groups):
-            release_param_grad(self._master_param_groups_of_current_rank[group_id])
-
-        # update working partition updated by the current rank
         device = get_accelerator().get_current_device()
         for group_id in range(self.num_param_groups):
-            master_working_param = self.optim.param_groups[group_id]["params"]
-            for idx, splited_param in enumerate(master_working_param):
-                working_param = real_working_params[group_id][idx]
-                if self.moe_extra_dp_pg is not None and is_moe_tensor(working_param):
-                    all_splited_param = [
-                        torch.zeros(splited_param.shape, device=device, dtype=self._dtype)
-                        for _ in range(self.moe_extra_dp_pg_size)
-                    ]
-                    dist.all_gather(
-                        all_splited_param, splited_param.to(device).to(self._dtype), group=self.moe_extra_dp_pg
-                    )
-                else:
+
+            def load_grad(num: int):
+                """copy grads to the same device and dtype as the master weights"""
+                for i in range(num):
+                    grad = grad_partition_groups.pop(0)
+                    # no need to copy fp32 grad if master_weights is False
+                    if self._master_weights:
+                        grad = grad.to(real_master_params[group_id][i].dtype).to(
+                            real_master_params[group_id][i].device, non_blocking=True
+                        )
+                    yield grad
+
+            def load_param(num: int):
+                """copy params back to the accelerator"""
+                for i in range(num):
+                    splited_param = real_master_params[group_id][i].to(device, non_blocking=True).to(self._dtype)
+                    yield splited_param
+
+            """
+            grad (device) --> grad (host or device) --> optim.step() --> param (host or device) --> param (device)
+            """
+            grad_pre_fetcher, param_pre_fetcher = None, None
+            for idx in range(len(real_master_params[group_id]) + 1):
+                is_first_step = idx == 0
+                is_last_step = idx == len(real_master_params[group_id])
+
+                if not is_last_step:
+                    # update the parameters
+                    if grad_pre_fetcher is None:
+                        grad_pre_fetcher = DataPrefetcher(load_grad(len(real_master_params[group_id])))
+
+                    real_master_params[group_id][idx].grad = grad_pre_fetcher.next()
+                    # HACK: torch optim would skip tensor whose grad is None
+                    self.optim.step()
+                    real_master_params[group_id][idx].grad = None
+                    torch.cuda.current_stream().synchronize()
+
+                if not is_first_step:
+                    # update working partition updated by the current rank
+                    if param_pre_fetcher is None:
+                        param_pre_fetcher = DataPrefetcher(load_param(len(real_master_params[group_id])))
+
+                    working_param = real_working_params[group_id][idx - 1]
+                    splited_param = param_pre_fetcher.next()
+
                     all_splited_param = [
                         torch.zeros(splited_param.shape, device=device, dtype=self._dtype)
                         for _ in range(self._world_size)
                     ]
-                    dist.all_gather(all_splited_param, splited_param.to(device).to(self._dtype), group=self.dp_pg)
-                working_param.data.copy_(flatten(all_splited_param)[: working_param.numel()].reshape_as(working_param))
+                    dist.all_gather(all_splited_param, splited_param, group=self.dp_pg)
+
+                    working_param.data.copy_(
+                        flatten(all_splited_param)[: working_param.numel()].reshape_as(working_param)
+                    )
+                    torch.cuda.current_stream().synchronize()
+
+            # release the grad
+            release_param_grad(self._master_param_groups_of_current_rank[group_id])
             self.optim.param_groups[group_id]["params"] = self._master_param_groups_of_current_rank[group_id]
 
+        assert len(grad_partition_groups) == 0, "grad_partition_groups should be empty after step()"
+
     def _compute_grad_norm(self, gradients: List[Tensor], norm_type: int = 2) -> float:
         r"""
         Compute and return the gradient norm for gradient clipping.
@@ -687,7 +580,9 @@ def _compute_grad_norm(self, gradients: List[Tensor], norm_type: int = 2) -> flo
             total_norm_cuda = torch.tensor(
                 [float(total_norm)], device=get_accelerator().get_current_device(), dtype=torch.float
             )
-            dist.all_reduce(total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=self.dp_pg)
+            dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=self.dp_pg)
+            if self.sub_dp_pg is not None:
+                dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=self.sub_dp_pg)
             total_norm = total_norm_cuda.item()
 
         else:
@@ -700,9 +595,9 @@ def _compute_grad_norm(self, gradients: List[Tensor], norm_type: int = 2) -> flo
             total_norm_exponentiated_cuda = torch.tensor(
                 [float(total_norm_exponentiated)], device=get_accelerator().get_current_device(), dtype=torch.float
             )
-            torch.distributed.all_reduce(
-                total_norm_exponentiated_cuda, op=torch.distributed.ReduceOp.SUM, group=self.dp_pg
-            )
+            dist.all_reduce(total_norm_exponentiated_cuda, op=dist.ReduceOp.SUM, group=self.dp_pg)
+            if self.sub_dp_pg is not None:
+                dist.all_reduce(total_norm_exponentiated_cuda, op=dist.ReduceOp.AVG, group=self.sub_dp_pg)
             total_norm = total_norm_exponentiated_cuda.item() ** (1.0 / norm_type)
 
         return total_norm
@@ -796,16 +691,10 @@ def state_dict(self) -> Dict:
             for k, v in state.items():
                 if isinstance(v, torch.Tensor) and k != "step":
                     working_param = self._param_store.master_to_working_param[id(param)]
-                    if self.moe_extra_dp_pg is not None and is_moe_tensor(v):
-                        gather_tensor = [
-                            torch.zeros(v.shape, device=device, dtype=v.dtype) for _ in range(self.moe_extra_dp_pg_size)
-                        ]
-                        dist.all_gather(gather_tensor, v.to(device), group=self.moe_extra_dp_pg)
-                    else:
-                        gather_tensor = [
-                            torch.zeros(v.shape, device=device, dtype=v.dtype) for _ in range(self._world_size)
-                        ]
-                        dist.all_gather(gather_tensor, v.to(device), group=self.dp_pg)
+                    gather_tensor = [
+                        torch.zeros(v.shape, device=device, dtype=v.dtype) for _ in range(self._world_size)
+                    ]
+                    dist.all_gather(gather_tensor, v.to(device), group=self.dp_pg)
                     param_state = (
                         torch.stack(gather_tensor).view(-1)[: working_param.numel()].reshape_as(working_param).cpu()
                     )
@@ -830,12 +719,8 @@ def load_state_dict(self, state_dict: Dict):
                         v = v.flatten()
                         if padding_size > 0:
                             v = torch.nn.functional.pad(v, [0, padding_size])
-                        if self.moe_extra_dp_pg is not None and is_moe_tensor(v):
-                            v_list = v.split(v.numel() // self.moe_extra_dp_pg_size)
-                            zero_state_dict["state"][param_idx][k] = v_list[self.moe_extra_dp_pg_rank].detach().clone()
-                        else:
-                            v_list = v.split(v.numel() // self._world_size)
-                            zero_state_dict["state"][param_idx][k] = v_list[self._local_rank].detach().clone()
+                        v_list = v.split(v.numel() // self._world_size)
+                        zero_state_dict["state"][param_idx][k] = v_list[self._local_rank].detach().clone()
 
         self.optim.load_state_dict(zero_state_dict)
 
@@ -867,16 +752,8 @@ def state_dict_shard(self, max_shard_size: int = 1024) -> Iterator[Tuple[Dict, i
 
             for k, v in states.items():
                 if isinstance(v, torch.Tensor) and k != "step":
-                    if self.moe_extra_dp_pg is not None and is_moe_tensor(v):
-                        state_tensor = [
-                            torch.zeros(v.shape, device=device, dtype=v.dtype) for _ in range(self.moe_extra_dp_pg_size)
-                        ]
-                        dist.all_gather(state_tensor, v.to(device), group=self.moe_extra_dp_pg)
-                    else:
-                        state_tensor = [
-                            torch.zeros(v.shape, device=device, dtype=v.dtype) for _ in range(self._world_size)
-                        ]
-                        dist.all_gather(state_tensor, v.to(device), group=self.dp_pg)
+                    state_tensor = [torch.zeros(v.shape, device=device, dtype=v.dtype) for _ in range(self._world_size)]
+                    dist.all_gather(state_tensor, v.to(device), group=self.dp_pg)
                     state_tensor = (
                         torch.stack(state_tensor).view(-1)[: working_param.numel()].reshape_as(working_param).cpu()
                     )
@@ -907,18 +784,10 @@ def update_master_params(self, model: nn.Module) -> None:
                 working_param = p.data.view(-1)
                 if padding_size > 0:
                     working_param = torch.nn.functional.pad(working_param, [0, padding_size])
-                if self.moe_extra_dp_pg is not None and is_moe_tensor(p):
-                    master_param.copy_(working_param.chunk(self.extra_dp_pg_size)[self.extra_dp_pg_rank])
-                else:
-                    master_param.copy_(working_param.chunk(self._world_size)[self._local_rank])
-        if hasattr(self, "master_moe_params"):
-            for master_moe_param, working_moe_param in zip(self.master_moe_params, self.working_moe_params):
-                master_moe_param.copy_(working_moe_param)
+                master_param.copy_(working_param.chunk(self._world_size)[self._local_rank])
 
     def get_working_to_master_map(self) -> Dict[int, torch.Tensor]:
         return self._param_store.working_to_master_param
 
     def get_master_to_working_map(self) -> Dict[int, torch.Tensor]:
-        if hasattr(self, "moe_master_to_working_map"):
-            return {**self._param_store.master_to_working_param, **self.moe_master_to_working_map}
         return self._param_store.master_to_working_param


Test patch:
diff --git a/tests/test_zero/test_low_level/test_grad_acc.py b/tests/test_zero/test_low_level/test_grad_acc.py
--- a/tests/test_zero/test_low_level/test_grad_acc.py
+++ b/tests/test_zero/test_low_level/test_grad_acc.py
@@ -8,12 +8,28 @@
 
 import colossalai
 from colossalai.accelerator import get_accelerator
-from colossalai.testing import spawn
+from colossalai.testing import parameterize, spawn
 from colossalai.testing.random import seed_all
 from colossalai.utils import conditional_context
 from colossalai.zero import LowLevelZeroOptimizer
 
 
+def loose_close(a, b, dtype: torch.dtype = torch.float32):
+    rtol = None
+    atol = None
+    if dtype is torch.float16:
+        rtol = 5e-2
+        atol = 5e-4
+    elif dtype is torch.bfloat16:
+        rtol = 4e-3
+        atol = 4e-3
+
+    a = a.detach().to(dtype)
+    b = b.detach().to(dtype)
+
+    assert_close(a, b, rtol=rtol, atol=atol)
+
+
 class MlpModel(nn.Module):
     def __init__(self):
         super(MlpModel, self).__init__()
@@ -26,7 +42,9 @@ def forward(self, x):
         return x
 
 
-def exam_zero_1_2_grad_acc():
+@parameterize("sub_dp_size", [1, 2])
+def exam_zero_1_2_grad_acc(sub_dp_size: int):
+    assert torch.distributed.get_world_size() % sub_dp_size == 0
     local_rank = torch.distributed.get_rank()
     seed_all(2009)
     device = get_accelerator().get_current_device()
@@ -37,10 +55,20 @@ def exam_zero_1_2_grad_acc():
     zero1_optimizer = torch.optim.Adam(zero1_model.parameters(), lr=1)
     zero2_optimizer = torch.optim.Adam(zero2_model.parameters(), lr=1)
     zero1_optimizer = LowLevelZeroOptimizer(
-        zero1_optimizer, overlap_communication=True, initial_scale=32, clip_grad_norm=1.0, verbose=True
+        zero1_optimizer,
+        overlap_communication=True,
+        initial_scale=32,
+        clip_grad_norm=1.0,
+        verbose=True,
+        sub_dp_size=sub_dp_size,
     )
     zero2_optimizer = LowLevelZeroOptimizer(
-        zero2_optimizer, overlap_communication=True, partition_grad=True, initial_scale=32, clip_grad_norm=1.0
+        zero2_optimizer,
+        overlap_communication=True,
+        partition_grad=True,
+        initial_scale=32,
+        clip_grad_norm=1.0,
+        sub_dp_size=sub_dp_size,
     )
     # create data
     seed_all(2021 + local_rank)
@@ -51,7 +79,7 @@ def fwd_bwd_func(number, cur_data, check_flag):
         # zero-dp forward
         zero1_output = zero1_model(cur_data)
         zero2_output = zero2_model(cur_data)
-        assert torch.equal(zero1_output, zero2_output)
+        loose_close(zero1_output, zero2_output)
 
         # zero-dp backward
         zero1_optimizer.backward(zero1_output.sum().float())
@@ -66,10 +94,13 @@ def fwd_bwd_func(number, cur_data, check_flag):
 
     # check updated param
     for z1p, z2p in zip(zero1_model.parameters(), zero2_model.parameters()):
-        assert torch.equal(z1p.data, z2p.data)
+        loose_close(z1p.data, z2p.data)
 
 
-def exam_zero_1_grad_acc(sync):
+@parameterize("no_sync", [True, False])
+@parameterize("sub_dp_size", [1, 2])
+def exam_zero_1_grad_acc(no_sync: bool, sub_dp_size: int):
+    assert torch.distributed.get_world_size() % sub_dp_size == 0
     local_rank = torch.distributed.get_rank()
     seed_all(2008)
     device = get_accelerator().get_current_device()
@@ -89,7 +120,11 @@ def exam_zero_1_grad_acc(sync):
     # in `check_sharded_param_consistency.py`, we will test whether
     # level 1 and 2 will produce exactly the same results
     zero_optimizer = LowLevelZeroOptimizer(
-        zero_optimizer, overlap_communication=False, reduce_bucket_size=262144, clip_grad_norm=1.0
+        zero_optimizer,
+        overlap_communication=False,
+        reduce_bucket_size=262144,
+        clip_grad_norm=1.0,
+        sub_dp_size=sub_dp_size,
     )
 
     torch_optimizer = torch.optim.Adam(torch_model.parameters(), lr=1)
@@ -108,38 +143,37 @@ def fwd_bwd_func(no_sync, cur_data, check_flag):
         # torch-ddp fwd and bwd
         with conditional_context(torch_model.no_sync(), no_sync):
             torch_output = torch_model(cur_data)
-            assert torch.equal(zero_output, torch_output)
+            loose_close(zero_output, torch_output)
             torch_output.sum().backward()
 
         if check_flag:
             # check grad
             for (n, p), z1p in zip(torch_model.named_parameters(), zero_model.parameters()):
-                assert torch.equal(p.grad, z1p.grad)
+                loose_close(p.grad, z1p.grad)
 
-    fwd_bwd_func(sync, input_data1, sync)
+    fwd_bwd_func(no_sync, input_data1, no_sync)
     fwd_bwd_func(False, input_data2, False)
+    torch.nn.utils.clip_grad_norm_(torch_model.parameters(), 1.0)
 
     zero_optimizer.step()
-    torch.nn.utils.clip_grad_norm_(torch_model.parameters(), 1.0)
     torch_optimizer.step()
 
     # check updated param
     for (n, p), z1p in zip(torch_model.named_parameters(), zero_model.parameters()):
         # print(n, p.shape, torch.max(p.data), torch.max(z1p.data), torch.max(torch.abs(p.data - z1p.data)))
-        assert_close(p.data, z1p.data)
+        loose_close(p.data, z1p.data)
 
 
 def run_dist(rank, world_size, port):
     colossalai.launch(config=dict(), rank=rank, world_size=world_size, port=port, host="localhost")
 
-    exam_zero_1_grad_acc(sync=True)
-    exam_zero_1_grad_acc(sync=False)
+    exam_zero_1_grad_acc()
     exam_zero_1_2_grad_acc()
 
 
 @pytest.mark.dist
 def test_grad_accumulation():
-    spawn(run_dist, 2)
+    spawn(run_dist, 4)
 
 
 if __name__ == "__main__":
diff --git a/tests/test_zero/test_low_level/test_zero1_2.py b/tests/test_zero/test_low_level/test_zero1_2.py
--- a/tests/test_zero/test_low_level/test_zero1_2.py
+++ b/tests/test_zero/test_low_level/test_zero1_2.py
@@ -51,7 +51,8 @@ def split_ddp_grad(grad, world_size):
     return splited_grad
 
 
-def exam_zero_1_2():
+@parameterize("sub_dp_size", [1, 2])
+def exam_zero_1_2(sub_dp_size: int):
     """
     In this test, we want to test whether zero stage 1 and 2
     deliver the same numerical results despite different communication
@@ -62,6 +63,7 @@ def exam_zero_1_2():
     pg: partition gradients and optimizer states
 
     """
+    assert torch.distributed.get_world_size() % sub_dp_size == 0
     local_rank = torch.distributed.get_rank()
     seed_all(2001)
 
@@ -73,10 +75,10 @@ def exam_zero_1_2():
     zero1_optimizer = torch.optim.Adam(zero1_model.parameters(), lr=1)
     zero2_optimizer = torch.optim.Adam(zero2_model.parameters(), lr=1)
     zero1_optimizer = LowLevelZeroOptimizer(
-        zero1_optimizer, overlap_communication=True, initial_scale=128, verbose=True
+        zero1_optimizer, overlap_communication=True, initial_scale=128, verbose=True, sub_dp_size=sub_dp_size
     )
     zero2_optimizer = LowLevelZeroOptimizer(
-        zero2_optimizer, overlap_communication=True, partition_grad=True, initial_scale=128
+        zero2_optimizer, overlap_communication=True, partition_grad=True, initial_scale=128, sub_dp_size=sub_dp_size
     )
     # create data
     seed_all(2001 + local_rank)
@@ -94,7 +96,7 @@ def exam_zero_1_2():
     z1g_list = zero1_optimizer._grad_store.get_working_grads_by_group_id(0)
     z2g_list = zero2_optimizer._grad_store.get_working_grads_by_group_id(0)
     for z1g, z2g in zip(z1g_list, z2g_list):
-        assert torch.equal(z1g, z2g)
+        loose_close(z1g, z2g)
 
     # step
     zero1_optimizer.step()
@@ -102,12 +104,13 @@ def exam_zero_1_2():
 
     # check updated param
     for z1p, z2p in zip(zero1_model.parameters(), zero2_model.parameters()):
-        assert torch.equal(z1p.data, z2p.data)
+        loose_close(z1p.data, z2p.data)
 
 
 @parameterize("dtype", [torch.float16, torch.bfloat16])
 @parameterize("master_weights", [True, False])
-def exam_zero_1_torch_ddp(world_size, dtype: torch.dtype, master_weights: bool):
+@parameterize("sub_dp_size", [1, 2])
+def exam_zero_1_torch_ddp(world_size, dtype: torch.dtype, master_weights: bool, sub_dp_size: int):
     """
     In this test, two pairs of model and optimizers are created.
     1. zero: use sharded optimizer and fp16 parameters
@@ -116,6 +119,7 @@ def exam_zero_1_torch_ddp(world_size, dtype: torch.dtype, master_weights: bool):
     We feed these two sets of models with the same input and check if the
     differences in model output and updated parameters are within tolerance.
     """
+    assert world_size % sub_dp_size == 0
     local_rank = torch.distributed.get_rank()
     seed_all(1453)
 
@@ -137,6 +141,7 @@ def exam_zero_1_torch_ddp(world_size, dtype: torch.dtype, master_weights: bool):
         initial_scale=1,
         reduce_bucket_size=1024 * 1024,
         master_weights=master_weights,
+        sub_dp_size=sub_dp_size,
     )
 
     torch_optimizer = torch.optim.SGD(torch_model.parameters(), lr=1)
@@ -162,7 +167,7 @@ def exam_zero_1_torch_ddp(world_size, dtype: torch.dtype, master_weights: bool):
     for (n, p), z1p in zip(torch_model.named_parameters(), zero_model.parameters()):
         if p.grad is not None:
             zero_grad_list = zero_optimizer._grad_store.get_partitioned_gradients_by_param_id(0, id(z1p))
-            torch_grad_list = split_ddp_grad(p.grad, world_size)
+            torch_grad_list = split_ddp_grad(p.grad, world_size // sub_dp_size)
             for zero_grad, torch_grad in zip(zero_grad_list, torch_grad_list):
                 loose_close(zero_grad, torch_grad, dtype=dtype)
 
@@ -187,7 +192,7 @@ def run_dist(rank, world_size, port):
 @pytest.mark.dist
 @rerun_if_address_is_in_use()
 def test_zero_1_2():
-    spawn(run_dist, 2)
+    spawn(run_dist, 4)
 
 
 if __name__ == "__main__":