We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/language/opt/train_gemini_opt.py`
Content:
```
1 #!/usr/bin/env python
2 # coding=utf-8
3 # Copyright 2021 The HuggingFace Inc. team. All rights reserved.
4 #
5 # Licensed under the Apache License, Version 2.0 (the "License");
6 # you may not use this file except in compliance with the License.
7 # You may obtain a copy of the License at
8 #
9 #     http://www.apache.org/licenses/LICENSE-2.0
10 #
11 # Unless required by applicable law or agreed to in writing, software
12 # distributed under the License is distributed on an "AS IS" BASIS,
13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 # See the License for the specific language governing permissions and
15 # limitations under the License.
16 """
17 Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...)
18 on a text file or a dataset without using HuggingFace Trainer.
19 
20 Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
21 https://huggingface.co/models?filter=text-generation
22 """
23 # You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.
24 
25 import time
26 from functools import partial
27 
28 import datasets
29 import torch
30 import torch.distributed as dist
31 import transformers
32 from transformers import CONFIG_MAPPING, MODEL_MAPPING, AutoConfig, OPTForCausalLM
33 from transformers.utils.versions import require_version
34 
35 import colossalai
36 from colossalai.logging import disable_existing_loggers, get_dist_logger
37 from colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer
38 from colossalai.nn.parallel import GeminiDDP
39 from colossalai.utils import get_current_device
40 from colossalai.utils.model.colo_init_context import ColoInitContext
41 
42 
43 def get_data(batch_size, seq_len, vocab_size):
44     input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=torch.cuda.current_device())
45     attention_mask = torch.ones_like(input_ids)
46     return input_ids, attention_mask
47 
48 
49 require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")
50 
51 MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
52 MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
53 
54 
55 def get_time_stamp():
56     torch.cuda.synchronize()
57     return time.time()
58 
59 
60 def get_tflops(model_numel, batch_size, seq_len, step_time):
61     return model_numel * batch_size * seq_len * 8 / 1e12 / (step_time + 1e-12)
62 
63 
64 def parse_args():
65     parser = colossalai.get_default_parser()
66     parser.add_argument(
67         "--model_name_or_path",
68         type=str,
69         help="Path to pretrained model or model identifier from huggingface.co/models.",
70         required=True,
71     )
72     parser.add_argument(
73         "--config_name",
74         type=str,
75         default=None,
76         help="Pretrained config name or path if not the same as model_name",
77     )
78     parser.add_argument(
79         "--batch_size",
80         type=int,
81         default=8,
82         help="Batch size (per dp group) for the training dataloader.",
83     )
84     parser.add_argument(
85         "--learning_rate",
86         type=float,
87         default=5e-5,
88         help="Initial learning rate (after the potential warmup period) to use.",
89     )
90     parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.")
91     parser.add_argument(
92         "--max_train_steps",
93         type=int,
94         default=20,
95         help="Total number of training steps to perform.",
96     )
97     parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
98     parser.add_argument(
99         "--model_type",
100         type=str,
101         default=None,
102         help="Model type to use if training from scratch.",
103         choices=MODEL_TYPES,
104     )
105     parser.add_argument("--mem_cap", type=int, default=0, help="use mem cap")
106     parser.add_argument("--init_in_cpu", action='store_true', default=False, help="init training model in cpu")
107     args = parser.parse_args()
108 
109     return args
110 
111 
112 def colo_memory_cap(size_in_GB):
113     from colossalai.utils import colo_device_memory_capacity, colo_set_process_memory_fraction, get_current_device
114     cuda_capacity = colo_device_memory_capacity(get_current_device())
115     if size_in_GB * (1024**3) < cuda_capacity:
116         colo_set_process_memory_fraction(size_in_GB * (1024**3) / cuda_capacity)
117         print("Using {} GB of GPU memory".format(size_in_GB))
118 
119 
120 def main():
121     args = parse_args()
122     disable_existing_loggers()
123     colossalai.launch_from_torch({})
124     logger = get_dist_logger()
125     is_main_process = dist.get_rank() == 0
126 
127     if is_main_process:
128         datasets.utils.logging.set_verbosity_warning()
129         transformers.utils.logging.set_verbosity_info()
130     else:
131         datasets.utils.logging.set_verbosity_error()
132         transformers.utils.logging.set_verbosity_error()
133 
134     if args.mem_cap > 0:
135         colo_memory_cap(args.mem_cap)
136 
137     # If passed along, set the training seed now.
138     if args.seed is not None:
139         torch.mannul_seed(args.seed)
140         logger.info(f"Rank {dist.get_rank()}: random seed is set to {args.seed}")
141 
142     # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
143     # https://huggingface.co/docs/datasets/loading_datasets.html.
144 
145     # Load pretrained model
146     # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently
147     # download model & vocab.
148     if args.config_name:
149         config = AutoConfig.from_pretrained(args.config_name)
150     elif args.model_name_or_path:
151         config = AutoConfig.from_pretrained(args.model_name_or_path)
152     else:
153         config = CONFIG_MAPPING[args.model_type]()
154         logger.warning("You are instantiating a new config instance from scratch.")
155     logger.info("Model config has been created", ranks=[0])
156 
157     if args.init_in_cpu:
158         init_dev = torch.device('cpu')
159     else:
160         init_dev = get_current_device()
161 
162     # build model
163     if args.model_name_or_path is None or args.model_name_or_path == 'facebook/opt-13b':
164         # currently, there has a bug in pretrained opt-13b
165         # we can not import it until huggingface fix it
166         logger.info("Train a new model from scratch", ranks=[0])
167         with ColoInitContext(device=init_dev, dtype=torch.half):
168             model = OPTForCausalLM(config)
169     else:
170         logger.info("Finetune a pre-trained model", ranks=[0])
171         with ColoInitContext(device=init_dev, dtype=torch.half):
172             model = OPTForCausalLM.from_pretrained(args.model_name_or_path,
173                                                    from_tf=bool(".ckpt" in args.model_name_or_path),
174                                                    config=config,
175                                                    local_files_only=False)
176 
177     # enable graident checkpointing
178     model.gradient_checkpointing_enable()
179 
180     numel = sum([p.numel() for p in model.parameters()])
181     PLACEMENT_POLICY = 'cpu'
182     model = GeminiDDP(model, device=get_current_device(), placement_policy=PLACEMENT_POLICY, pin_memory=True)
183     optimizer = GeminiAdamOptimizer(model, lr=args.learning_rate, initial_scale=2**14, gpu_margin_mem_ratio=0.0)
184 
185     SEQ_LEN = 1024
186     VOCAB_SIZE = 50257
187 
188     get_tflops_func = partial(get_tflops, numel, args.batch_size, SEQ_LEN)
189 
190     model.train()
191     for step in range(args.max_train_steps):
192         st_time = time.time()
193         input_ids, attn_mask = get_data(args.batch_size, SEQ_LEN, VOCAB_SIZE)
194 
195         outputs = model(input_ids=input_ids, attention_mask=attn_mask, labels=input_ids, use_cache=False)
196         loss = outputs['loss']
197         optimizer.backward(loss)
198 
199         optimizer.step()
200         optimizer.zero_grad()
201         torch.cuda.synchronize()
202         step_time = time.time() - st_time
203         step_tflops = get_tflops_func(step_time)
204 
205         logger.info("step {} finished, Tflops {}".format(step, step_tflops), ranks=[0])
206 
207     logger.info("Training finished", ranks=[0])
208 
209 
210 if __name__ == "__main__":
211     main()
212
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         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 @@
     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 @@
 
     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

{"golden_diff": "diff --git a/examples/language/opt/train_gemini_opt.py b/examples/language/opt/train_gemini_opt.py\n--- a/examples/language/opt/train_gemini_opt.py\n+++ b/examples/language/opt/train_gemini_opt.py\n@@ -39,6 +39,8 @@\n from colossalai.utils import get_current_device\n from colossalai.utils.model.colo_init_context import ColoInitContext\n \n+from colossalai.tensor import ProcessGroup, ShardSpec\n+\n \n def get_data(batch_size, seq_len, vocab_size):\n     input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=torch.cuda.current_device())\n@@ -102,6 +104,11 @@\n         help=\"Model type to use if training from scratch.\",\n         choices=MODEL_TYPES,\n     )\n+    parser.add_argument(\n+        \"--shardinit\",\n+        action=\"store_true\",\n+        help=\"Initialize the model with tensor parallel\",\n+    )\n     parser.add_argument(\"--mem_cap\", type=int, default=0, help=\"use mem cap\")\n     parser.add_argument(\"--init_in_cpu\", action='store_true', default=False, help=\"init training model in cpu\")\n     args = parser.parse_args()\n@@ -159,16 +166,30 @@\n     else:\n         init_dev = get_current_device()\n \n+    # shard init prameters\n+    if args.shardinit:\n+        logger.info(\"Sharding initialization !\", ranks=[0])\n+    else:\n+        logger.info(\"Skipping sharding initialization\", ranks=[0])\n+\n+    world_size = torch.distributed.get_world_size()\n+    shard_pg = ProcessGroup(tp_degree=world_size) if args.shardinit else None\n+    default_dist_spec = ShardSpec([-1], [world_size]) if args.shardinit else None\n+\n     # build model\n     if args.model_name_or_path is None or args.model_name_or_path == 'facebook/opt-13b':\n         # currently, there has a bug in pretrained opt-13b\n         # we can not import it until huggingface fix it\n         logger.info(\"Train a new model from scratch\", ranks=[0])\n-        with ColoInitContext(device=init_dev, dtype=torch.half):\n+        with ColoInitContext(device=init_dev, dtype=torch.half,\n+                             default_dist_spec=default_dist_spec,\n+                             default_pg=shard_pg):\n             model = OPTForCausalLM(config)\n     else:\n         logger.info(\"Finetune a pre-trained model\", ranks=[0])\n-        with ColoInitContext(device=init_dev, dtype=torch.half):\n+        with ColoInitContext(device=init_dev, dtype=torch.half,\n+                             default_dist_spec=default_dist_spec,\n+                             default_pg=shard_pg):\n             model = OPTForCausalLM.from_pretrained(args.model_name_or_path,\n                                                    from_tf=bool(\".ckpt\" in args.model_name_or_path),\n                                                    config=config,\n@@ -179,7 +200,8 @@\n \n     numel = sum([p.numel() for p in model.parameters()])\n     PLACEMENT_POLICY = 'cpu'\n-    model = GeminiDDP(model, device=get_current_device(), placement_policy=PLACEMENT_POLICY, pin_memory=True)\n+    model = GeminiDDP(model, device=get_current_device(), placement_policy=PLACEMENT_POLICY,\n+                      pin_memory=True, strict_ddp_mode=args.shardinit)\n     optimizer = GeminiAdamOptimizer(model, lr=args.learning_rate, initial_scale=2**14, gpu_margin_mem_ratio=0.0)\n \n     SEQ_LEN = 1024\n", "issue": "[BUG]: \u75283\u5f20A100\u663e\u5361\u4f7f\u7528examples\u7684opt demo\u4f1a\u7206\u663e\u5b58\n### \ud83d\udc1b Describe the bug\n\n\u4f7f\u7528\u7684\u6a21\u578b\u662fopt-66b\uff0c\u4f7f\u7528\u7684\u662f80gb\u663e\u5b58\u7684A100\u663e\u5361\r\n\u4f7f\u7528\u7684\u811a\u672cbash ./run_gemini.sh 1 0 66b 3\r\n\u6211\u4eec\u76f4\u63a5\u8dd1\u811a\u672c\u4e4b\u540e\uff0c\u5728\u52a0\u8f7d\u9884\u8bad\u7ec3\u6a21\u578b\u65f6\uff0c3\u5f20\u663e\u5361\u4f1a\u540c\u65f6\u52a0\u523080gb\uff0c\u7136\u540e\u62a5CUDA out of memory\r\n\u540e\u6765\u6539\u6210\u6a21\u578b\u52a0\u8f7d\u5230\u5185\u5b58\u7684\u65f6\u5019\u4f7f\u7528\u4e86400gb\u7684\u5185\u5b58\uff0c\u7136\u540e\u7206\u53e6\u5916\u7684\u9519\u8bef\r\nDetected CUDA files, patching ldflags\r\nEmitting ninja build file /root/.cache/colossalai/torch_extensions/torch1.10_cu11.3/build.ninja...\r\nBuilding extension module fused_optim...\r\nAllowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)\r\nninja: no work to do.\r\nLoading extension module fused_optim...\r\nTime to load fused_optim op: 1.2893073558807373 seconds\r\n[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                  \r\n[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                  \r\n[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                  \r\n                    INFO     colossalai - colossalai - INFO: Found overflow. Skip step                                                                                        \r\n                    INFO     colossalai - colossalai - INFO: Found overflow. Skip step                                                                                        \r\n                    INFO     colossalai - colossalai - INFO: Found overflow. Skip step                                                                                        \r\n                    INFO     colossalai - colossalai - INFO: train_gemini_opt.py:205 main                                                                                     \r\n                    INFO     colossalai - colossalai - INFO: step 0 finished, Tflops 25.13966826409652                                                                        \r\nWARNING:torch.distributed.elastic.multiprocessing.api:Sending process 505485 closing signal SIGTERM\r\nWARNING:torch.distributed.elastic.multiprocessing.api:Sending process 505487 closing signal SIGTERM\r\nWARNING:torch.distributed.elastic.multiprocessing.api:Unable to shutdown process 505487 via 15, forcefully exitting via 9\r\nERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: -9) local_rank: 1 (pid: 505486) of binary: /root/miniconda3/bin/python\r\nTraceback (most recent call last):\r\n  File \"/root/miniconda3/bin/torchrun\", line 8, in <module>\r\n    sys.exit(main())\r\n  File \"/root/miniconda3/lib/python3.8/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py\", line 345, in wrapper\r\n    return f(*args, **kwargs)\r\n  File \"/root/miniconda3/lib/python3.8/site-packages/torch/distributed/run.py\", line 719, in main\r\n    run(args)\r\n  File \"/root/miniconda3/lib/python3.8/site-packages/torch/distributed/run.py\", line 710, in run\r\n    elastic_launch(\r\n  File \"/root/miniconda3/lib/python3.8/site-packages/torch/distributed/launcher/api.py\", line 131, in __call__\r\n    return launch_agent(self._config, self._entrypoint, list(args))\r\n  File \"/root/miniconda3/lib/python3.8/site-packages/torch/distributed/launcher/api.py\", line 259, in launch_agent\r\n    raise ChildFailedError(\r\ntorch.distributed.elastic.multiprocessing.errors.ChildFailedError: \r\n=======================================================\r\ntrain_gemini_opt.py FAILED\r\n-------------------------------------------------------\r\n\r\n\n\n### Environment\n\n_No response_\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "#!/usr/bin/env python\n# coding=utf-8\n# Copyright 2021 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nFine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...)\non a text file or a dataset without using HuggingFace Trainer.\n\nHere is the full list of checkpoints on the hub that can be fine-tuned by this script:\nhttps://huggingface.co/models?filter=text-generation\n\"\"\"\n# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.\n\nimport time\nfrom functools import partial\n\nimport datasets\nimport torch\nimport torch.distributed as dist\nimport transformers\nfrom transformers import CONFIG_MAPPING, MODEL_MAPPING, AutoConfig, OPTForCausalLM\nfrom transformers.utils.versions import require_version\n\nimport colossalai\nfrom colossalai.logging import disable_existing_loggers, get_dist_logger\nfrom colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer\nfrom colossalai.nn.parallel import GeminiDDP\nfrom colossalai.utils import get_current_device\nfrom colossalai.utils.model.colo_init_context import ColoInitContext\n\n\ndef get_data(batch_size, seq_len, vocab_size):\n    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=torch.cuda.current_device())\n    attention_mask = torch.ones_like(input_ids)\n    return input_ids, attention_mask\n\n\nrequire_version(\"datasets>=1.8.0\", \"To fix: pip install -r examples/pytorch/language-modeling/requirements.txt\")\n\nMODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())\nMODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)\n\n\ndef get_time_stamp():\n    torch.cuda.synchronize()\n    return time.time()\n\n\ndef get_tflops(model_numel, batch_size, seq_len, step_time):\n    return model_numel * batch_size * seq_len * 8 / 1e12 / (step_time + 1e-12)\n\n\ndef parse_args():\n    parser = colossalai.get_default_parser()\n    parser.add_argument(\n        \"--model_name_or_path\",\n        type=str,\n        help=\"Path to pretrained model or model identifier from huggingface.co/models.\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--config_name\",\n        type=str,\n        default=None,\n        help=\"Pretrained config name or path if not the same as model_name\",\n    )\n    parser.add_argument(\n        \"--batch_size\",\n        type=int,\n        default=8,\n        help=\"Batch size (per dp group) for the training dataloader.\",\n    )\n    parser.add_argument(\n        \"--learning_rate\",\n        type=float,\n        default=5e-5,\n        help=\"Initial learning rate (after the potential warmup period) to use.\",\n    )\n    parser.add_argument(\"--weight_decay\", type=float, default=0.0, help=\"Weight decay to use.\")\n    parser.add_argument(\n        \"--max_train_steps\",\n        type=int,\n        default=20,\n        help=\"Total number of training steps to perform.\",\n    )\n    parser.add_argument(\"--seed\", type=int, default=None, help=\"A seed for reproducible training.\")\n    parser.add_argument(\n        \"--model_type\",\n        type=str,\n        default=None,\n        help=\"Model type to use if training from scratch.\",\n        choices=MODEL_TYPES,\n    )\n    parser.add_argument(\"--mem_cap\", type=int, default=0, help=\"use mem cap\")\n    parser.add_argument(\"--init_in_cpu\", action='store_true', default=False, help=\"init training model in cpu\")\n    args = parser.parse_args()\n\n    return args\n\n\ndef colo_memory_cap(size_in_GB):\n    from colossalai.utils import colo_device_memory_capacity, colo_set_process_memory_fraction, get_current_device\n    cuda_capacity = colo_device_memory_capacity(get_current_device())\n    if size_in_GB * (1024**3) < cuda_capacity:\n        colo_set_process_memory_fraction(size_in_GB * (1024**3) / cuda_capacity)\n        print(\"Using {} GB of GPU memory\".format(size_in_GB))\n\n\ndef main():\n    args = parse_args()\n    disable_existing_loggers()\n    colossalai.launch_from_torch({})\n    logger = get_dist_logger()\n    is_main_process = dist.get_rank() == 0\n\n    if is_main_process:\n        datasets.utils.logging.set_verbosity_warning()\n        transformers.utils.logging.set_verbosity_info()\n    else:\n        datasets.utils.logging.set_verbosity_error()\n        transformers.utils.logging.set_verbosity_error()\n\n    if args.mem_cap > 0:\n        colo_memory_cap(args.mem_cap)\n\n    # If passed along, set the training seed now.\n    if args.seed is not None:\n        torch.mannul_seed(args.seed)\n        logger.info(f\"Rank {dist.get_rank()}: random seed is set to {args.seed}\")\n\n    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at\n    # https://huggingface.co/docs/datasets/loading_datasets.html.\n\n    # Load pretrained model\n    # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently\n    # download model & vocab.\n    if args.config_name:\n        config = AutoConfig.from_pretrained(args.config_name)\n    elif args.model_name_or_path:\n        config = AutoConfig.from_pretrained(args.model_name_or_path)\n    else:\n        config = CONFIG_MAPPING[args.model_type]()\n        logger.warning(\"You are instantiating a new config instance from scratch.\")\n    logger.info(\"Model config has been created\", ranks=[0])\n\n    if args.init_in_cpu:\n        init_dev = torch.device('cpu')\n    else:\n        init_dev = get_current_device()\n\n    # build model\n    if args.model_name_or_path is None or args.model_name_or_path == 'facebook/opt-13b':\n        # currently, there has a bug in pretrained opt-13b\n        # we can not import it until huggingface fix it\n        logger.info(\"Train a new model from scratch\", ranks=[0])\n        with ColoInitContext(device=init_dev, dtype=torch.half):\n            model = OPTForCausalLM(config)\n    else:\n        logger.info(\"Finetune a pre-trained model\", ranks=[0])\n        with ColoInitContext(device=init_dev, dtype=torch.half):\n            model = OPTForCausalLM.from_pretrained(args.model_name_or_path,\n                                                   from_tf=bool(\".ckpt\" in args.model_name_or_path),\n                                                   config=config,\n                                                   local_files_only=False)\n\n    # enable graident checkpointing\n    model.gradient_checkpointing_enable()\n\n    numel = sum([p.numel() for p in model.parameters()])\n    PLACEMENT_POLICY = 'cpu'\n    model = GeminiDDP(model, device=get_current_device(), placement_policy=PLACEMENT_POLICY, pin_memory=True)\n    optimizer = GeminiAdamOptimizer(model, lr=args.learning_rate, initial_scale=2**14, gpu_margin_mem_ratio=0.0)\n\n    SEQ_LEN = 1024\n    VOCAB_SIZE = 50257\n\n    get_tflops_func = partial(get_tflops, numel, args.batch_size, SEQ_LEN)\n\n    model.train()\n    for step in range(args.max_train_steps):\n        st_time = time.time()\n        input_ids, attn_mask = get_data(args.batch_size, SEQ_LEN, VOCAB_SIZE)\n\n        outputs = model(input_ids=input_ids, attention_mask=attn_mask, labels=input_ids, use_cache=False)\n        loss = outputs['loss']\n        optimizer.backward(loss)\n\n        optimizer.step()\n        optimizer.zero_grad()\n        torch.cuda.synchronize()\n        step_time = time.time() - st_time\n        step_tflops = get_tflops_func(step_time)\n\n        logger.info(\"step {} finished, Tflops {}\".format(step, step_tflops), ranks=[0])\n\n    logger.info(\"Training finished\", ranks=[0])\n\n\nif __name__ == \"__main__\":\n    main()\n", "path": "examples/language/opt/train_gemini_opt.py"}], "after_files": [{"content": "#!/usr/bin/env python\n# coding=utf-8\n# Copyright 2021 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nFine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...)\non a text file or a dataset without using HuggingFace Trainer.\n\nHere is the full list of checkpoints on the hub that can be fine-tuned by this script:\nhttps://huggingface.co/models?filter=text-generation\n\"\"\"\n# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.\n\nimport time\nfrom functools import partial\n\nimport datasets\nimport torch\nimport torch.distributed as dist\nimport transformers\nfrom transformers import CONFIG_MAPPING, MODEL_MAPPING, AutoConfig, OPTForCausalLM\nfrom transformers.utils.versions import require_version\n\nimport colossalai\nfrom colossalai.logging import disable_existing_loggers, get_dist_logger\nfrom colossalai.nn.optimizer.gemini_optimizer import GeminiAdamOptimizer\nfrom colossalai.nn.parallel import GeminiDDP\nfrom colossalai.utils import get_current_device\nfrom colossalai.utils.model.colo_init_context import ColoInitContext\n\nfrom colossalai.tensor import ProcessGroup, ShardSpec\n\n\ndef get_data(batch_size, seq_len, vocab_size):\n    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=torch.cuda.current_device())\n    attention_mask = torch.ones_like(input_ids)\n    return input_ids, attention_mask\n\n\nrequire_version(\"datasets>=1.8.0\", \"To fix: pip install -r examples/pytorch/language-modeling/requirements.txt\")\n\nMODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())\nMODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)\n\n\ndef get_time_stamp():\n    torch.cuda.synchronize()\n    return time.time()\n\n\ndef get_tflops(model_numel, batch_size, seq_len, step_time):\n    return model_numel * batch_size * seq_len * 8 / 1e12 / (step_time + 1e-12)\n\n\ndef parse_args():\n    parser = colossalai.get_default_parser()\n    parser.add_argument(\n        \"--model_name_or_path\",\n        type=str,\n        help=\"Path to pretrained model or model identifier from huggingface.co/models.\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--config_name\",\n        type=str,\n        default=None,\n        help=\"Pretrained config name or path if not the same as model_name\",\n    )\n    parser.add_argument(\n        \"--batch_size\",\n        type=int,\n        default=8,\n        help=\"Batch size (per dp group) for the training dataloader.\",\n    )\n    parser.add_argument(\n        \"--learning_rate\",\n        type=float,\n        default=5e-5,\n        help=\"Initial learning rate (after the potential warmup period) to use.\",\n    )\n    parser.add_argument(\"--weight_decay\", type=float, default=0.0, help=\"Weight decay to use.\")\n    parser.add_argument(\n        \"--max_train_steps\",\n        type=int,\n        default=20,\n        help=\"Total number of training steps to perform.\",\n    )\n    parser.add_argument(\"--seed\", type=int, default=None, help=\"A seed for reproducible training.\")\n    parser.add_argument(\n        \"--model_type\",\n        type=str,\n        default=None,\n        help=\"Model type to use if training from scratch.\",\n        choices=MODEL_TYPES,\n    )\n    parser.add_argument(\n        \"--shardinit\",\n        action=\"store_true\",\n        help=\"Initialize the model with tensor parallel\",\n    )\n    parser.add_argument(\"--mem_cap\", type=int, default=0, help=\"use mem cap\")\n    parser.add_argument(\"--init_in_cpu\", action='store_true', default=False, help=\"init training model in cpu\")\n    args = parser.parse_args()\n\n    return args\n\n\ndef colo_memory_cap(size_in_GB):\n    from colossalai.utils import colo_device_memory_capacity, colo_set_process_memory_fraction, get_current_device\n    cuda_capacity = colo_device_memory_capacity(get_current_device())\n    if size_in_GB * (1024**3) < cuda_capacity:\n        colo_set_process_memory_fraction(size_in_GB * (1024**3) / cuda_capacity)\n        print(\"Using {} GB of GPU memory\".format(size_in_GB))\n\n\ndef main():\n    args = parse_args()\n    disable_existing_loggers()\n    colossalai.launch_from_torch({})\n    logger = get_dist_logger()\n    is_main_process = dist.get_rank() == 0\n\n    if is_main_process:\n        datasets.utils.logging.set_verbosity_warning()\n        transformers.utils.logging.set_verbosity_info()\n    else:\n        datasets.utils.logging.set_verbosity_error()\n        transformers.utils.logging.set_verbosity_error()\n\n    if args.mem_cap > 0:\n        colo_memory_cap(args.mem_cap)\n\n    # If passed along, set the training seed now.\n    if args.seed is not None:\n        torch.mannul_seed(args.seed)\n        logger.info(f\"Rank {dist.get_rank()}: random seed is set to {args.seed}\")\n\n    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at\n    # https://huggingface.co/docs/datasets/loading_datasets.html.\n\n    # Load pretrained model\n    # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently\n    # download model & vocab.\n    if args.config_name:\n        config = AutoConfig.from_pretrained(args.config_name)\n    elif args.model_name_or_path:\n        config = AutoConfig.from_pretrained(args.model_name_or_path)\n    else:\n        config = CONFIG_MAPPING[args.model_type]()\n        logger.warning(\"You are instantiating a new config instance from scratch.\")\n    logger.info(\"Model config has been created\", ranks=[0])\n\n    if args.init_in_cpu:\n        init_dev = torch.device('cpu')\n    else:\n        init_dev = get_current_device()\n\n    # shard init prameters\n    if args.shardinit:\n        logger.info(\"Sharding initialization !\", ranks=[0])\n    else:\n        logger.info(\"Skipping sharding initialization\", ranks=[0])\n\n    world_size = torch.distributed.get_world_size()\n    shard_pg = ProcessGroup(tp_degree=world_size) if args.shardinit else None\n    default_dist_spec = ShardSpec([-1], [world_size]) if args.shardinit else None\n\n    # build model\n    if args.model_name_or_path is None or args.model_name_or_path == 'facebook/opt-13b':\n        # currently, there has a bug in pretrained opt-13b\n        # we can not import it until huggingface fix it\n        logger.info(\"Train a new model from scratch\", ranks=[0])\n        with ColoInitContext(device=init_dev, dtype=torch.half,\n                             default_dist_spec=default_dist_spec,\n                             default_pg=shard_pg):\n            model = OPTForCausalLM(config)\n    else:\n        logger.info(\"Finetune a pre-trained model\", ranks=[0])\n        with ColoInitContext(device=init_dev, dtype=torch.half,\n                             default_dist_spec=default_dist_spec,\n                             default_pg=shard_pg):\n            model = OPTForCausalLM.from_pretrained(args.model_name_or_path,\n                                                   from_tf=bool(\".ckpt\" in args.model_name_or_path),\n                                                   config=config,\n                                                   local_files_only=False)\n\n    # enable graident checkpointing\n    model.gradient_checkpointing_enable()\n\n    numel = sum([p.numel() for p in model.parameters()])\n    PLACEMENT_POLICY = 'cpu'\n    model = GeminiDDP(model, device=get_current_device(), placement_policy=PLACEMENT_POLICY,\n                      pin_memory=True, strict_ddp_mode=args.shardinit)\n    optimizer = GeminiAdamOptimizer(model, lr=args.learning_rate, initial_scale=2**14, gpu_margin_mem_ratio=0.0)\n\n    SEQ_LEN = 1024\n    VOCAB_SIZE = 50257\n\n    get_tflops_func = partial(get_tflops, numel, args.batch_size, SEQ_LEN)\n\n    model.train()\n    for step in range(args.max_train_steps):\n        st_time = time.time()\n        input_ids, attn_mask = get_data(args.batch_size, SEQ_LEN, VOCAB_SIZE)\n\n        outputs = model(input_ids=input_ids, attention_mask=attn_mask, labels=input_ids, use_cache=False)\n        loss = outputs['loss']\n        optimizer.backward(loss)\n\n        optimizer.step()\n        optimizer.zero_grad()\n        torch.cuda.synchronize()\n        step_time = time.time() - st_time\n        step_tflops = get_tflops_func(step_time)\n\n        logger.info(\"step {} finished, Tflops {}\".format(step, step_tflops), ranks=[0])\n\n    logger.info(\"Training finished\", ranks=[0])\n\n\nif __name__ == \"__main__\":\n    main()\n", "path": "examples/language/opt/train_gemini_opt.py"}]}