Repository: hpcaitech/ColossalAI
Pull request: [zero] add zero wrappers
### What's New

I added two wrapper functions to unify current ZeRO APIs. In these functions, users can simply set the arguement `zero_stage` to switch to different types of ZeRO implementation.

Problem statement:
[FEATURE]: Unify interface of Gemini(offload) + ZeRO1 + ZeRO2 + ZeRO3
### Describe the feature

## Current States

Currently, GeminiDDP has to shard the optimizer state (OS), and it has covered zero3 and zero3+offload. However, It doesn't actually cover zero1 and zero2, which do not shard the OS. Considering zero1/2 performs very well for small models, we should definitely support them in ColossalAI.

We have some individual APIs for each optimizer, such as `GeminiDDP`, `GeminiAdamOptimizer`(zero3 + zero-offload), `LowLevelZeroOptimizer` (zero1+2).

## Purpose

We intend to unify zero0/2/3/3+offload into a single API interface. The API can be designed like this: 

```python
# define model
with ColoInitContext(device=get_current_device(), dtype=torch.half):
    model = model_builder(args.model_type)(checkpoint=True)

# define Optimizer
optimizer = ColoOptimizer(model.parameters(), optimizer_class=torch.optim.Adam, lr=0.01, type = [zero1 | zero2 | zero3 | gemini])

# maybe a model as an argument is more practical for implementation.
optimizer = ColoOptimizer(model, optimizer_class=torch.optim.Adam, lr=0.01, type = [zero1 | zero2 | zero3 | gemini])
```

Hints from issue discussion:

Commit message: [zero] add zero wrappers

