Repository: hpcaitech/ColossalAI

Issue:
[BUG]: 用3张A100显卡使用examples的opt demo会爆显存
### 🐛 Describe the bug

使用的模型是opt-66b，使用的是80gb显存的A100显卡
使用的脚本bash ./run_gemini.sh 1 0 66b 3
我们直接跑脚本之后，在加载预训练模型时，3张显卡会同时加到80gb，然后报CUDA out of memory
后来改成模型加载到内存的时候使用了400gb的内存，然后爆另外的错误
Detected CUDA files, patching ldflags
Emitting ninja build file /root/.cache/colossalai/torch_extensions/torch1.10_cu11.3/build.ninja...
Building extension module fused_optim...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module fused_optim...
Time to load fused_optim op: 1.2893073558807373 seconds
[02/16/23 15:09:50] INFO     colossalai - colossalai - INFO: /root/miniconda3/lib/python3.8/site-packages/colossalai/nn/optimizer/zero_optimizer.py:217 step                  
[02/16/23 15:09:50] INFO     colossalai - colossalai - INFO: /root/miniconda3/lib/python3.8/site-packages/colossalai/nn/optimizer/zero_optimizer.py:217 step                  
[02/16/23 15:09:50] INFO     colossalai - colossalai - INFO: /root/miniconda3/lib/python3.8/site-packages/colossalai/nn/optimizer/zero_optimizer.py:217 step                  
                    INFO     colossalai - colossalai - INFO: Found overflow. Skip step                                                                                        
                    INFO     colossalai - colossalai - INFO: Found overflow. Skip step                                                                                        
                    INFO     colossalai - colossalai - INFO: Found overflow. Skip step                                                                                        
                    INFO     colossalai - colossalai - INFO: train_gemini_opt.py:205 main                                                                                     
                    INFO     colossalai - colossalai - INFO: step 0 finished, Tflops 25.13966826409652                                                                        
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 505485 closing signal SIGTERM
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 505487 closing signal SIGTERM
WARNING:torch.distributed.elastic.multiprocessing.api:Unable to shutdown process 505487 via 15, forcefully exitting via 9
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: -9) local_rank: 1 (pid: 505486) of binary: /root/miniconda3/bin/python
Traceback (most recent call last):
  File "/root/miniconda3/bin/torchrun", line 8, in <module>
    sys.exit(main())
  File "/root/miniconda3/lib/python3.8/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 345, in wrapper
    return f(*args, **kwargs)
  File "/root/miniconda3/lib/python3.8/site-packages/torch/distributed/run.py", line 719, in main
    run(args)
  File "/root/miniconda3/lib/python3.8/site-packages/torch/distributed/run.py", line 710, in run
    elastic_launch(
  File "/root/miniconda3/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 131, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/root/miniconda3/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 259, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
=======================================================
train_gemini_opt.py FAILED
-------------------------------------------------------



### Environment

_No response_
[tensor] fix some unittests

[tensor] fix some unittests



Discussion:
 
> Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿      
----  
**Title:** [BUG]: Using the opt demo of examples with 3 A100 graphics cards will burst the video memory    
      
I have validated this issue: in both cases (loaded to cuda memory and cpu memory), there was OOM. 
Would u please try a smaller model for now, such as 1.3b? 




Resolving patch:
diff --git a/examples/language/opt/train_gemini_opt.py b/examples/language/opt/train_gemini_opt.py
--- a/examples/language/opt/train_gemini_opt.py
+++ b/examples/language/opt/train_gemini_opt.py
@@ -39,6 +39,8 @@
 from colossalai.utils import get_current_device
 from colossalai.utils.model.colo_init_context import ColoInitContext
 
+from colossalai.tensor import ProcessGroup, ShardSpec
+
 
 def get_data(batch_size, seq_len, vocab_size):
     input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=torch.cuda.current_device())
@@ -102,6 +104,11 @@ def parse_args():
         help="Model type to use if training from scratch.",
         choices=MODEL_TYPES,
     )
+    parser.add_argument(
+        "--shardinit",
+        action="store_true",
+        help="Initialize the model with tensor parallel",
+    )
     parser.add_argument("--mem_cap", type=int, default=0, help="use mem cap")
     parser.add_argument("--init_in_cpu", action='store_true', default=False, help="init training model in cpu")
     args = parser.parse_args()
@@ -159,16 +166,30 @@ def main():
     else:
         init_dev = get_current_device()
 
+    # shard init prameters
+    if args.shardinit:
+        logger.info("Sharding initialization !", ranks=[0])
+    else:
+        logger.info("Skipping sharding initialization", ranks=[0])
+
+    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
+
     # build model
     if args.model_name_or_path is None or args.model_name_or_path == 'facebook/opt-13b':
         # currently, there has a bug in pretrained opt-13b
         # we can not import it until huggingface fix it
         logger.info("Train a new model from scratch", ranks=[0])
-        with ColoInitContext(device=init_dev, dtype=torch.half):
+        with ColoInitContext(device=init_dev, dtype=torch.half,
+                             default_dist_spec=default_dist_spec,
+                             default_pg=shard_pg):
             model = OPTForCausalLM(config)
     else:
         logger.info("Finetune a pre-trained model", ranks=[0])
-        with ColoInitContext(device=init_dev, dtype=torch.half):
+        with ColoInitContext(device=init_dev, dtype=torch.half,
+                             default_dist_spec=default_dist_spec,
+                             default_pg=shard_pg):
             model = OPTForCausalLM.from_pretrained(args.model_name_or_path,
                                                    from_tf=bool(".ckpt" in args.model_name_or_path),
                                                    config=config,
@@ -179,7 +200,8 @@ def main():
 
     numel = sum([p.numel() for p in model.parameters()])
     PLACEMENT_POLICY = 'cpu'
-    model = GeminiDDP(model, device=get_current_device(), placement_policy=PLACEMENT_POLICY, pin_memory=True)
+    model = GeminiDDP(model, device=get_current_device(), placement_policy=PLACEMENT_POLICY,
+                      pin_memory=True, strict_ddp_mode=args.shardinit)
     optimizer = GeminiAdamOptimizer(model, lr=args.learning_rate, initial_scale=2**14, gpu_margin_mem_ratio=0.0)
 
     SEQ_LEN = 1024


Test patch: