[gemini] update the gpt example (#2527)

colossalai/nn/parallel/zero_wrapper.py

@@ -32,16 +32,19 @@ def zero_model_wrapper(model: nn.Module, zero_stage: int = 1, gemini_config: Opt
                 >>> 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
+        wrapped_model = model
     else:
-        return GeminiDDP(model, **gemini_config)
+        wrapped_model = GeminiDDP(model, **gemini_config)
+
+    setattr(wrapped_model, "_colo_zero_stage", zero_stage)
+
+    return wrapped_model
 
 
 def zero_optim_wrapper(model: nn.Module,

examples/language/gpt/gemini/benchmark_gemini.sh

 for MODEL_TYPE in "gpt2_medium"; do
-  for DISTPLAN in "colossalai"; do
+  for DISTPLAN in "CAI_Gemini"; do
     for BATCH_SIZE in 16; do
       for GPUNUM in 1 2 4 8; do
         for TPDEGREE in 1 2 4 8; do

examples/language/gpt/gemini/run_gemini.sh

 set -x
-# distplan in ["colossalai", "zero1", "zero2", "torch_ddp", "torch_zero"]
-export DISTPLAN=${DISTPLAN:-"colossalai"}
+# distplan in ["CAI_ZeRO1", "CAI_ZeRO2", "CAI_Gemini", "Pytorch_DDP", "Pytorch_ZeRO"]
+export DISTPLAN=${DISTPLAN:-"CAI_Gemini"}
 
 # The following options only valid when DISTPLAN="colossalai"
 export GPUNUM=${GPUNUM:-1}
@@ -12,6 +12,12 @@ export MODEL_TYPE=${MODEL_TYPE:-"gpt2_medium"}
 export TRAIN_STEP=${TRAIN_STEP:-10}
 # export PYTHONPATH=$PWD:$PYTHONPATH
 
+if [ ${USE_SHARD_INIT} = "True" ]; then
+  USE_SHARD_INIT="--shardinit"
+else
+  USE_SHARD_INIT=""
+fi
+
 mkdir -p gemini_logs
 
 torchrun --standalone --nproc_per_node=${GPUNUM} ./train_gpt_demo.py \
@@ -19,7 +25,7 @@ torchrun --standalone --nproc_per_node=${GPUNUM} ./train_gpt_demo.py \
 --model_type=${MODEL_TYPE} \
 --batch_size=${BATCH_SIZE} \
 --placement=${PLACEMENT} \
---shardinit=${USE_SHARD_INIT} \
+${USE_SHARD_INIT} \
 --distplan=${DISTPLAN} \
 --train_step=${TRAIN_STEP} \
 2>&1 | tee ./gemini_logs/${MODEL_TYPE}_${DISTPLAN}_gpu_${GPUNUM}_bs_${BATCH_SIZE}_tp_${TPDEGREE}_${PLACEMENT}.log

examples/language/gpt/gemini/train_gpt_demo.py

@@ -12,26 +12,21 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 
 import colossalai
 from colossalai.logging import disable_existing_loggers, get_dist_logger
-from colossalai.nn.parallel import ZeroDDP
+from colossalai.nn.optimizer import HybridAdam
+from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper
 from colossalai.tensor import ColoParameter, ComputePattern, ComputeSpec, ProcessGroup, ReplicaSpec, ShardSpec
 from colossalai.utils import get_current_device
 from colossalai.utils.model.colo_init_context import ColoInitContext
 
 CAI_VERSION = colossalai.__version__
 
-if version.parse(CAI_VERSION) > version.parse("0.1.10"):
-    # These are added after 0.1.10
-    from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer
-    from colossalai.nn.parallel import GeminiDDP
-    from colossalai.zero.sharded_optim import LowLevelZeroOptimizer
-
 
 def parse_args():
     parser = colossalai.get_default_parser()
     parser.add_argument(
         "--distplan",
         type=str,
-        default='colossalai',
+        default='CAI_Gemini',
         help="The distributed plan [colossalai, zero1, zero2, torch_ddp, torch_zero].",
     )
     parser.add_argument(
@@ -48,8 +43,7 @@ def parse_args():
     )
     parser.add_argument(
         "--shardinit",
-        type=bool,
-        default=False,
+        action='store_true',
         help=
         "Shard the tensors when init the model to shrink peak memory size on the assigned device. Valid when using colossalai as dist plan.",
     )
@@ -186,57 +180,16 @@ def tensor_parallelize(model: torch.nn.Module, pg: ProcessGroup):
             param.visited = True
 
 
-# Gemini + ZeRO DDP
-def build_gemini(model: torch.nn.Module, pg: ProcessGroup, placement_policy: str = "auto", ddp_flag: bool = True):
-    fp16_init_scale = 2**5
-    gpu_margin_mem_ratio_for_auto = 0
-
-    if version.parse(CAI_VERSION) > version.parse("0.1.10"):
-        model = GeminiDDP(model,
-                          strict_ddp_mode=ddp_flag,
-                          device=get_current_device(),
-                          placement_policy=placement_policy,
-                          pin_memory=True,
-                          hidden_dim=model.config.n_embd,
-                          search_range_mb=128)
-        # configure the const policy
-        if placement_policy == 'const':
-            model.gemini_manager._placement_policy.set_const_memory_boundary(2 * 1024)
-        # build a highly optimized cpu optimizer
-        optimizer = GeminiAdamOptimizer(model,
-                                        lr=1e-3,
-                                        initial_scale=fp16_init_scale,
-                                        gpu_margin_mem_ratio=gpu_margin_mem_ratio_for_auto)
-    elif version.parse("0.1.9") <= version.parse(CAI_VERSION) <= version.parse("0.1.10"):
-        from colossalai.gemini import ChunkManager, GeminiManager
-        from colossalai.nn.optimizer import HybridAdam
-        from colossalai.zero import ZeroOptimizer
-        chunk_size = ChunkManager.search_chunk_size(model, 64 * 1024**2, 1024, filter_exlarge_params=True)
-        chunk_manager = ChunkManager(chunk_size,
-                                     pg,
-                                     enable_distributed_storage=True,
-                                     init_device=GeminiManager.get_default_device(placement_policy))
-        gemini_manager = GeminiManager(placement_policy, chunk_manager)
-        model = ZeroDDP(model, gemini_manager)
-        optimizer = HybridAdam(model.parameters(), lr=1e-3)
-        optimizer = ZeroOptimizer(optimizer,
-                                  model,
-                                  initial_scale=fp16_init_scale,
-                                  gpu_margin_mem_ratio=gpu_margin_mem_ratio_for_auto)
-    else:
-        raise NotImplemented(f"CAI version {CAI_VERSION} is not supported")
-    return model, optimizer
-
-
 def main():
     # version check
-    # this example is supposed to work for versions greater than 0.1.9
-    assert version.parse(CAI_VERSION) >= version.parse("0.1.9")
+    # this example is supposed to work for versions greater than 0.2.0
+    assert version.parse(CAI_VERSION) >= version.parse("0.2.0")
 
     set_cpu_maximum_parallelism()
     args = parse_args()
 
-    if args.distplan not in ["colossalai", "torch_ddp", "torch_zero", "zero1", "zero2"]:
+    # if args.distplan not in ["colossalai", "torch_ddp", "torch_zero", "zero1", "zero2"]:
+    if args.distplan not in ["CAI_ZeRO1", "CAI_ZeRO2", "CAI_Gemini", "Pytorch_DDP", "Pytorch_ZeRO"]:
         raise TypeError(f"{args.distplan} is error")
 
     # batch size per DP degree
@@ -260,22 +213,21 @@ def main():
     criterion = GPTLMLoss()
 
     torch.manual_seed(123)
-    if args.distplan == "colossalai":
+    if args.distplan.startswith("CAI"):
         # all param must use the same process group.
         world_size = torch.distributed.get_world_size()
         shard_pg = ProcessGroup(tp_degree=world_size) if args.shardinit else None
         default_dist_spec = ShardSpec([-1], [world_size]) if args.shardinit else None
 
+        if args.shardinit and args.distplan != "CAI_Gemini":
+            raise RuntimeError("You can only use shardinit with CAI_Gemini")
+
         # build GPT model
-        if version.parse(CAI_VERSION) > version.parse("0.1.10"):
-            with ColoInitContext(device=get_current_device(),
-                                 dtype=torch.half,
-                                 default_dist_spec=default_dist_spec,
-                                 default_pg=shard_pg):
-                model = model_builder(args.model_type)(checkpoint=True)
-        else:
-            with ColoInitContext(device=get_current_device()):
-                model = model_builder(args.model_type)(checkpoint=True)
+        with ColoInitContext(device=get_current_device(),
+                             dtype=torch.half,
+                             default_dist_spec=default_dist_spec,
+                             default_pg=shard_pg):
+            model = model_builder(args.model_type)(checkpoint=True)
 
         tp_pg = ProcessGroup(tp_degree=args.tp_degree)
         # Tensor Parallelism (TP)
@@ -283,34 +235,49 @@ def main():
         if args.tp_degree > 1:
             tensor_parallelize(model, tp_pg)
 
-        # build a Gemini model and a highly optimized cpu optimizer
-        # Gemini + ZeRO DP, Note it must be used after TP
-        model, optimizer = build_gemini(model, tp_pg, args.placement, args.tp_degree == 1)
+        # asign running configurations
+        gemini_config = None
+        if args.distplan.startswith("CAI_ZeRO"):
+            optim_config = dict(reduce_bucket_size=12 * 1024 * 1024, overlap_communication=True, verbose=True)
+        elif args.distplan == "CAI_Gemini":
+            gemini_config = dict(strict_ddp_mode=args.tp_degree == 1,
+                                 device=get_current_device(),
+                                 placement_policy=args.placement,
+                                 pin_memory=True,
+                                 hidden_dim=model.config.n_embd,
+                                 search_range_mb=128)
+            optim_config = dict(gpu_margin_mem_ratio=0.)
+        else:
+            raise RuntimeError
+
+        # build a highly optimized gpu/cpu optimizer
+        optimizer = HybridAdam(model.parameters(), lr=1e-3)
+
+        if args.distplan == "CAI_ZeRO1":
+            zero_stage = 1
+        elif args.distplan == "CAI_ZeRO2":
+            zero_stage = 2
+        elif args.distplan == "CAI_Gemini":
+            zero_stage = 3
+        else:
+            raise RuntimeError
+
+        # wrap your model and optimizer
+        model = zero_model_wrapper(model, zero_stage, gemini_config)
+        optimizer = zero_optim_wrapper(model, optimizer, optim_config=optim_config)
 
         logger.info(get_mem_info(prefix='After init optim, '), ranks=[0])
-    else:
+    elif args.distplan.startswith("Pytorch"):
         assert args.tp_degree == 1, "The degree of TP should be 1 for DDP examples."
         model = model_builder(args.model_type)(checkpoint=True).cuda()
-
-    if args.distplan.startswith("torch"):
         model = DDP(model)
-        if args.distplan.endswith("ddp"):
-            optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
-        elif args.distplan.endswith("zero"):
+        if args.distplan.endswith("DDP"):
+            optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+        elif args.distplan.endswith("ZeRO"):
             from torch.distributed.optim import ZeroRedundancyOptimizer
-            optimizer = ZeroRedundancyOptimizer(model.parameters(), optimizer_class=torch.optim.Adam, lr=0.01)
-    elif args.distplan.startswith("zero"):
-        model = model.half()
-        partition_flag = (args.distplan == "zero2")
-        optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
-
-        optimizer = LowLevelZeroOptimizer(
-            optimizer,
-            reduce_bucket_size=12 * 1024 * 1024,
-            overlap_communication=True,
-            partition_grad=partition_flag,
-            verbose=True,
-        )
+            optimizer = ZeroRedundancyOptimizer(model.parameters(), optimizer_class=torch.optim.Adam, lr=1e-3)
+    else:
+        raise RuntimeError
 
     # model is shared after TP
     numel = get_model_size(model)
@@ -338,17 +305,18 @@ def main():
         fwd_time = fwd_end - start
         logger.info(get_mem_info(prefix=f'[{n + 1}/{NUM_STEPS}] Forward '), ranks=[0])
 
-        if args.distplan in ["colossalai", "zero1", "zero2"]:
+        if args.distplan.startswith("CAI"):
             optimizer.backward(loss)
-        elif args.distplan in ["torch_ddp", "torch_zero"]:
+        elif args.distplan.startswith("Pytorch"):
             loss.backward()
+        else:
+            raise RuntimeError
+
         torch.cuda.synchronize()
         bwd_end = time()
         bwd_time = bwd_end - fwd_end
         logger.info(get_mem_info(prefix=f'[{n + 1}/{NUM_STEPS}] Backward '), ranks=[0])
 
-        if args.distplan in ["zero1", "zero2"]:
-            optimizer.sync_grad()
         optimizer.step()
         torch.cuda.synchronize()
         optim_time = time() - bwd_end