Patch to review:
diff --git a/colossalai/nn/parallel/zero_wrapper.py b/colossalai/nn/parallel/zero_wrapper.py
new file mode 100644
index 000000000000..49e1ae29d193
--- /dev/null
+++ b/colossalai/nn/parallel/zero_wrapper.py
@@ -0,0 +1,106 @@
+from copy import copy
+from typing import Dict, Optional
+
+import torch
+import torch.nn as nn
+
+from .gemini_parallel import GeminiDDP
+
+
+def zero_model_wrapper(model: nn.Module, zero_stage: int = 1, gemini_config: Optional[Dict] = None):
+    """This wrapper function is used to wrap your training model for ZeRO DDP.
+
+    Example:
+
+        >>> with ColoInitContext():
+        >>>     my_model = Bert()
+        >>> my_optim = SGD(my_model.parameters(), lr = 1e-3)
+        >>> zero_model = zero_model_wrapper(my_model, zero_stage=1)
+        >>> zero_optim = zero_optim_wrapper(zero_model, my_optim)
+
+    Args:
+        model (nn.Module): The model used in ZeRO DDP.
+        zero_stage (int, optional): The stage of ZeRO DDP. You can find more information in ZeRO's paper.
+            https://arxiv.org/abs/1910.02054
+        gemini_config (dict, optional): The configuration dictionary of `GeminiDDP`. `GeminiDDP` is enabled
+            when the stage is set to 3. You can set the arguemnts of `GeminiDDP` in the gemini_config.
+            Here is an example where we set the device of the model, the placement policy of Gemini, and the
+            size of hidden dimension to help Gemini find out a unified chunk size.
+
+            Example:
+
+                >>> config_dict = dict(device=torch.cuda.current_device(), hidden_dim=1024, placement_policy='auto')
+                >>> model = zero_model_wrapper(model, zero_stage=3, gemini_config=config_dict)
+    """
+    setattr(model, "_colo_zero_stage", zero_stage)
+    assert zero_stage in [1, 2, 3], "The stage of ZeRO should be 1, 2 or 3"
+
+    if gemini_config is None:
+        gemini_config = dict()
+
+    if zero_stage in [1, 2]:
+        return model
+    else:
+        return GeminiDDP(model, **gemini_config)
+
+
+def zero_optim_wrapper(model: nn.Module,
+                       optimizer: torch.optim.Optimizer,
+                       initial_scale: float = 2**16,
+                       growth_factor: float = 2,
+                       backoff_factor: float = 0.5,
+                       growth_interval: int = 1000,
+                       hysteresis: int = 2,
+                       min_scale: float = 1,
+                       max_scale: float = 2**32,
+                       max_norm: float = 0.0,
+                       norm_type: float = 2.0,
+                       zero_config: Optional[Dict] = None):
+    """This wrapper function is used to wrap your training optimizer for ZeRO DDP.
+
+    Args:
+        model (nn.Module): Your model wrapped by `zero_model_wrapper`
+        optimizer (torch.optim.Optimizer): Your initialized optimizer
+        initial_scale (float, optional): initial_scale used by DynamicGradScaler.
+        min_scale (float, optional): min_scale used by DynamicGradScaler.
+        growth_factor (float, optional): growth_factor used by DynamicGradScaler.
+        backoff_factor (float, optional): backoff_factor used by DynamicGradScaler.
+        growth_interval (float, optional): growth_interval used by DynamicGradScaler.
+        hysteresis (float, optional): hysteresis used by DynamicGradScaler.
+        max_scale (int, optional): max_scale used by DynamicGradScaler.
+        max_norm (float, optional): max_norm used for `clip_grad_norm`. You should notice that you shall not do
+            clip_grad_norm by yourself when using ZeRO DDP. The ZeRO optimizer will take care of clip_grad_norm.
+        norm_type (float, optional): norm_type used for `clip_grad_norm`.
+        zero_config (dict, optinoal): The configuration used for the ZeRO optimizer.
+            Example:
+
+                >>> zero2_config = dict(reduce_bucket_size=12 * 1024 * 1024, overlap_communication=True)
+                >>> optim = zero_optim_wrapper(model, optim, zero_config=zero2_config)
+    """
+    assert hasattr(model, "_colo_zero_stage"), "You should use `zero_ddp_wrapper` first"
+    zero_stage = getattr(model, "_colo_zero_stage")
+
+    assert norm_type == 2.0, "Current ZeRO optimizers only support 'norm_type=2'"
+
+    if zero_config is None:
+        config_dict = dict()
+    else:
+        config_dict = copy(zero_config)
+
+    config_dict['initial_scale'] = initial_scale
+    config_dict['growth_factor'] = growth_factor
+    config_dict['backoff_factor'] = backoff_factor
+    config_dict['growth_interval'] = growth_interval
+    config_dict['hysteresis'] = hysteresis
+    config_dict['min_scale'] = min_scale
+    config_dict['max_scale'] = max_scale
+
+    if zero_stage in [1, 2]:
+        from colossalai.zero.sharded_optim.low_level_optim import LowLevelZeroOptimizer
+        config_dict['partition_grad'] = zero_stage == 2
+        config_dict['clip_grad_norm'] = max_norm
+        return LowLevelZeroOptimizer(optimizer, **zero_config)
+    else:
+        from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer
+        config_dict['clipping_norm'] = max_norm
+        return ZeroOptimizer(optimizer, model, **zero_config)
diff --git a/colossalai/zero/sharded_optim/low_level_optim.py b/colossalai/zero/sharded_optim/low_level_optim.py
index f45b5e200a61..d174fc6ac138 100644
--- a/colossalai/zero/sharded_optim/low_level_optim.py
+++ b/colossalai/zero/sharded_optim/low_level_optim.py
@@ -17,7 +17,6 @@
     calculate_global_norm_from_list,
     compute_norm,
     flatten,
-    get_grad_accumulate_object,
     has_inf_or_nan,
     reduce_tensor_dp_group,
     release_param_grad,
@@ -386,7 +385,7 @@ def _add_to_reduction_bucket(self, param, reduce_rank=None):
     # torch.optim.Optimizer methods
     ################################
 
-    def backward(self, loss, retain_graph=False):
+    def backward(self, loss, retain_graph=False, sync_grad=True):
         loss = self.loss_scale * loss
         loss.backward(retain_graph=retain_graph)
 
@@ -402,6 +401,10 @@ def backward(self, loss, retain_graph=False):
             torch.cuda.synchronize()
             self._param_store.clear_grads_of_previous_reduced_params()
 
