We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[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: `applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py`
Content:
```
1 import warnings
2 from typing import Optional, Union
3 
4 import torch
5 import torch.distributed as dist
6 import torch.nn as nn
7 import torch.optim as optim
8 from chatgpt.models.base import Actor
9 from chatgpt.models.lora import LoraLinear
10 from torch.optim import Optimizer
11 
12 from transformers.modeling_utils import PreTrainedModel
13 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
14 
15 import colossalai
16 from colossalai.nn.optimizer import CPUAdam, HybridAdam
17 from colossalai.nn.parallel import ZeroDDP, zero_model_wrapper, zero_optim_wrapper
18 from colossalai.nn.parallel.utils import get_static_torch_model
19 from colossalai.tensor import ProcessGroup, ShardSpec
20 from colossalai.utils import get_current_device
21 from colossalai.utils.model.colo_init_context import ColoInitContext
22 
23 from .base import Strategy
24 from .ddp import DDPStrategy
25 
26 
27 class ColossalAIStrategy(DDPStrategy):
28     """
29         The strategy for training with ColossalAI.
30 
31     Args:
32         stage(int): The stage to use in ZeRO. Choose in (1, 2, 3)
33         seed(int): The seed for the random number generator.
34         shard_init(bool): Whether to shard the model parameters during initialization. Only for ZeRO-3.
35             This is not compativle with `from_pretrained()`. We temporarily disable this and will support it in the future.
36         placement_policy(str): The placement policy for gemini. Choose in ('cpu', 'cuda')
37                           If it is “cpu”, parameters, gradients and optimizer states will be offloaded to CPU,
38                           If it is “cuda”, they will not be offloaded, which means max CUDA memory will be used. It is the fastest.
39         pin_memory(bool): Whether to pin the memory for the data loader. Only for ZeRO-3.
40         force_outputs_fp32(bool): Whether to force the outputs to be fp32. Only for ZeRO-3.
41         search_range_mb(int): The search range in MB for the chunk size. Only for ZeRO-3.
42         hidden_dim(optional, int): The hidden dimension for the gemini. Only for ZeRO-3.
43         min_chunk_size_mb(float): The minimum chunk size in MB. Only for ZeRO-3.
44         gpu_margin_mem_ratio(float): The margin memory ratio for the GPU. Only for ZeRO-3.
45         reduce_bugket_size(int): The reduce bucket size in bytes. Only for ZeRO-1 and ZeRO-2.
46         overlap_communication(bool): Whether to overlap communication and computation. Only for ZeRO-1 and ZeRO-2.
47         initial_scale(float): The initial scale for the optimizer.
48         growth_factor(float): The growth factor for the optimizer.
49         backoff_factor(float): The backoff factor for the optimizer.
50         growth_interval(int): The growth interval for the optimizer.
51         hysteresis(int): The hysteresis for the optimizer.
52         min_scale(float): The minimum scale for the optimizer.
53         max_scale(float): The maximum scale for the optimizer.
54         max_norm(float): The maximum norm for the optimizer.
55         norm_type(float): The norm type for the optimizer.
56 
57     """
58 
59     def __init__(
60             self,
61             stage: int = 3,
62             seed: int = 42,
63             shard_init: bool = False,    # only for stage 3
64             placement_policy: str = 'cuda',
65             pin_memory: bool = True,    # only for stage 3
66             force_outputs_fp32: bool = False,    # only for stage 3
67             search_range_mb: int = 32,    # only for stage 3
68             hidden_dim: Optional[int] = None,    # only for stage 3
69             min_chunk_size_mb: float = 32,    # only for stage 3
70             gpu_margin_mem_ratio: float = 0.0,    # only for stage 3
71             reduce_bucket_size: int = 12 * 1024**2,    # only for stage 1&2
72             overlap_communication: bool = True,    # only for stage 1&2
73             initial_scale: float = 2**16,
74             growth_factor: float = 2,
75             backoff_factor: float = 0.5,
76             growth_interval: int = 1000,
77             hysteresis: int = 2,
78             min_scale: float = 1,
79             max_scale: float = 2**32,
80             max_norm: float = 0.0,
81             norm_type: float = 2.0) -> None:
82         super().__init__(seed)
83         assert placement_policy in ('cpu', 'cuda'), f'Unsupported placement policy "{placement_policy}"'
84         self.stage = stage
85         # TODO(ver217): support shard_init when using from_pretrained()
86         if shard_init:
87             warnings.warn(
88                 f'Shard init is not supported model.from_pretrained() yet. Please load weights after strategy.prepare()'
89             )
90         self.shard_init = shard_init
91         self.gemini_config = dict(device=get_current_device(),
92                                   placement_policy=placement_policy,
93                                   pin_memory=pin_memory,
94                                   force_outputs_fp32=force_outputs_fp32,
95                                   strict_ddp_mode=shard_init,
96                                   search_range_mb=search_range_mb,
97                                   hidden_dim=hidden_dim,
98                                   min_chunk_size_mb=min_chunk_size_mb)
99         if stage == 3:
100             self.zero_optim_config = dict(gpu_margin_mem_ratio=gpu_margin_mem_ratio)
101         else:
102             self.zero_optim_config = dict(reduce_bucket_size=reduce_bucket_size,
103                                           overlap_communication=overlap_communication,
104                                           cpu_offload=(placement_policy == 'cpu'))
105         self.optim_kwargs = dict(initial_scale=initial_scale,
106                                  growth_factor=growth_factor,
107                                  backoff_factor=backoff_factor,
108                                  growth_interval=growth_interval,
109                                  hysteresis=hysteresis,
110                                  min_scale=min_scale,
111                                  max_scale=max_scale,
112                                  max_norm=max_norm,
113                                  norm_type=norm_type)
114 
115     def setup_distributed(self) -> None:
116         colossalai.launch_from_torch({}, seed=self.seed)
117 
118     def model_init_context(self):
119         if self.stage == 3:
120             world_size = dist.get_world_size()
121             shard_pg = ProcessGroup(tp_degree=world_size) if self.shard_init else None
122             default_dist_spec = ShardSpec([-1], [world_size]) if self.shard_init else None
123             return ColoInitContext(device=get_current_device(),
124                                    dtype=torch.half,
125                                    default_pg=shard_pg,
126                                    default_dist_spec=default_dist_spec)
127         return super().model_init_context()
128 
129     def setup_model(self, model: nn.Module) -> nn.Module:
130         return zero_model_wrapper(model, zero_stage=self.stage, gemini_config=self.gemini_config)
131 
132     def setup_optimizer(self, optimizer: optim.Optimizer, model: nn.Module) -> optim.Optimizer:
133         assert isinstance(optimizer, (CPUAdam, HybridAdam)), f'Unsupported optimizer {type(optimizer)}'
134         return zero_optim_wrapper(model, optimizer, optim_config=self.zero_optim_config, **self.optim_kwargs)
135 
136     def backward(self, loss: torch.Tensor, model: nn.Module, optimizer: optim.Optimizer, **kwargs) -> None:
137         optimizer.backward(loss)
138 
139     def optimizer_step(self, optimizer: optim.Optimizer, **kwargs) -> None:
140         optimizer.step()
141 
142     @staticmethod
143     def _unwrap_actor(actor: Actor) -> nn.Module:
144         model: Union[nn.Module, ZeroDDP] = Strategy._unwrap_actor(actor)
145         if isinstance(model, ZeroDDP):
146             return model.module
147         return model
148 
149     def save_model(self, model: nn.Module, path: str, only_rank0: bool = False, tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None:
150         unwrapped_model = self._unwrap_model(model)
151         # TODO : better way to get torch model from gemini model
152         # to get torch model from gemini model
153         if isinstance(unwrapped_model, ZeroDDP):
154             state_dict = unwrapped_model.state_dict()
155             unwrapped_model = get_static_torch_model(unwrapped_model)
156             if only_rank0 and dist.get_rank() != 0:
157                 return
158             unwrapped_model.load_state_dict(state_dict)
159         # merge lora_weights into weights
160         for module in unwrapped_model.modules():
161             if isinstance(module, LoraLinear):
162                 module.merge_weights=True
163                 module.eval()
164         # get state_dict and save
165 
166         if not isinstance(self.model, PreTrainedModel):
167             state_dict = unwrapped_model.state_dict()
168             if only_rank0 and dist.get_rank() != 0:
169                 return
170             torch.save(state_dict, path)
171         else:
172             self.model.save_pretrained(path)
173             if tokenizer is not None:
174                 tokenizer.save_pretrained(path)
175 
176     def save_optimizer(self, optimizer: Optimizer, path: str, only_rank0: bool = False) -> None:
177         if only_rank0:
178             raise RuntimeError(
179                 f'Optimizer states are sharded when using ColossalAIStrategy. Only rank0 is not supported.')
180         torch.save(optimizer.state_dict(), path)
181
```
--- 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/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py b/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py
--- a/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py
+++ b/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py
@@ -27,6 +27,7 @@
 
     Args:
         stage(int): The stage to use in ZeRO. Choose in (1, 2, 3)
+        precision(str): The precision to use. Choose in ('fp32', 'fp16'). Stage 3 only supports fp16.
         seed(int): The seed for the random number generator.
         shard_init(bool): Whether to shard the model parameters during initialization. Only for ZeRO-3.
             This is not compativle with `from_pretrained()`. We temporarily disable this and will support it in the future.
@@ -56,6 +57,7 @@
     def __init__(
             self,
             stage: int = 3,
+            precision: str = 'fp16',
             seed: int = 42,
             shard_init: bool = False,    # only for stage 3
             placement_policy: str = 'cuda',
@@ -78,12 +80,17 @@
             norm_type: float = 2.0) -> None:
         super().__init__(seed)
         assert placement_policy in ('cpu', 'cuda'), f'Unsupported placement policy "{placement_policy}"'
+        assert precision in ('fp32', 'fp16'), f'Unsupported precision "{precision}"'
         self.stage = stage
         # TODO(ver217): support shard_init when using from_pretrained()
         if shard_init:
             warnings.warn(
                 f'Shard init is not supported model.from_pretrained() yet. Please load weights after strategy.prepare()'
             )
+        if stage == 3 and precision == 'fp32':
+            warnings.warn(f'Stage 3 only supports fp16. Precision is set to fp16.')
+            precision = 'fp16'
+        self.precision = precision
         self.shard_init = shard_init
         self.gemini_config = dict(device=get_current_device(),
                                   placement_policy=placement_policy,
@@ -124,7 +131,10 @@
         return super().model_init_context()
 
     def setup_model(self, model: nn.Module) -> nn.Module:
-        return zero_model_wrapper(model, zero_stage=self.stage, gemini_config=self.gemini_config)
+        model = zero_model_wrapper(model, zero_stage=self.stage, gemini_config=self.gemini_config)
+        if self.stage != 3 and self.precision == 'fp16':
+            model = model.half()
+        return model
 
     def setup_optimizer(self, optimizer: optim.Optimizer, model: nn.Module) -> optim.Optimizer:
         assert isinstance(optimizer, (CPUAdam, HybridAdam)), f'Unsupported optimizer {type(optimizer)}'
@@ -156,7 +166,7 @@
         # merge lora_weights into weights
         for module in unwrapped_model.modules():
             if isinstance(module, LoraLinear):
-                module.merge_weights=True
+                module.merge_weights = True
                 module.eval()
         # get state_dict and save
         state_dict = unwrapped_model.state_dict()

{"golden_diff": "diff --git a/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py b/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py\n--- a/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py\n+++ b/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py\n@@ -27,6 +27,7 @@\n \n     Args:\n         stage(int): The stage to use in ZeRO. Choose in (1, 2, 3)\n+        precision(str): The precision to use. Choose in ('fp32', 'fp16'). Stage 3 only supports fp16.\n         seed(int): The seed for the random number generator.\n         shard_init(bool): Whether to shard the model parameters during initialization. Only for ZeRO-3.\n             This is not compativle with `from_pretrained()`. We temporarily disable this and will support it in the future.\n@@ -56,6 +57,7 @@\n     def __init__(\n             self,\n             stage: int = 3,\n+            precision: str = 'fp16',\n             seed: int = 42,\n             shard_init: bool = False,    # only for stage 3\n             placement_policy: str = 'cuda',\n@@ -78,12 +80,17 @@\n             norm_type: float = 2.0) -> None:\n         super().__init__(seed)\n         assert placement_policy in ('cpu', 'cuda'), f'Unsupported placement policy \"{placement_policy}\"'\n+        assert precision in ('fp32', 'fp16'), f'Unsupported precision \"{precision}\"'\n         self.stage = stage\n         # TODO(ver217): support shard_init when using from_pretrained()\n         if shard_init:\n             warnings.warn(\n                 f'Shard init is not supported model.from_pretrained() yet. Please load weights after strategy.prepare()'\n             )\n+        if stage == 3 and precision == 'fp32':\n+            warnings.warn(f'Stage 3 only supports fp16. Precision is set to fp16.')\n+            precision = 'fp16'\n+        self.precision = precision\n         self.shard_init = shard_init\n         self.gemini_config = dict(device=get_current_device(),\n                                   placement_policy=placement_policy,\n@@ -124,7 +131,10 @@\n         return super().model_init_context()\n \n     def setup_model(self, model: nn.Module) -> nn.Module:\n-        return zero_model_wrapper(model, zero_stage=self.stage, gemini_config=self.gemini_config)\n+        model = zero_model_wrapper(model, zero_stage=self.stage, gemini_config=self.gemini_config)\n+        if self.stage != 3 and self.precision == 'fp16':\n+            model = model.half()\n+        return model\n \n     def setup_optimizer(self, optimizer: optim.Optimizer, model: nn.Module) -> optim.Optimizer:\n         assert isinstance(optimizer, (CPUAdam, HybridAdam)), f'Unsupported optimizer {type(optimizer)}'\n@@ -156,7 +166,7 @@\n         # merge lora_weights into weights\n         for module in unwrapped_model.modules():\n             if isinstance(module, LoraLinear):\n-                module.merge_weights=True\n+                module.merge_weights = True\n                 module.eval()\n         # get state_dict and save\n         state_dict = unwrapped_model.state_dict()\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import warnings\nfrom typing import Optional, Union\n\nimport torch\nimport torch.distributed as dist\nimport torch.nn as nn\nimport torch.optim as optim\nfrom chatgpt.models.base import Actor\nfrom chatgpt.models.lora import LoraLinear\nfrom torch.optim import Optimizer\n\nfrom transformers.modeling_utils import PreTrainedModel\nfrom transformers.tokenization_utils_base import PreTrainedTokenizerBase\n\nimport colossalai\nfrom colossalai.nn.optimizer import CPUAdam, HybridAdam\nfrom colossalai.nn.parallel import ZeroDDP, zero_model_wrapper, zero_optim_wrapper\nfrom colossalai.nn.parallel.utils import get_static_torch_model\nfrom colossalai.tensor import ProcessGroup, ShardSpec\nfrom colossalai.utils import get_current_device\nfrom colossalai.utils.model.colo_init_context import ColoInitContext\n\nfrom .base import Strategy\nfrom .ddp import DDPStrategy\n\n\nclass ColossalAIStrategy(DDPStrategy):\n    \"\"\"\n        The strategy for training with ColossalAI.\n\n    Args:\n        stage(int): The stage to use in ZeRO. Choose in (1, 2, 3)\n        seed(int): The seed for the random number generator.\n        shard_init(bool): Whether to shard the model parameters during initialization. Only for ZeRO-3.\n            This is not compativle with `from_pretrained()`. We temporarily disable this and will support it in the future.\n        placement_policy(str): The placement policy for gemini. Choose in ('cpu', 'cuda')\n                          If it is \u201ccpu\u201d, parameters, gradients and optimizer states will be offloaded to CPU,\n                          If it is \u201ccuda\u201d, they will not be offloaded, which means max CUDA memory will be used. It is the fastest.\n        pin_memory(bool): Whether to pin the memory for the data loader. Only for ZeRO-3.\n        force_outputs_fp32(bool): Whether to force the outputs to be fp32. Only for ZeRO-3.\n        search_range_mb(int): The search range in MB for the chunk size. Only for ZeRO-3.\n        hidden_dim(optional, int): The hidden dimension for the gemini. Only for ZeRO-3.\n        min_chunk_size_mb(float): The minimum chunk size in MB. Only for ZeRO-3.\n        gpu_margin_mem_ratio(float): The margin memory ratio for the GPU. Only for ZeRO-3.\n        reduce_bugket_size(int): The reduce bucket size in bytes. Only for ZeRO-1 and ZeRO-2.\n        overlap_communication(bool): Whether to overlap communication and computation. Only for ZeRO-1 and ZeRO-2.\n        initial_scale(float): The initial scale for the optimizer.\n        growth_factor(float): The growth factor for the optimizer.\n        backoff_factor(float): The backoff factor for the optimizer.\n        growth_interval(int): The growth interval for the optimizer.\n        hysteresis(int): The hysteresis for the optimizer.\n        min_scale(float): The minimum scale for the optimizer.\n        max_scale(float): The maximum scale for the optimizer.\n        max_norm(float): The maximum norm for the optimizer.\n        norm_type(float): The norm type for the optimizer.\n\n    \"\"\"\n\n    def __init__(\n            self,\n            stage: int = 3,\n            seed: int = 42,\n            shard_init: bool = False,    # only for stage 3\n            placement_policy: str = 'cuda',\n            pin_memory: bool = True,    # only for stage 3\n            force_outputs_fp32: bool = False,    # only for stage 3\n            search_range_mb: int = 32,    # only for stage 3\n            hidden_dim: Optional[int] = None,    # only for stage 3\n            min_chunk_size_mb: float = 32,    # only for stage 3\n            gpu_margin_mem_ratio: float = 0.0,    # only for stage 3\n            reduce_bucket_size: int = 12 * 1024**2,    # only for stage 1&2\n            overlap_communication: bool = True,    # only for stage 1&2\n            initial_scale: float = 2**16,\n            growth_factor: float = 2,\n            backoff_factor: float = 0.5,\n            growth_interval: int = 1000,\n            hysteresis: int = 2,\n            min_scale: float = 1,\n            max_scale: float = 2**32,\n            max_norm: float = 0.0,\n            norm_type: float = 2.0) -> None:\n        super().__init__(seed)\n        assert placement_policy in ('cpu', 'cuda'), f'Unsupported placement policy \"{placement_policy}\"'\n        self.stage = stage\n        # TODO(ver217): support shard_init when using from_pretrained()\n        if shard_init:\n            warnings.warn(\n                f'Shard init is not supported model.from_pretrained() yet. Please load weights after strategy.prepare()'\n            )\n        self.shard_init = shard_init\n        self.gemini_config = dict(device=get_current_device(),\n                                  placement_policy=placement_policy,\n                                  pin_memory=pin_memory,\n                                  force_outputs_fp32=force_outputs_fp32,\n                                  strict_ddp_mode=shard_init,\n                                  search_range_mb=search_range_mb,\n                                  hidden_dim=hidden_dim,\n                                  min_chunk_size_mb=min_chunk_size_mb)\n        if stage == 3:\n            self.zero_optim_config = dict(gpu_margin_mem_ratio=gpu_margin_mem_ratio)\n        else:\n            self.zero_optim_config = dict(reduce_bucket_size=reduce_bucket_size,\n                                          overlap_communication=overlap_communication,\n                                          cpu_offload=(placement_policy == 'cpu'))\n        self.optim_kwargs = dict(initial_scale=initial_scale,\n                                 growth_factor=growth_factor,\n                                 backoff_factor=backoff_factor,\n                                 growth_interval=growth_interval,\n                                 hysteresis=hysteresis,\n                                 min_scale=min_scale,\n                                 max_scale=max_scale,\n                                 max_norm=max_norm,\n                                 norm_type=norm_type)\n\n    def setup_distributed(self) -> None:\n        colossalai.launch_from_torch({}, seed=self.seed)\n\n    def model_init_context(self):\n        if self.stage == 3:\n            world_size = dist.get_world_size()\n            shard_pg = ProcessGroup(tp_degree=world_size) if self.shard_init else None\n            default_dist_spec = ShardSpec([-1], [world_size]) if self.shard_init else None\n            return ColoInitContext(device=get_current_device(),\n                                   dtype=torch.half,\n                                   default_pg=shard_pg,\n                                   default_dist_spec=default_dist_spec)\n        return super().model_init_context()\n\n    def setup_model(self, model: nn.Module) -> nn.Module:\n        return zero_model_wrapper(model, zero_stage=self.stage, gemini_config=self.gemini_config)\n\n    def setup_optimizer(self, optimizer: optim.Optimizer, model: nn.Module) -> optim.Optimizer:\n        assert isinstance(optimizer, (CPUAdam, HybridAdam)), f'Unsupported optimizer {type(optimizer)}'\n        return zero_optim_wrapper(model, optimizer, optim_config=self.zero_optim_config, **self.optim_kwargs)\n\n    def backward(self, loss: torch.Tensor, model: nn.Module, optimizer: optim.Optimizer, **kwargs) -> None:\n        optimizer.backward(loss)\n\n    def optimizer_step(self, optimizer: optim.Optimizer, **kwargs) -> None:\n        optimizer.step()\n\n    @staticmethod\n    def _unwrap_actor(actor: Actor) -> nn.Module:\n        model: Union[nn.Module, ZeroDDP] = Strategy._unwrap_actor(actor)\n        if isinstance(model, ZeroDDP):\n            return model.module\n        return model\n\n    def save_model(self, model: nn.Module, path: str, only_rank0: bool = False, tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None:\n        unwrapped_model = self._unwrap_model(model)\n        # TODO : better way to get torch model from gemini model\n        # to get torch model from gemini model\n        if isinstance(unwrapped_model, ZeroDDP):\n            state_dict = unwrapped_model.state_dict()\n            unwrapped_model = get_static_torch_model(unwrapped_model)\n            if only_rank0 and dist.get_rank() != 0:\n                return\n            unwrapped_model.load_state_dict(state_dict)\n        # merge lora_weights into weights\n        for module in unwrapped_model.modules():\n            if isinstance(module, LoraLinear):\n                module.merge_weights=True\n                module.eval()\n        # get state_dict and save\n\n        if not isinstance(self.model, PreTrainedModel):\n            state_dict = unwrapped_model.state_dict()\n            if only_rank0 and dist.get_rank() != 0:\n                return\n            torch.save(state_dict, path)\n        else:\n            self.model.save_pretrained(path)\n            if tokenizer is not None:\n                tokenizer.save_pretrained(path)\n\n    def save_optimizer(self, optimizer: Optimizer, path: str, only_rank0: bool = False) -> None:\n        if only_rank0:\n            raise RuntimeError(\n                f'Optimizer states are sharded when using ColossalAIStrategy. Only rank0 is not supported.')\n        torch.save(optimizer.state_dict(), path)\n", "path": "applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py"}], "after_files": [{"content": "import warnings\nfrom typing import Optional, Union\n\nimport torch\nimport torch.distributed as dist\nimport torch.nn as nn\nimport torch.optim as optim\nfrom chatgpt.models.base import Actor\nfrom chatgpt.models.lora import LoraLinear\nfrom torch.optim import Optimizer\n\nimport colossalai\nfrom colossalai.nn.optimizer import CPUAdam, HybridAdam\nfrom colossalai.nn.parallel import ZeroDDP, zero_model_wrapper, zero_optim_wrapper\nfrom colossalai.nn.parallel.utils import get_static_torch_model\nfrom colossalai.tensor import ProcessGroup, ShardSpec\nfrom colossalai.utils import get_current_device\nfrom colossalai.utils.model.colo_init_context import ColoInitContext\n\nfrom .base import Strategy\nfrom .ddp import DDPStrategy\n\n\nclass ColossalAIStrategy(DDPStrategy):\n    \"\"\"\n        The strategy for training with ColossalAI.\n\n    Args:\n        stage(int): The stage to use in ZeRO. Choose in (1, 2, 3)\n        precision(str): The precision to use. Choose in ('fp32', 'fp16'). Stage 3 only supports fp16.\n        seed(int): The seed for the random number generator.\n        shard_init(bool): Whether to shard the model parameters during initialization. Only for ZeRO-3.\n            This is not compativle with `from_pretrained()`. We temporarily disable this and will support it in the future.\n        placement_policy(str): The placement policy for gemini. Choose in ('cpu', 'cuda')\n                          If it is \u201ccpu\u201d, parameters, gradients and optimizer states will be offloaded to CPU,\n                          If it is \u201ccuda\u201d, they will not be offloaded, which means max CUDA memory will be used. It is the fastest.\n        pin_memory(bool): Whether to pin the memory for the data loader. Only for ZeRO-3.\n        force_outputs_fp32(bool): Whether to force the outputs to be fp32. Only for ZeRO-3.\n        search_range_mb(int): The search range in MB for the chunk size. Only for ZeRO-3.\n        hidden_dim(optional, int): The hidden dimension for the gemini. Only for ZeRO-3.\n        min_chunk_size_mb(float): The minimum chunk size in MB. Only for ZeRO-3.\n        gpu_margin_mem_ratio(float): The margin memory ratio for the GPU. Only for ZeRO-3.\n        reduce_bugket_size(int): The reduce bucket size in bytes. Only for ZeRO-1 and ZeRO-2.\n        overlap_communication(bool): Whether to overlap communication and computation. Only for ZeRO-1 and ZeRO-2.\n        initial_scale(float): The initial scale for the optimizer.\n        growth_factor(float): The growth factor for the optimizer.\n        backoff_factor(float): The backoff factor for the optimizer.\n        growth_interval(int): The growth interval for the optimizer.\n        hysteresis(int): The hysteresis for the optimizer.\n        min_scale(float): The minimum scale for the optimizer.\n        max_scale(float): The maximum scale for the optimizer.\n        max_norm(float): The maximum norm for the optimizer.\n        norm_type(float): The norm type for the optimizer.\n\n    \"\"\"\n\n    def __init__(\n            self,\n            stage: int = 3,\n            precision: str = 'fp16',\n            seed: int = 42,\n            shard_init: bool = False,    # only for stage 3\n            placement_policy: str = 'cuda',\n            pin_memory: bool = True,    # only for stage 3\n            force_outputs_fp32: bool = False,    # only for stage 3\n            search_range_mb: int = 32,    # only for stage 3\n            hidden_dim: Optional[int] = None,    # only for stage 3\n            min_chunk_size_mb: float = 32,    # only for stage 3\n            gpu_margin_mem_ratio: float = 0.0,    # only for stage 3\n            reduce_bucket_size: int = 12 * 1024**2,    # only for stage 1&2\n            overlap_communication: bool = True,    # only for stage 1&2\n            initial_scale: float = 2**16,\n            growth_factor: float = 2,\n            backoff_factor: float = 0.5,\n            growth_interval: int = 1000,\n            hysteresis: int = 2,\n            min_scale: float = 1,\n            max_scale: float = 2**32,\n            max_norm: float = 0.0,\n            norm_type: float = 2.0) -> None:\n        super().__init__(seed)\n        assert placement_policy in ('cpu', 'cuda'), f'Unsupported placement policy \"{placement_policy}\"'\n        assert precision in ('fp32', 'fp16'), f'Unsupported precision \"{precision}\"'\n        self.stage = stage\n        # TODO(ver217): support shard_init when using from_pretrained()\n        if shard_init:\n            warnings.warn(\n                f'Shard init is not supported model.from_pretrained() yet. Please load weights after strategy.prepare()'\n            )\n        if stage == 3 and precision == 'fp32':\n            warnings.warn(f'Stage 3 only supports fp16. Precision is set to fp16.')\n            precision = 'fp16'\n        self.precision = precision\n        self.shard_init = shard_init\n        self.gemini_config = dict(device=get_current_device(),\n                                  placement_policy=placement_policy,\n                                  pin_memory=pin_memory,\n                                  force_outputs_fp32=force_outputs_fp32,\n                                  strict_ddp_mode=shard_init,\n                                  search_range_mb=search_range_mb,\n                                  hidden_dim=hidden_dim,\n                                  min_chunk_size_mb=min_chunk_size_mb)\n        if stage == 3:\n            self.zero_optim_config = dict(gpu_margin_mem_ratio=gpu_margin_mem_ratio)\n        else:\n            self.zero_optim_config = dict(reduce_bucket_size=reduce_bucket_size,\n                                          overlap_communication=overlap_communication,\n                                          cpu_offload=(placement_policy == 'cpu'))\n        self.optim_kwargs = dict(initial_scale=initial_scale,\n                                 growth_factor=growth_factor,\n                                 backoff_factor=backoff_factor,\n                                 growth_interval=growth_interval,\n                                 hysteresis=hysteresis,\n                                 min_scale=min_scale,\n                                 max_scale=max_scale,\n                                 max_norm=max_norm,\n                                 norm_type=norm_type)\n\n    def setup_distributed(self) -> None:\n        colossalai.launch_from_torch({}, seed=self.seed)\n\n    def model_init_context(self):\n        if self.stage == 3:\n            world_size = dist.get_world_size()\n            shard_pg = ProcessGroup(tp_degree=world_size) if self.shard_init else None\n            default_dist_spec = ShardSpec([-1], [world_size]) if self.shard_init else None\n            return ColoInitContext(device=get_current_device(),\n                                   dtype=torch.half,\n                                   default_pg=shard_pg,\n                                   default_dist_spec=default_dist_spec)\n        return super().model_init_context()\n\n    def setup_model(self, model: nn.Module) -> nn.Module:\n        model = zero_model_wrapper(model, zero_stage=self.stage, gemini_config=self.gemini_config)\n        if self.stage != 3 and self.precision == 'fp16':\n            model = model.half()\n        return model\n\n    def setup_optimizer(self, optimizer: optim.Optimizer, model: nn.Module) -> optim.Optimizer:\n        assert isinstance(optimizer, (CPUAdam, HybridAdam)), f'Unsupported optimizer {type(optimizer)}'\n        return zero_optim_wrapper(model, optimizer, optim_config=self.zero_optim_config, **self.optim_kwargs)\n\n    def backward(self, loss: torch.Tensor, model: nn.Module, optimizer: optim.Optimizer, **kwargs) -> None:\n        optimizer.backward(loss)\n\n    def optimizer_step(self, optimizer: optim.Optimizer, **kwargs) -> None:\n        optimizer.step()\n\n    @staticmethod\n    def _unwrap_actor(actor: Actor) -> nn.Module:\n        model: Union[nn.Module, ZeroDDP] = Strategy._unwrap_actor(actor)\n        if isinstance(model, ZeroDDP):\n            return model.module\n        return model\n\n    def save_model(self, model: nn.Module, path: str, only_rank0: bool = False) -> None:\n        unwrapped_model = self._unwrap_model(model)\n        # TODO : better way to get torch model from gemini model\n        # to get torch model from gemini model\n        if isinstance(unwrapped_model, ZeroDDP):\n            state_dict = unwrapped_model.state_dict()\n            unwrapped_model = get_static_torch_model(unwrapped_model)\n            if only_rank0 and dist.get_rank() != 0:\n                return\n            unwrapped_model.load_state_dict(state_dict)\n        # merge lora_weights into weights\n        for module in unwrapped_model.modules():\n            if isinstance(module, LoraLinear):\n                module.merge_weights = True\n                module.eval()\n        # get state_dict and save\n        state_dict = unwrapped_model.state_dict()\n        if only_rank0 and dist.get_rank() != 0:\n            return\n        torch.save(state_dict, path)\n\n    def save_optimizer(self, optimizer: Optimizer, path: str, only_rank0: bool = False) -> None:\n        if only_rank0:\n            raise RuntimeError(\n                f'Optimizer states are sharded when using ColossalAIStrategy. Only rank0 is not supported.')\n        torch.save(optimizer.state_dict(), path)\n", "path": "applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py"}]}