We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[DOC]: A bug about usage of booster api in the document
### 📚 The doc issue

[zero_with_chunk](https://colossalai.org/docs/features/zero_with_chunk) and [booster_api](https://colossalai.org/docs/basics/booster_api) both used torch.optim.Adam as the type of optimizer in the code examples implemented with booster api. However, this will trigger an AssertionError when passing this kind of optimizer to booster.boost:

![20230605-161741](https://github.com/hpcaitech/ColossalAI/assets/56809903/789fda8d-40eb-489c-9026-56d3d872e1c1)

[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: `colossalai/booster/booster.py`
Content:
```
1 import warnings
2 from contextlib import contextmanager
3 from typing import Callable, Iterator, List, Optional, Tuple, Union
4 
5 import torch
6 import torch.nn as nn
7 from torch.optim import Optimizer
8 from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
9 from torch.utils.data import DataLoader
10 
11 from colossalai.checkpoint_io import GeneralCheckpointIO
12 
13 from .accelerator import Accelerator
14 from .mixed_precision import MixedPrecision, mixed_precision_factory
15 from .plugin import Plugin
16 
17 __all__ = ['Booster']
18 
19 
20 class Booster:
21     """
22     Booster is a high-level API for training neural networks. It provides a unified interface for
23     training with different precision, accelerator, and plugin.
24 
25     Examples:
26         ```python
27         colossalai.launch(...)
28         plugin = GeminiPlugin(stage=3, ...)
29         booster = Booster(precision='fp16', plugin=plugin)
30 
31         model = GPT2()
32         optimizer = Adam(model.parameters())
33         dataloader = Dataloader(Dataset)
34         lr_scheduler = LinearWarmupScheduler()
35         criterion = GPTLMLoss()
36 
37         model, optimizer, lr_scheduler, dataloader = booster.boost(model, optimizer, lr_scheduler, dataloader)
38 
39         for epoch in range(max_epochs):
40             for input_ids, attention_mask in dataloader:
41                 outputs = model(input_ids, attention_mask)
42                 loss = criterion(outputs.logits, input_ids)
43                 booster.backward(loss, optimizer)
44                 optimizer.step()
45                 lr_scheduler.step()
46                 optimizer.zero_grad()
47         ```
48 
49     Args:
50         device (str or torch.device): The device to run the training. Default: 'cuda'.
51         mixed_precision (str or MixedPrecision): The mixed precision to run the training. Default: None.
52                                 If the argument is a string, it can be 'fp16', 'fp16_apex', 'bf16', or 'fp8'.
53                                 'fp16' would use PyTorch AMP while `fp16_apex` would use Nvidia Apex.
54         plugin (Plugin): The plugin to run the training. Default: None.
55     """
56 
57     def __init__(self,
58                  device: str = 'cuda',
59                  mixed_precision: Union[MixedPrecision, str] = None,
60                  plugin: Optional[Plugin] = None) -> None:
61         if plugin is not None:
62             assert isinstance(
63                 plugin, Plugin), f'Expected the argument plugin to be an instance of Plugin, but got {type(plugin)}.'
64         self.plugin = plugin
65 
66         # set accelerator
67         if self.plugin and self.plugin.control_device():
68             self.accelerator = None
69             warnings.warn('The plugin will control the accelerator, so the device argument will be ignored.')
70         else:
71             self.accelerator = Accelerator(device)
72 
73         # set precision
74         if self.plugin and self.plugin.control_precision():
75             warnings.warn('The plugin will control the precision, so the mixed_precision argument will be ignored.')
76             self.mixed_precision = None
77         elif mixed_precision is None:
78             self.mixed_precision = None
79         else:
80             # validate and set precision
81             if isinstance(mixed_precision, str):
82                 # the user will take the default arguments for amp training
83                 self.mixed_precision = mixed_precision_factory(mixed_precision)
84             elif isinstance(mixed_precision, MixedPrecision):
85                 # the user can customize the arguments by passing the precision object
86                 self.mixed_precision = mixed_precision
87             else:
88                 raise ValueError(
89                     f'Expected the argument mixed_precision to be a string or an instance of Precision, but got {type(mixed_precision)}.'
90                 )
91 
92         if self.plugin is not None and self.plugin.control_checkpoint_io():
93             self.checkpoint_io = self.plugin.get_checkpoint_io()
94         else:
95             self.checkpoint_io = GeneralCheckpointIO()
96 
97     def boost(
98         self,
99         model: nn.Module,
100         optimizer: Optimizer,
101         criterion: Callable = None,
102         dataloader: DataLoader = None,
103         lr_scheduler: LRScheduler = None,
104     ) -> List[Union[nn.Module, Optimizer, LRScheduler, DataLoader]]:
105         """
106         Boost the model, optimizer, criterion, lr_scheduler, and dataloader.
107 
108         Args:
109             model (nn.Module): The model to be boosted.
110             optimizer (Optimizer): The optimizer to be boosted.
111             criterion (Callable): The criterion to be boosted.
112             dataloader (DataLoader): The dataloader to be boosted.
113             lr_scheduler (LRScheduler): The lr_scheduler to be boosted.
114         """
115         # TODO(FrankLeeeee): consider multi-model and multi-optimizer case
116         # TODO(FrankLeeeee): consider multi-dataloader case
117         # transform model for mixed precision
118         if self.plugin:
119             model, optimizer, criterion, dataloader, lr_scheduler = self.plugin.configure(
120                 model, optimizer, criterion, dataloader, lr_scheduler)
121 
122         if self.plugin and not self.plugin.control_device():
123             # transform model for accelerator
124             model = self.accelerator.configure(model)
125 
126         if self.mixed_precision and (self.plugin is None or self.plugin and not self.plugin.control_precision()):
127             # transform model for mixed precision
128             # when mixed_precision is specified and the plugin is not given or does not control the precision
129             model, optimizer, criterion = self.mixed_precision.configure(model, optimizer, criterion)
130 
131         return model, optimizer, criterion, dataloader, lr_scheduler
132 
133     def backward(self, loss: torch.Tensor, optimizer: Optimizer) -> None:
134         """Backward pass.
135 
136         Args:
137             loss (torch.Tensor): The loss to be backpropagated.
138             optimizer (Optimizer): The optimizer to be updated.
139         """
140         # TODO: implement this method with plugin
141         optimizer.backward(loss)
142 
143     def execute_pipeline(self,
144                          data_iter: Iterator,
145                          model: nn.Module,
146                          criterion: Callable[[torch.Tensor], torch.Tensor],
147                          optimizer: Optimizer,
148                          return_loss: bool = True,
149                          return_outputs: bool = False) -> Tuple[Optional[torch.Tensor], ...]:
150         # TODO: implement this method
151         # run pipeline forward backward pass
152         # return loss or outputs if needed
153         pass
154 
155     def no_sync(self, model: nn.Module) -> contextmanager:
156         """Context manager to disable gradient synchronization across DP process groups.
157 
158         Args:
159             model (nn.Module): The model to be disabled gradient synchronization.
160 
161         Returns:
162             contextmanager: Context to disable gradient synchronization.
163         """
164         assert self.plugin is not None, f'no_sync is only enabled when a plugin is provided and the plugin supports no_sync.'
165         assert self.plugin.support_no_sync, f'The plugin {self.plugin.__class__.__name__} does not support no_sync.'
166         return self.plugin.no_sync(model)
167 
168     def load_model(self, model: nn.Module, checkpoint: str, strict: bool = True):
169         """Load model from checkpoint.
170 
171         Args:
172             model (nn.Module): A model boosted by Booster.
173             checkpoint (str): Path to the checkpoint. It must be a local path.
174                 It should be a directory path if the checkpoint is sharded. Otherwise, it should be a file path.
175             strict (bool, optional): whether to strictly enforce that the keys
176                 in :attr:`state_dict` match the keys returned by this module's
177                 :meth:`~torch.nn.Module.state_dict` function. Defaults to True.
178         """
179         self.checkpoint_io.load_model(model, checkpoint, strict)
180 
181     def save_model(self,
182                    model: nn.Module,
183                    checkpoint: str,
184                    prefix: str = None,
185                    shard: bool = False,
186                    size_per_shard: int = 1024):
187         """Save model to checkpoint.
188 
189         Args:
190             model (nn.Module): A model boosted by Booster.
191             checkpoint (str): Path to the checkpoint. It must be a local path.
192                 It is a file path if ``shard=False``. Otherwise, it is a directory path.
193             prefix (str, optional): A prefix added to parameter and buffer
194                 names to compose the keys in state_dict. Defaults to None.
195             shard (bool, optional): Whether to save checkpoint a sharded way.
196                 If true, the checkpoint will be a folder. Otherwise, it will be a single file. Defaults to False.
197             size_per_shard (int, optional): Maximum size of checkpoint shard file in MB. This is useful only when ``shard=True``. Defaults to 1024.
198         """
199         self.checkpoint_io.save_model(model, checkpoint=checkpoint, shard=shard, size_per_shard=size_per_shard)
200 
201     def load_optimizer(self, optimizer: Optimizer, checkpoint: str):
202         """Load optimizer from checkpoint.
203 
204         Args:
205             optimizer (Optimizer): An optimizer boosted by Booster.
206             checkpoint (str): Path to the checkpoint. It must be a local path.
207                 It should be a directory path if the checkpoint is sharded. Otherwise, it should be a file path.
208         """
209         self.checkpoint_io.load_optimizer(optimizer, checkpoint)
210 
211     def save_optimizer(self, optimizer: Optimizer, checkpoint: str, shard: bool = False, size_per_shard: int = 1024):
212         """Save optimizer to checkpoint.
213         Warning: Saving sharded optimizer checkpoint is not supported yet.
214 
215         Args:
216             optimizer (Optimizer): An optimizer boosted by Booster.
217             checkpoint (str): Path to the checkpoint. It must be a local path.
218                 It is a file path if ``shard=False``. Otherwise, it is a directory path.
219             shard (bool, optional): Whether to save checkpoint a sharded way.
220                 If true, the checkpoint will be a folder. Otherwise, it will be a single file. Defaults to False.
221             size_per_shard (int, optional): Maximum size of checkpoint shard file in MB. This is useful only when ``shard=True``. Defaults to 1024.
222         """
223         self.checkpoint_io.save_optimizer(optimizer, checkpoint, shard, size_per_shard)
224 
225     def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):
226         """Save lr scheduler to checkpoint.
227 
228         Args:
229             lr_scheduler (LRScheduler): A lr scheduler boosted by Booster.
230             checkpoint (str): Path to the checkpoint. It must be a local file path.
231         """
232         self.checkpoint_io.save_lr_scheduler(lr_scheduler, checkpoint)
233 
234     def load_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):
235         """Load lr scheduler from checkpoint.
236 
237         Args:
238             lr_scheduler (LRScheduler): A lr scheduler boosted by Booster.
239             checkpoint (str): Path to the checkpoint. It must be a local file path.
240         """
241         self.checkpoint_io.load_lr_scheduler(lr_scheduler, checkpoint)
242
```
--- 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/colossalai/booster/booster.py b/colossalai/booster/booster.py
--- a/colossalai/booster/booster.py
+++ b/colossalai/booster/booster.py
@@ -25,11 +25,11 @@
     Examples:
         ```python
         colossalai.launch(...)
-        plugin = GeminiPlugin(stage=3, ...)
+        plugin = GeminiPlugin(...)
         booster = Booster(precision='fp16', plugin=plugin)
 
         model = GPT2()
-        optimizer = Adam(model.parameters())
+        optimizer = HybridAdam(model.parameters())
         dataloader = Dataloader(Dataset)
         lr_scheduler = LinearWarmupScheduler()
         criterion = GPTLMLoss()

{"golden_diff": "diff --git a/colossalai/booster/booster.py b/colossalai/booster/booster.py\n--- a/colossalai/booster/booster.py\n+++ b/colossalai/booster/booster.py\n@@ -25,11 +25,11 @@\n     Examples:\n         ```python\n         colossalai.launch(...)\n-        plugin = GeminiPlugin(stage=3, ...)\n+        plugin = GeminiPlugin(...)\n         booster = Booster(precision='fp16', plugin=plugin)\n \n         model = GPT2()\n-        optimizer = Adam(model.parameters())\n+        optimizer = HybridAdam(model.parameters())\n         dataloader = Dataloader(Dataset)\n         lr_scheduler = LinearWarmupScheduler()\n         criterion = GPTLMLoss()\n", "issue": "[DOC]: A bug about usage of booster api in the document\n### \ud83d\udcda The doc issue\n\n[zero_with_chunk](https://colossalai.org/docs/features/zero_with_chunk) and [booster_api](https://colossalai.org/docs/basics/booster_api) both used torch.optim.Adam as the type of optimizer in the code examples implemented with booster api. However, this will trigger an AssertionError when passing this kind of optimizer to booster.boost:\r\n\r\n![20230605-161741](https://github.com/hpcaitech/ColossalAI/assets/56809903/789fda8d-40eb-489c-9026-56d3d872e1c1)\r\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import warnings\nfrom contextlib import contextmanager\nfrom typing import Callable, Iterator, List, Optional, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import _LRScheduler as LRScheduler\nfrom torch.utils.data import DataLoader\n\nfrom colossalai.checkpoint_io import GeneralCheckpointIO\n\nfrom .accelerator import Accelerator\nfrom .mixed_precision import MixedPrecision, mixed_precision_factory\nfrom .plugin import Plugin\n\n__all__ = ['Booster']\n\n\nclass Booster:\n    \"\"\"\n    Booster is a high-level API for training neural networks. It provides a unified interface for\n    training with different precision, accelerator, and plugin.\n\n    Examples:\n        ```python\n        colossalai.launch(...)\n        plugin = GeminiPlugin(stage=3, ...)\n        booster = Booster(precision='fp16', plugin=plugin)\n\n        model = GPT2()\n        optimizer = Adam(model.parameters())\n        dataloader = Dataloader(Dataset)\n        lr_scheduler = LinearWarmupScheduler()\n        criterion = GPTLMLoss()\n\n        model, optimizer, lr_scheduler, dataloader = booster.boost(model, optimizer, lr_scheduler, dataloader)\n\n        for epoch in range(max_epochs):\n            for input_ids, attention_mask in dataloader:\n                outputs = model(input_ids, attention_mask)\n                loss = criterion(outputs.logits, input_ids)\n                booster.backward(loss, optimizer)\n                optimizer.step()\n                lr_scheduler.step()\n                optimizer.zero_grad()\n        ```\n\n    Args:\n        device (str or torch.device): The device to run the training. Default: 'cuda'.\n        mixed_precision (str or MixedPrecision): The mixed precision to run the training. Default: None.\n                                If the argument is a string, it can be 'fp16', 'fp16_apex', 'bf16', or 'fp8'.\n                                'fp16' would use PyTorch AMP while `fp16_apex` would use Nvidia Apex.\n        plugin (Plugin): The plugin to run the training. Default: None.\n    \"\"\"\n\n    def __init__(self,\n                 device: str = 'cuda',\n                 mixed_precision: Union[MixedPrecision, str] = None,\n                 plugin: Optional[Plugin] = None) -> None:\n        if plugin is not None:\n            assert isinstance(\n                plugin, Plugin), f'Expected the argument plugin to be an instance of Plugin, but got {type(plugin)}.'\n        self.plugin = plugin\n\n        # set accelerator\n        if self.plugin and self.plugin.control_device():\n            self.accelerator = None\n            warnings.warn('The plugin will control the accelerator, so the device argument will be ignored.')\n        else:\n            self.accelerator = Accelerator(device)\n\n        # set precision\n        if self.plugin and self.plugin.control_precision():\n            warnings.warn('The plugin will control the precision, so the mixed_precision argument will be ignored.')\n            self.mixed_precision = None\n        elif mixed_precision is None:\n            self.mixed_precision = None\n        else:\n            # validate and set precision\n            if isinstance(mixed_precision, str):\n                # the user will take the default arguments for amp training\n                self.mixed_precision = mixed_precision_factory(mixed_precision)\n            elif isinstance(mixed_precision, MixedPrecision):\n                # the user can customize the arguments by passing the precision object\n                self.mixed_precision = mixed_precision\n            else:\n                raise ValueError(\n                    f'Expected the argument mixed_precision to be a string or an instance of Precision, but got {type(mixed_precision)}.'\n                )\n\n        if self.plugin is not None and self.plugin.control_checkpoint_io():\n            self.checkpoint_io = self.plugin.get_checkpoint_io()\n        else:\n            self.checkpoint_io = GeneralCheckpointIO()\n\n    def boost(\n        self,\n        model: nn.Module,\n        optimizer: Optimizer,\n        criterion: Callable = None,\n        dataloader: DataLoader = None,\n        lr_scheduler: LRScheduler = None,\n    ) -> List[Union[nn.Module, Optimizer, LRScheduler, DataLoader]]:\n        \"\"\"\n        Boost the model, optimizer, criterion, lr_scheduler, and dataloader.\n\n        Args:\n            model (nn.Module): The model to be boosted.\n            optimizer (Optimizer): The optimizer to be boosted.\n            criterion (Callable): The criterion to be boosted.\n            dataloader (DataLoader): The dataloader to be boosted.\n            lr_scheduler (LRScheduler): The lr_scheduler to be boosted.\n        \"\"\"\n        # TODO(FrankLeeeee): consider multi-model and multi-optimizer case\n        # TODO(FrankLeeeee): consider multi-dataloader case\n        # transform model for mixed precision\n        if self.plugin:\n            model, optimizer, criterion, dataloader, lr_scheduler = self.plugin.configure(\n                model, optimizer, criterion, dataloader, lr_scheduler)\n\n        if self.plugin and not self.plugin.control_device():\n            # transform model for accelerator\n            model = self.accelerator.configure(model)\n\n        if self.mixed_precision and (self.plugin is None or self.plugin and not self.plugin.control_precision()):\n            # transform model for mixed precision\n            # when mixed_precision is specified and the plugin is not given or does not control the precision\n            model, optimizer, criterion = self.mixed_precision.configure(model, optimizer, criterion)\n\n        return model, optimizer, criterion, dataloader, lr_scheduler\n\n    def backward(self, loss: torch.Tensor, optimizer: Optimizer) -> None:\n        \"\"\"Backward pass.\n\n        Args:\n            loss (torch.Tensor): The loss to be backpropagated.\n            optimizer (Optimizer): The optimizer to be updated.\n        \"\"\"\n        # TODO: implement this method with plugin\n        optimizer.backward(loss)\n\n    def execute_pipeline(self,\n                         data_iter: Iterator,\n                         model: nn.Module,\n                         criterion: Callable[[torch.Tensor], torch.Tensor],\n                         optimizer: Optimizer,\n                         return_loss: bool = True,\n                         return_outputs: bool = False) -> Tuple[Optional[torch.Tensor], ...]:\n        # TODO: implement this method\n        # run pipeline forward backward pass\n        # return loss or outputs if needed\n        pass\n\n    def no_sync(self, model: nn.Module) -> contextmanager:\n        \"\"\"Context manager to disable gradient synchronization across DP process groups.\n\n        Args:\n            model (nn.Module): The model to be disabled gradient synchronization.\n\n        Returns:\n            contextmanager: Context to disable gradient synchronization.\n        \"\"\"\n        assert self.plugin is not None, f'no_sync is only enabled when a plugin is provided and the plugin supports no_sync.'\n        assert self.plugin.support_no_sync, f'The plugin {self.plugin.__class__.__name__} does not support no_sync.'\n        return self.plugin.no_sync(model)\n\n    def load_model(self, model: nn.Module, checkpoint: str, strict: bool = True):\n        \"\"\"Load model from checkpoint.\n\n        Args:\n            model (nn.Module): A model boosted by Booster.\n            checkpoint (str): Path to the checkpoint. It must be a local path.\n                It should be a directory path if the checkpoint is sharded. Otherwise, it should be a file path.\n            strict (bool, optional): whether to strictly enforce that the keys\n                in :attr:`state_dict` match the keys returned by this module's\n                :meth:`~torch.nn.Module.state_dict` function. Defaults to True.\n        \"\"\"\n        self.checkpoint_io.load_model(model, checkpoint, strict)\n\n    def save_model(self,\n                   model: nn.Module,\n                   checkpoint: str,\n                   prefix: str = None,\n                   shard: bool = False,\n                   size_per_shard: int = 1024):\n        \"\"\"Save model to checkpoint.\n\n        Args:\n            model (nn.Module): A model boosted by Booster.\n            checkpoint (str): Path to the checkpoint. It must be a local path.\n                It is a file path if ``shard=False``. Otherwise, it is a directory path.\n            prefix (str, optional): A prefix added to parameter and buffer\n                names to compose the keys in state_dict. Defaults to None.\n            shard (bool, optional): Whether to save checkpoint a sharded way.\n                If true, the checkpoint will be a folder. Otherwise, it will be a single file. Defaults to False.\n            size_per_shard (int, optional): Maximum size of checkpoint shard file in MB. This is useful only when ``shard=True``. Defaults to 1024.\n        \"\"\"\n        self.checkpoint_io.save_model(model, checkpoint=checkpoint, shard=shard, size_per_shard=size_per_shard)\n\n    def load_optimizer(self, optimizer: Optimizer, checkpoint: str):\n        \"\"\"Load optimizer from checkpoint.\n\n        Args:\n            optimizer (Optimizer): An optimizer boosted by Booster.\n            checkpoint (str): Path to the checkpoint. It must be a local path.\n                It should be a directory path if the checkpoint is sharded. Otherwise, it should be a file path.\n        \"\"\"\n        self.checkpoint_io.load_optimizer(optimizer, checkpoint)\n\n    def save_optimizer(self, optimizer: Optimizer, checkpoint: str, shard: bool = False, size_per_shard: int = 1024):\n        \"\"\"Save optimizer to checkpoint.\n        Warning: Saving sharded optimizer checkpoint is not supported yet.\n\n        Args:\n            optimizer (Optimizer): An optimizer boosted by Booster.\n            checkpoint (str): Path to the checkpoint. It must be a local path.\n                It is a file path if ``shard=False``. Otherwise, it is a directory path.\n            shard (bool, optional): Whether to save checkpoint a sharded way.\n                If true, the checkpoint will be a folder. Otherwise, it will be a single file. Defaults to False.\n            size_per_shard (int, optional): Maximum size of checkpoint shard file in MB. This is useful only when ``shard=True``. Defaults to 1024.\n        \"\"\"\n        self.checkpoint_io.save_optimizer(optimizer, checkpoint, shard, size_per_shard)\n\n    def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):\n        \"\"\"Save lr scheduler to checkpoint.\n\n        Args:\n            lr_scheduler (LRScheduler): A lr scheduler boosted by Booster.\n            checkpoint (str): Path to the checkpoint. It must be a local file path.\n        \"\"\"\n        self.checkpoint_io.save_lr_scheduler(lr_scheduler, checkpoint)\n\n    def load_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):\n        \"\"\"Load lr scheduler from checkpoint.\n\n        Args:\n            lr_scheduler (LRScheduler): A lr scheduler boosted by Booster.\n            checkpoint (str): Path to the checkpoint. It must be a local file path.\n        \"\"\"\n        self.checkpoint_io.load_lr_scheduler(lr_scheduler, checkpoint)\n", "path": "colossalai/booster/booster.py"}], "after_files": [{"content": "import warnings\nfrom contextlib import contextmanager\nfrom typing import Callable, Iterator, List, Optional, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import _LRScheduler as LRScheduler\nfrom torch.utils.data import DataLoader\n\nfrom colossalai.checkpoint_io import GeneralCheckpointIO\n\nfrom .accelerator import Accelerator\nfrom .mixed_precision import MixedPrecision, mixed_precision_factory\nfrom .plugin import Plugin\n\n__all__ = ['Booster']\n\n\nclass Booster:\n    \"\"\"\n    Booster is a high-level API for training neural networks. It provides a unified interface for\n    training with different precision, accelerator, and plugin.\n\n    Examples:\n        ```python\n        colossalai.launch(...)\n        plugin = GeminiPlugin(...)\n        booster = Booster(precision='fp16', plugin=plugin)\n\n        model = GPT2()\n        optimizer = HybridAdam(model.parameters())\n        dataloader = Dataloader(Dataset)\n        lr_scheduler = LinearWarmupScheduler()\n        criterion = GPTLMLoss()\n\n        model, optimizer, lr_scheduler, dataloader = booster.boost(model, optimizer, lr_scheduler, dataloader)\n\n        for epoch in range(max_epochs):\n            for input_ids, attention_mask in dataloader:\n                outputs = model(input_ids, attention_mask)\n                loss = criterion(outputs.logits, input_ids)\n                booster.backward(loss, optimizer)\n                optimizer.step()\n                lr_scheduler.step()\n                optimizer.zero_grad()\n        ```\n\n    Args:\n        device (str or torch.device): The device to run the training. Default: 'cuda'.\n        mixed_precision (str or MixedPrecision): The mixed precision to run the training. Default: None.\n                                If the argument is a string, it can be 'fp16', 'fp16_apex', 'bf16', or 'fp8'.\n                                'fp16' would use PyTorch AMP while `fp16_apex` would use Nvidia Apex.\n        plugin (Plugin): The plugin to run the training. Default: None.\n    \"\"\"\n\n    def __init__(self,\n                 device: str = 'cuda',\n                 mixed_precision: Union[MixedPrecision, str] = None,\n                 plugin: Optional[Plugin] = None) -> None:\n        if plugin is not None:\n            assert isinstance(\n                plugin, Plugin), f'Expected the argument plugin to be an instance of Plugin, but got {type(plugin)}.'\n        self.plugin = plugin\n\n        # set accelerator\n        if self.plugin and self.plugin.control_device():\n            self.accelerator = None\n            warnings.warn('The plugin will control the accelerator, so the device argument will be ignored.')\n        else:\n            self.accelerator = Accelerator(device)\n\n        # set precision\n        if self.plugin and self.plugin.control_precision():\n            warnings.warn('The plugin will control the precision, so the mixed_precision argument will be ignored.')\n            self.mixed_precision = None\n        elif mixed_precision is None:\n            self.mixed_precision = None\n        else:\n            # validate and set precision\n            if isinstance(mixed_precision, str):\n                # the user will take the default arguments for amp training\n                self.mixed_precision = mixed_precision_factory(mixed_precision)\n            elif isinstance(mixed_precision, MixedPrecision):\n                # the user can customize the arguments by passing the precision object\n                self.mixed_precision = mixed_precision\n            else:\n                raise ValueError(\n                    f'Expected the argument mixed_precision to be a string or an instance of Precision, but got {type(mixed_precision)}.'\n                )\n\n        if self.plugin is not None and self.plugin.control_checkpoint_io():\n            self.checkpoint_io = self.plugin.get_checkpoint_io()\n        else:\n            self.checkpoint_io = GeneralCheckpointIO()\n\n    def boost(\n        self,\n        model: nn.Module,\n        optimizer: Optimizer,\n        criterion: Callable = None,\n        dataloader: DataLoader = None,\n        lr_scheduler: LRScheduler = None,\n    ) -> List[Union[nn.Module, Optimizer, LRScheduler, DataLoader]]:\n        \"\"\"\n        Boost the model, optimizer, criterion, lr_scheduler, and dataloader.\n\n        Args:\n            model (nn.Module): The model to be boosted.\n            optimizer (Optimizer): The optimizer to be boosted.\n            criterion (Callable): The criterion to be boosted.\n            dataloader (DataLoader): The dataloader to be boosted.\n            lr_scheduler (LRScheduler): The lr_scheduler to be boosted.\n        \"\"\"\n        # TODO(FrankLeeeee): consider multi-model and multi-optimizer case\n        # TODO(FrankLeeeee): consider multi-dataloader case\n        # transform model for mixed precision\n        if self.plugin:\n            model, optimizer, criterion, dataloader, lr_scheduler = self.plugin.configure(\n                model, optimizer, criterion, dataloader, lr_scheduler)\n\n        if self.plugin and not self.plugin.control_device():\n            # transform model for accelerator\n            model = self.accelerator.configure(model)\n\n        if self.mixed_precision and (self.plugin is None or self.plugin and not self.plugin.control_precision()):\n            # transform model for mixed precision\n            # when mixed_precision is specified and the plugin is not given or does not control the precision\n            model, optimizer, criterion = self.mixed_precision.configure(model, optimizer, criterion)\n\n        return model, optimizer, criterion, dataloader, lr_scheduler\n\n    def backward(self, loss: torch.Tensor, optimizer: Optimizer) -> None:\n        \"\"\"Backward pass.\n\n        Args:\n            loss (torch.Tensor): The loss to be backpropagated.\n            optimizer (Optimizer): The optimizer to be updated.\n        \"\"\"\n        # TODO: implement this method with plugin\n        optimizer.backward(loss)\n\n    def execute_pipeline(self,\n                         data_iter: Iterator,\n                         model: nn.Module,\n                         criterion: Callable[[torch.Tensor], torch.Tensor],\n                         optimizer: Optimizer,\n                         return_loss: bool = True,\n                         return_outputs: bool = False) -> Tuple[Optional[torch.Tensor], ...]:\n        # TODO: implement this method\n        # run pipeline forward backward pass\n        # return loss or outputs if needed\n        pass\n\n    def no_sync(self, model: nn.Module) -> contextmanager:\n        \"\"\"Context manager to disable gradient synchronization across DP process groups.\n\n        Args:\n            model (nn.Module): The model to be disabled gradient synchronization.\n\n        Returns:\n            contextmanager: Context to disable gradient synchronization.\n        \"\"\"\n        assert self.plugin is not None, f'no_sync is only enabled when a plugin is provided and the plugin supports no_sync.'\n        assert self.plugin.support_no_sync, f'The plugin {self.plugin.__class__.__name__} does not support no_sync.'\n        return self.plugin.no_sync(model)\n\n    def load_model(self, model: nn.Module, checkpoint: str, strict: bool = True):\n        \"\"\"Load model from checkpoint.\n\n        Args:\n            model (nn.Module): A model boosted by Booster.\n            checkpoint (str): Path to the checkpoint. It must be a local path.\n                It should be a directory path if the checkpoint is sharded. Otherwise, it should be a file path.\n            strict (bool, optional): whether to strictly enforce that the keys\n                in :attr:`state_dict` match the keys returned by this module's\n                :meth:`~torch.nn.Module.state_dict` function. Defaults to True.\n        \"\"\"\n        self.checkpoint_io.load_model(model, checkpoint, strict)\n\n    def save_model(self,\n                   model: nn.Module,\n                   checkpoint: str,\n                   prefix: str = None,\n                   shard: bool = False,\n                   size_per_shard: int = 1024):\n        \"\"\"Save model to checkpoint.\n\n        Args:\n            model (nn.Module): A model boosted by Booster.\n            checkpoint (str): Path to the checkpoint. It must be a local path.\n                It is a file path if ``shard=False``. Otherwise, it is a directory path.\n            prefix (str, optional): A prefix added to parameter and buffer\n                names to compose the keys in state_dict. Defaults to None.\n            shard (bool, optional): Whether to save checkpoint a sharded way.\n                If true, the checkpoint will be a folder. Otherwise, it will be a single file. Defaults to False.\n            size_per_shard (int, optional): Maximum size of checkpoint shard file in MB. This is useful only when ``shard=True``. Defaults to 1024.\n        \"\"\"\n        self.checkpoint_io.save_model(model, checkpoint=checkpoint, shard=shard, size_per_shard=size_per_shard)\n\n    def load_optimizer(self, optimizer: Optimizer, checkpoint: str):\n        \"\"\"Load optimizer from checkpoint.\n\n        Args:\n            optimizer (Optimizer): An optimizer boosted by Booster.\n            checkpoint (str): Path to the checkpoint. It must be a local path.\n                It should be a directory path if the checkpoint is sharded. Otherwise, it should be a file path.\n        \"\"\"\n        self.checkpoint_io.load_optimizer(optimizer, checkpoint)\n\n    def save_optimizer(self, optimizer: Optimizer, checkpoint: str, shard: bool = False, size_per_shard: int = 1024):\n        \"\"\"Save optimizer to checkpoint.\n        Warning: Saving sharded optimizer checkpoint is not supported yet.\n\n        Args:\n            optimizer (Optimizer): An optimizer boosted by Booster.\n            checkpoint (str): Path to the checkpoint. It must be a local path.\n                It is a file path if ``shard=False``. Otherwise, it is a directory path.\n            shard (bool, optional): Whether to save checkpoint a sharded way.\n                If true, the checkpoint will be a folder. Otherwise, it will be a single file. Defaults to False.\n            size_per_shard (int, optional): Maximum size of checkpoint shard file in MB. This is useful only when ``shard=True``. Defaults to 1024.\n        \"\"\"\n        self.checkpoint_io.save_optimizer(optimizer, checkpoint, shard, size_per_shard)\n\n    def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):\n        \"\"\"Save lr scheduler to checkpoint.\n\n        Args:\n            lr_scheduler (LRScheduler): A lr scheduler boosted by Booster.\n            checkpoint (str): Path to the checkpoint. It must be a local file path.\n        \"\"\"\n        self.checkpoint_io.save_lr_scheduler(lr_scheduler, checkpoint)\n\n    def load_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):\n        \"\"\"Load lr scheduler from checkpoint.\n\n        Args:\n            lr_scheduler (LRScheduler): A lr scheduler boosted by Booster.\n            checkpoint (str): Path to the checkpoint. It must be a local file path.\n        \"\"\"\n        self.checkpoint_io.load_lr_scheduler(lr_scheduler, checkpoint)\n", "path": "colossalai/booster/booster.py"}]}