+        # gradient synchronization
+        if sync_grad:
+            self._sync_grad()
+
     def zero_grad(self, set_to_none=True):
         """
         Set parameter gradients to zero. If set_to_none = True, gradient
@@ -537,7 +540,7 @@ def _unscale_and_clip_grads(self, grad_groups_flat, total_norm):
     # Gradient Synchronization #
     ############################
 
-    def sync_grad(self):
+    def _sync_grad(self):
         # update param already reduced flag
         reduction_states = self._param_store.get_param_reduction_states()
         for tensor, state in reduction_states.items():
diff --git a/tests/test_zero/low_level_zero/test_grad_acc.py b/tests/test_zero/low_level_zero/test_grad_acc.py
index 1e157c70a7ef..504df202e168 100644
--- a/tests/test_zero/low_level_zero/test_grad_acc.py
+++ b/tests/test_zero/low_level_zero/test_grad_acc.py
@@ -9,7 +9,6 @@
 from torch.testing import assert_close
 
 import colossalai
-from colossalai.tensor import ProcessGroup
 from colossalai.testing.random import seed_all
 from colossalai.utils import free_port
 from colossalai.zero import LowLevelZeroOptimizer
@@ -60,16 +59,16 @@ def fwd_bwd_func(number, cur_data):
         assert torch.equal(zero1_output, zero2_output)
 
         # zero-dp backward
-        zero1_optimizer.backward(zero1_output.sum().float())
-        zero2_optimizer.backward(zero2_output.sum().float())
+        zero1_optimizer.backward(zero1_output.sum().float(), sync_grad=False)
+        zero2_optimizer.backward(zero2_output.sum().float(), sync_grad=False)
 
         for (n, z1p), z2p in zip(zero1_model.named_parameters(), zero2_model.parameters()):
             if z2p.grad is not None:
                 # print(local_rank, n, z1p.shape, torch.max(z2p.grad), torch.max(torch.abs(z1p.grad - z2p.grad)))
                 assert torch.equal(z1p.grad, z2p.grad)
 
-        zero1_optimizer.sync_grad()
-        zero2_optimizer.sync_grad()
+        zero1_optimizer._sync_grad()
+        zero2_optimizer._sync_grad()
 
     fwd_bwd_func(0, input_data1)
     fwd_bwd_func(1, input_data2)
@@ -124,7 +123,7 @@ def fwd_bwd_func(number, cur_data, check_flag):
         assert torch.equal(zero_output, torch_output)
 
         # zero-dp backward
-        zero_optimizer.backward(zero_output.sum().float())
+        zero_optimizer.backward(zero_output.sum().float(), sync_grad=False)
         # torch-ddp backward
         torch_output.sum().backward()
 
@@ -135,7 +134,7 @@ def fwd_bwd_func(number, cur_data, check_flag):
                 # print(n, p.shape, torch.max(torch.abs(p.grad - unscale_grad)))
                 assert torch.equal(p.grad, unscale_grad)
 
-        zero_optimizer.sync_grad()
+        zero_optimizer._sync_grad()
 
     fwd_bwd_func(0, input_data1, True)
     fwd_bwd_func(1, input_data2, False)
diff --git a/tests/test_zero/low_level_zero/test_zero1_2.py b/tests/test_zero/low_level_zero/test_zero1_2.py
index 4949630725ca..930b6129174e 100644
--- a/tests/test_zero/low_level_zero/test_zero1_2.py
+++ b/tests/test_zero/low_level_zero/test_zero1_2.py
@@ -78,16 +78,16 @@ def exam_zero_1_2():
     assert torch.equal(zero1_output, zero2_output)
 
     # zero-dp backward
-    zero1_optimizer.backward(zero1_output.mean().float())
-    zero2_optimizer.backward(zero2_output.mean().float())
+    zero1_optimizer.backward(zero1_output.mean().float(), sync_grad=False)
+    zero2_optimizer.backward(zero2_output.mean().float(), sync_grad=False)
 
     for (n, z1p), z2p in zip(zero1_model.named_parameters(), zero2_model.parameters()):
         if z2p.grad is not None:
             # print(local_rank, n, z1p.shape, torch.max(z2p.grad), torch.max(torch.abs(z1p.grad - z2p.grad)))
             assert torch.equal(z1p.grad, z2p.grad)
 
-    zero1_optimizer.sync_grad()
-    zero2_optimizer.sync_grad()
+    zero1_optimizer._sync_grad()
+    zero2_optimizer._sync_grad()
 
     # step
     zero1_optimizer.step()
@@ -146,7 +146,7 @@ def exam_zero_1_torch_ddp():
     half_close(zero_output, torch_output, loose=True)
 
     # zero-dp backward
-    zero_optimizer.backward(zero_output.mean().float())
+    zero_optimizer.backward(zero_output.mean().float(), sync_grad=False)
 
     # torch-ddp backward
     torch_output.mean().backward()
@@ -156,7 +156,7 @@ def exam_zero_1_torch_ddp():
         half_close(p.grad, z1p.grad, loose=True)
 
     # zero-dp step
-    zero_optimizer.sync_grad()
+    zero_optimizer._sync_grad()
     zero_optimizer.step()
 
     # torch ddp step
diff --git a/tests/test_zero/low_level_zero/test_zero_tp.py b/tests/test_zero/low_level_zero/test_zero_tp.py
index ea8e3a0a3369..15d3530ff90a 100644
--- a/tests/test_zero/low_level_zero/test_zero_tp.py
+++ b/tests/test_zero/low_level_zero/test_zero_tp.py
@@ -74,7 +74,6 @@ def exam_zero_with_tp(overlap_flag, partition_flag):
     torch_loss.backward()
     torch.nn.utils.clip_grad_norm_(torch_model.parameters(), 1.0)
     hybrid_optim.backward(hybrid_loss)
-    hybrid_optim.sync_grad()
 
     torch_optim.step()
     hybrid_optim.step()


Code review comments:
diff_hunk: @@ -0,0 +1,106 @@
+from copy import copy
+from typing import Dict, Optional
+
+import torch
+import torch.nn as nn
+
+from .gemini_parallel import GeminiDDP
+
+
+def zero_model_wrapper(model: nn.Module, zero_stage: int = 1, gemini_config: Optional[Dict] = None):
+    """This wrapper function is used to wrap your training model for ZeRO DDP.
+
+    Example:
+
+        >>> with ColoInitContext():
+        >>>     my_model = Bert()
+        >>> my_optim = SGD(my_model.parameters(), lr = 1e-3)
+        >>> zero_model = zero_model_wrapper(my_model, zero_stage=1)
+        >>> zero_optim = zero_optim_wrapper(zero_model, my_optim)
+
+    Args:
+        model (nn.Module): The model used in ZeRO DDP.
+        zero_stage (int, optional): The stage of ZeRO DDP. You can find more information in ZeRO's paper.
+            https://arxiv.org/abs/1910.02054
+        gemini_config (dict, optional): The configuration dictionary of `GeminiDDP`. `GeminiDDP` is enabled
+            when the stage is set to 3. You can set the arguemnts of `GeminiDDP` in the gemini_config.
+            Here is an example where we set the device of the model, the placement policy of Gemini, and the
+            size of hidden dimension to help Gemini find out a unified chunk size.
+
+            Example:
+
+                >>> config_dict = dict(device=torch.cuda.current_device(), hidden_dim=1024, placement_policy='auto')
+                >>> model = zero_model_wrapper(model, zero_stage=3, gemini_config=config_dict)
+    """
+    setattr(model, "_colo_zero_stage", zero_stage)
+    assert zero_stage in [1, 2, 3], "The stage of ZeRO should be 1, 2 or 3"
+
+    if gemini_config is None:
+        gemini_config = dict()
+
+    if zero_stage in [1, 2]:
+        return model
+    else:
+        return GeminiDDP(model, **gemini_config)
+
+
+def zero_optim_wrapper(model: nn.Module,
+                       optimizer: torch.optim.Optimizer,
+                       initial_scale: float = 2**16,
+                       growth_factor: float = 2,
+                       backoff_factor: float = 0.5,
+                       growth_interval: int = 1000,
+                       hysteresis: int = 2,
+                       min_scale: float = 1,
+                       max_scale: float = 2**32,
+                       max_norm: float = 0.0,
+                       norm_type: float = 2.0,
+                       zero_config: Optional[Dict] = None):
path: colossalai/nn/parallel/zero_wrapper.py
text: @user1:
zero_config->config

diff_hunk: @@ -0,0 +1,106 @@
+from copy import copy
+from typing import Dict, Optional
+
+import torch
+import torch.nn as nn
+
+from .gemini_parallel import GeminiDDP
+
+
+def zero_model_wrapper(model: nn.Module, zero_stage: int = 1, gemini_config: Optional[Dict] = None):
+    """This wrapper function is used to wrap your training model for ZeRO DDP.
+
+    Example:
+
+        >>> with ColoInitContext():
+        >>>     my_model = Bert()
+        >>> my_optim = SGD(my_model.parameters(), lr = 1e-3)
+        >>> zero_model = zero_model_wrapper(my_model, zero_stage=1)
+        >>> zero_optim = zero_optim_wrapper(zero_model, my_optim)
+
+    Args:
+        model (nn.Module): The model used in ZeRO DDP.
+        zero_stage (int, optional): The stage of ZeRO DDP. You can find more information in ZeRO's paper.
+            https://arxiv.org/abs/1910.02054
+        gemini_config (dict, optional): The configuration dictionary of `GeminiDDP`. `GeminiDDP` is enabled
+            when the stage is set to 3. You can set the arguemnts of `GeminiDDP` in the gemini_config.
+            Here is an example where we set the device of the model, the placement policy of Gemini, and the
+            size of hidden dimension to help Gemini find out a unified chunk size.
+
+            Example:
+
+                >>> config_dict = dict(device=torch.cuda.current_device(), hidden_dim=1024, placement_policy='auto')
+                >>> model = zero_model_wrapper(model, zero_stage=3, gemini_config=config_dict)
+    """
+    setattr(model, "_colo_zero_stage", zero_stage)
+    assert zero_stage in [1, 2, 3], "The stage of ZeRO should be 1, 2 or 3"
+
+    if gemini_config is None:
+        gemini_config = dict()
+
+    if zero_stage in [1, 2]:
+        return model
+    else:
+        return GeminiDDP(model, **gemini_config)
+
+
+def zero_optim_wrapper(model: nn.Module,
+                       optimizer: torch.optim.Optimizer,
+                       initial_scale: float = 2**16,
+                       growth_factor: float = 2,
+                       backoff_factor: float = 0.5,
+                       growth_interval: int = 1000,
+                       hysteresis: int = 2,
+                       min_scale: float = 1,
+                       max_scale: float = 2**32,
+                       max_norm: float = 0.0,
+                       norm_type: float = 2.0,
+                       zero_config: Optional[Dict] = None):
+    """This wrapper function is used to wrap your training optimizer for ZeRO DDP.
+
+    Args:
+        model (nn.Module): Your model wrapped by `zero_model_wrapper`
+        optimizer (torch.optim.Optimizer): Your initialized optimizer
+        initial_scale (float, optional): initial_scale used by DynamicGradScaler.
+        min_scale (float, optional): min_scale used by DynamicGradScaler.
+        growth_factor (float, optional): growth_factor used by DynamicGradScaler.
+        backoff_factor (float, optional): backoff_factor used by DynamicGradScaler.
+        growth_interval (float, optional): growth_interval used by DynamicGradScaler.
+        hysteresis (float, optional): hysteresis used by DynamicGradScaler.
+        max_scale (int, optional): max_scale used by DynamicGradScaler.
+        max_norm (float, optional): max_norm used for `clip_grad_norm`. You should notice that you shall not do
+            clip_grad_norm by yourself when using ZeRO DDP. The ZeRO optimizer will take care of clip_grad_norm.
+        norm_type (float, optional): norm_type used for `clip_grad_norm`.
+        zero_config (dict, optinoal): The configuration used for the ZeRO optimizer.
+            Example:
+
+                >>> zero2_config = dict(reduce_bucket_size=12 * 1024 * 1024, overlap_communication=True)
+                >>> optim = zero_optim_wrapper(model, optim, zero_config=zero2_config)
+    """
+    assert hasattr(model, "_colo_zero_stage"), "You should use `zero_ddp_wrapper` first"
+    zero_stage = getattr(model, "_colo_zero_stage")
+
+    assert norm_type == 2.0, "Current ZeRO optimizers only support 'norm_type=2'"
+
+    if zero_config is None:
+        config_dict = dict()
+    else:
+        config_dict = copy(zero_config)
+
+    config_dict['initial_scale'] = initial_scale
+    config_dict['growth_factor'] = growth_factor
+    config_dict['backoff_factor'] = backoff_factor
+    config_dict['growth_interval'] = growth_interval
+    config_dict['hysteresis'] = hysteresis
+    config_dict['min_scale'] = min_scale
+    config_dict['max_scale'] = max_scale
+
+    if zero_stage in [1, 2]:
+        from colossalai.zero.sharded_optim.low_level_optim import LowLevelZeroOptimizer
+        config_dict['partition_grad'] = zero_stage == 2
+        config_dict['clip_grad_norm'] = max_norm
+        return LowLevelZeroOptimizer(optimizer, **zero_config)
+    else:
+        from colossalai.nn.optimizer.zero_optimizer import ZeroOptimizer
+        config_dict['clipping_norm'] = max_norm
+        return ZeroOptimizer(optimizer, model, **zero_config)
path: colossalai/nn/parallel/zero_wrapper.py
text: @user1:
optimizer must be a HybridAdam here? Did you check the type?

@author:
The ZeroOptimizer will check the type of the used optimizer. If the user doesn't use the pointed optimizer, it will print a available optimizer list.