prompt: 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
  
  [booster] implement accelerator
  # Overview
  
  
  `Accelerator` is an abstraction for the hardware devices used for training. It can be CPU, GPU, TPU, etc. With this abstraction, we can provide room for future extension to more hardware which is compatible with PyTorch.
  
  Wanna track the development progress? Take a look at
  
  proposal: https://github.com/hpcaitech/ColossalAI/discussions/3046
  project kanban: https://github.com/orgs/hpcaitech/projects/19
  
  # Goal
  
  The `Accelerator` can move models to the target device. A sample definition is given below and subject to future changes.
  
  ```python
  
  class Accelerator:
  
          def __init__(self, device):
                self.device = device
  
          def setup_model(self, model) -> torch.nn.Module
  ```
  
  
  
  --- 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/accelerator.py`
  Content:
  ```
  1 import torch
  2 import torch.nn as nn
  3 
  4 __all__ = ['Accelerator']
  5 
  6 
  7 class Accelerator:
  8 
  9     def __init__(self, device: torch.device):
  10         self.device = device
  11 
  12     def setup_model(self, model: nn.Module) -> nn.Module:
  13         # TODO: implement this method
  14         pass
  15
  ```
  
  Path: `colossalai/booster/booster.py`
  Content:
  ```
  1 from contextlib import contextmanager
  2 from typing import Callable, Iterable, Iterator, List, Optional, Tuple, Union
  3 
  4 import torch
  5 import torch.nn as nn
  6 from torch import Tensor
  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 .mixed_precision import MixedPrecision, mixed_precision_factory
  12 from .plugin import Plugin
  13 
  14 __all__ = ['Booster']
  15 
  16 
  17 class Booster:
  18     """
  19     Booster is a high-level API for training neural networks. It provides a unified interface for
  20     training with different precisio, accelerator, and plugin.
  21 
  22     Examples:
  23         >>> colossalai.launch(...)
  24         >>> plugin = GeminiPlugin(stage=3, ...)
  25         >>> booster = Booster(precision='fp16', plugin=plugin)
  26         >>>
  27         >>> model = GPT2()
  28         >>> optimizer = Adam(model.parameters())
  29         >>> dataloader = Dataloader(Dataset)
  30         >>> lr_scheduler = LinearWarmupScheduler()
  31         >>> criterion = GPTLMLoss()
  32         >>>
  33         >>> model, optimizer, lr_scheduler, dataloader = booster.boost(model, optimizer, lr_scheduler, dataloader)
  34         >>>
  35         >>> for epoch in range(max_epochs):
  36         >>>     for input_ids, attention_mask in dataloader:
  37         >>>         outputs = model(input_ids, attention_mask)
  38         >>>         loss = criterion(outputs.logits, input_ids)
  39         >>>         booster.backward(loss, optimizer)
  40         >>>         optimizer.step()
  41         >>>         lr_scheduler.step()
  42         >>>         optimizer.zero_grad()
  43 
  44 
  45     Args:
  46         device (str or torch.device): The device to run the training. Default: 'cuda'.
  47         mixed_precision (str or MixedPrecision): The mixed precision to run the training. Default: None.
  48                                 If the argument is a string, it can be 'fp16', 'fp16_apex', 'bf16', or 'fp8'.
  49                                 'fp16' would use PyTorch AMP while `fp16_apex` would use Nvidia Apex.
  50         plugin (Plugin): The plugin to run the training. Default: None.
  51     """
  52 
  53     def __init__(self,
  54                  device: Union[str, torch.device] = 'cuda',
  55                  mixed_precision: Union[MixedPrecision, str] = None,
  56                  plugin: Optional[Plugin] = None) -> None:
  57         # validate and set precision
  58         if isinstance(MixedPrecision, str):
  59             # the user will take the default arguments for amp training
  60             self.mixed_precision = mixed_precision_factory(mixed_precision)
  61         elif isinstance(mixed_precision, MixedPrecision):
  62             # the user can customize the arguments by passing the precision object
  63             self.mixed_precision = mixed_precision
  64         else:
  65             raise ValueError(
  66                 f'Expected the argument mixed_precision to be a string or an instance of Precision, but got {type(mixed_precision)}.'
  67             )
  68 
  69     def boost(self, model: nn.Module, optimizer: Optimizer, criterion: Callable, lr_scheduler: LRScheduler,
  70               dataloader: DataLoader) -> List[Union[nn.Module, Optimizer, LRScheduler, DataLoader]]:
  71         """
  72         Boost the model, optimizer, criterion, lr_scheduler, and dataloader.
  73 
  74         Args:
  75             model (nn.Module): The model to be boosted.
  76             optimizer (Optimizer): The optimizer to be boosted.
  77             criterion (Callable): The criterion to be boosted.
  78             lr_scheduler (LRScheduler): The lr_scheduler to be boosted.
  79             dataloader (DataLoader): The dataloader to be boosted.
  80         """
  81         # TODO(FrankLeeeee): consider multi-model and multi-optimizer case
  82         # TODO(lsg): Add plugin control logic
  83         # e.g.
  84         # if self.plugin is not None and self.plugin.control_boost:
  85         #    ...
  86         # transform model for mixed precision
  87         model, optimizer, criterion = self.mixed_precision.configure(model, optimizer, criterion)
  88         return model, optimizer, criterion, lr_scheduler, dataloader
  89 
  90     def backward(self, loss: torch.Tensor, optimizer: Optimizer) -> None:
  91         # TODO: implement this method with plugin
  92         optimizer.backward(loss)
  93 
  94     def execute_pipeline(self,
  95                          data_iter: Iterator,
  96                          model: nn.Module,
  97                          criterion: Callable[[torch.Tensor], torch.Tensor],
  98                          optimizer: Optimizer,
  99                          return_loss: bool = True,
  100                          return_outputs: bool = False) -> Tuple[Optional[torch.Tensor], ...]:
  101         # TODO: implement this method
  102         # run pipeline forward backward pass
  103         # return loss or outputs if needed
  104         pass
  105 
  106     def no_sync(self, model: nn.Module) -> contextmanager:
  107         # TODO: implement this method
  108         pass
  109 
  110     def save(self,
  111              obj: Union[nn.Module, Optimizer, LRScheduler],
  112              path_like: str,
  113              plan: str = 'torch',
  114              **kwargs) -> None:
  115         # TODO: implement this method
  116         pass
  117 
  118     def load(self,
  119              obj: Union[nn.Module, Optimizer, LRScheduler],
  120              path_like: str,
  121              plan: str = 'torch',
  122              **kwargs) -> None:
  123         # TODO: implement this method
  124         pass
  125
  ```
  --- 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()
  
  ```
  
golden_diff: diff --git a/colossalai/booster/accelerator.py b/colossalai/booster/accelerator.py
  --- a/colossalai/booster/accelerator.py
  +++ b/colossalai/booster/accelerator.py
  @@ -3,12 +3,52 @@
   
   __all__ = ['Accelerator']
   
  +_supported_devices = [
  +    'cpu',
  +    'cuda',
  +
  +    # To be supported
  +    # 'xpu',
  +    # 'npu',
  +    # 'tpu',
  +]
  +
   
   class Accelerator:
  +    """
  +    Accelerator is an abstraction for the hardware device that is used to run the model.
  +
  +    Args:
  +        device (str): The device to be used. Currently only support 'cpu' and 'gpu'.
  +    """
   
  -    def __init__(self, device: torch.device):
  +    def __init__(self, device: str):
           self.device = device
   
  -    def setup_model(self, model: nn.Module) -> nn.Module:
  -        # TODO: implement this method
  -        pass
  +        assert self.device in _supported_devices, f"Device {self.device} is not supported yet, supported devices include {_supported_devices}"
  +
  +    def bind(self):
  +        """
  +        Set the default device for the current process.
  +        """
  +        if self.device == 'cpu':
  +            pass
  +        elif self.device == 'cuda':
  +            # TODO(FrankLeeeee): use global environment to check if it is a dist job
  +            # if is_distributed:
  +            #     local_rank = EnvTable().get_local_rank()
  +            #     torch.cuda.set_device(torch.device(f'cuda:{local_rank}'))
  +            torch.cuda.set_device(torch.device('cuda'))
  +            pass
  +        else:
  +            raise ValueError(f"Device {self.device} is not supported yet")
  +
  +    def configure_model(self, model: nn.Module) -> nn.Module:
  +        """
  +        Move the model to the device.
  +
  +        Args:
  +            model (nn.Module): The model to be moved.
  +        """
  +        model = model.to(torch.device(self.device))
  +        return model
  diff --git a/colossalai/booster/booster.py b/colossalai/booster/booster.py
  --- a/colossalai/booster/booster.py
  +++ b/colossalai/booster/booster.py
  @@ -8,6 +8,7 @@
   from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
   from torch.utils.data import DataLoader
   
  +from .accelerator import Accelerator
   from .mixed_precision import MixedPrecision, mixed_precision_factory
   from .plugin import Plugin
   
  @@ -51,9 +52,16 @@
       """
   
       def __init__(self,
  -                 device: Union[str, torch.device] = 'cuda',
  +                 device: str = 'cuda',
                    mixed_precision: Union[MixedPrecision, str] = None,
                    plugin: Optional[Plugin] = None) -> None:
  +        # TODO(FrankLeeeee): add plugin control logic
  +        # if self.plugin is not None and self.plugin.control_accelerator:
  +        #     ...
  +        # create acclerator
  +        self.acceleartor = Accelerator(device)
  +        self.acceleartor.set_default_device()
  +
           # validate and set precision
           if isinstance(MixedPrecision, str):
               # the user will take the default arguments for amp training
  @@ -78,6 +86,11 @@
               lr_scheduler (LRScheduler): The lr_scheduler to be boosted.
               dataloader (DataLoader): The dataloader to be boosted.
           """
  +        # TODO(FrankLeeeee): add plugin control logic
  +        # if self.plugin is not None and self.plugin.control_accelerator:
  +        #     ...
  +        model = self.acceleartor.configure_model(model)
  +
           # TODO(FrankLeeeee): consider multi-model and multi-optimizer case
           # TODO(lsg): Add plugin control logic
           # e.g.
  
verification_info: {"golden_diff": "diff --git a/colossalai/booster/accelerator.py b/colossalai/booster/accelerator.py\n--- a/colossalai/booster/accelerator.py\n+++ b/colossalai/booster/accelerator.py\n@@ -3,12 +3,52 @@\n \n __all__ = ['Accelerator']\n \n+_supported_devices = [\n+    'cpu',\n+    'cuda',\n+\n+    # To be supported\n+    # 'xpu',\n+    # 'npu',\n+    # 'tpu',\n+]\n+\n \n class Accelerator:\n+    \"\"\"\n+    Accelerator is an abstraction for the hardware device that is used to run the model.\n+\n+    Args:\n+        device (str): The device to be used. Currently only support 'cpu' and 'gpu'.\n+    \"\"\"\n \n-    def __init__(self, device: torch.device):\n+    def __init__(self, device: str):\n         self.device = device\n \n-    def setup_model(self, model: nn.Module) -> nn.Module:\n-        # TODO: implement this method\n-        pass\n+        assert self.device in _supported_devices, f\"Device {self.device} is not supported yet, supported devices include {_supported_devices}\"\n+\n+    def bind(self):\n+        \"\"\"\n+        Set the default device for the current process.\n+        \"\"\"\n+        if self.device == 'cpu':\n+            pass\n+        elif self.device == 'cuda':\n+            # TODO(FrankLeeeee): use global environment to check if it is a dist job\n+            # if is_distributed:\n+            #     local_rank = EnvTable().get_local_rank()\n+            #     torch.cuda.set_device(torch.device(f'cuda:{local_rank}'))\n+            torch.cuda.set_device(torch.device('cuda'))\n+            pass\n+        else:\n+            raise ValueError(f\"Device {self.device} is not supported yet\")\n+\n+    def configure_model(self, model: nn.Module) -> nn.Module:\n+        \"\"\"\n+        Move the model to the device.\n+\n+        Args:\n+            model (nn.Module): The model to be moved.\n+        \"\"\"\n+        model = model.to(torch.device(self.device))\n+        return model\ndiff --git a/colossalai/booster/booster.py b/colossalai/booster/booster.py\n--- a/colossalai/booster/booster.py\n+++ b/colossalai/booster/booster.py\n@@ -8,6 +8,7 @@\n from torch.optim.lr_scheduler import _LRScheduler as LRScheduler\n from torch.utils.data import DataLoader\n \n+from .accelerator import Accelerator\n from .mixed_precision import MixedPrecision, mixed_precision_factory\n from .plugin import Plugin\n \n@@ -51,9 +52,16 @@\n     \"\"\"\n \n     def __init__(self,\n-                 device: Union[str, torch.device] = 'cuda',\n+                 device: str = 'cuda',\n                  mixed_precision: Union[MixedPrecision, str] = None,\n                  plugin: Optional[Plugin] = None) -> None:\n+        # TODO(FrankLeeeee): add plugin control logic\n+        # if self.plugin is not None and self.plugin.control_accelerator:\n+        #     ...\n+        # create acclerator\n+        self.acceleartor = Accelerator(device)\n+        self.acceleartor.set_default_device()\n+\n         # validate and set precision\n         if isinstance(MixedPrecision, str):\n             # the user will take the default arguments for amp training\n@@ -78,6 +86,11 @@\n             lr_scheduler (LRScheduler): The lr_scheduler to be boosted.\n             dataloader (DataLoader): The dataloader to be boosted.\n         \"\"\"\n+        # TODO(FrankLeeeee): add plugin control logic\n+        # if self.plugin is not None and self.plugin.control_accelerator:\n+        #     ...\n+        model = self.acceleartor.configure_model(model)\n+\n         # TODO(FrankLeeeee): consider multi-model and multi-optimizer case\n         # TODO(lsg): Add plugin control logic\n         # e.g.\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[booster] implement accelerator\n# Overview\n\n\n`Accelerator` is an abstraction for the hardware devices used for training. It can be CPU, GPU, TPU, etc. With this abstraction, we can provide room for future extension to more hardware which is compatible with PyTorch.\n\nWanna track the development progress? Take a look at\n\nproposal: https://github.com/hpcaitech/ColossalAI/discussions/3046\nproject kanban: https://github.com/orgs/hpcaitech/projects/19\n\n# Goal\n\nThe `Accelerator` can move models to the target device. A sample definition is given below and subject to future changes.\n\n```python\n\nclass Accelerator:\n\n        def __init__(self, device):\n              self.device = device\n\n        def setup_model(self, model) -> torch.nn.Module\n```\n\n\n", "before_files": [{"content": "import torch\nimport torch.nn as nn\n\n__all__ = ['Accelerator']\n\n\nclass Accelerator:\n\n    def __init__(self, device: torch.device):\n        self.device = device\n\n    def setup_model(self, model: nn.Module) -> nn.Module:\n        # TODO: implement this method\n        pass\n", "path": "colossalai/booster/accelerator.py"}, {"content": "from contextlib import contextmanager\nfrom typing import Callable, Iterable, Iterator, List, Optional, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nfrom torch import Tensor\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import _LRScheduler as LRScheduler\nfrom torch.utils.data import DataLoader\n\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 precisio, accelerator, and plugin.\n\n    Examples:\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: Union[str, torch.device] = 'cuda',\n                 mixed_precision: Union[MixedPrecision, str] = None,\n                 plugin: Optional[Plugin] = None) -> None:\n        # validate and set precision\n        if isinstance(MixedPrecision, 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    def boost(self, model: nn.Module, optimizer: Optimizer, criterion: Callable, lr_scheduler: LRScheduler,\n              dataloader: DataLoader) -> 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            lr_scheduler (LRScheduler): The lr_scheduler to be boosted.\n            dataloader (DataLoader): The dataloader to be boosted.\n        \"\"\"\n        # TODO(FrankLeeeee): consider multi-model and multi-optimizer case\n        # TODO(lsg): Add plugin control logic\n        # e.g.\n        # if self.plugin is not None and self.plugin.control_boost:\n        #    ...\n        # transform model for mixed precision\n        model, optimizer, criterion = self.mixed_precision.configure(model, optimizer, criterion)\n        return model, optimizer, criterion, lr_scheduler, dataloader\n\n    def backward(self, loss: torch.Tensor, optimizer: Optimizer) -> None:\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        # TODO: implement this method\n        pass\n\n    def save(self,\n             obj: Union[nn.Module, Optimizer, LRScheduler],\n             path_like: str,\n             plan: str = 'torch',\n             **kwargs) -> None:\n        # TODO: implement this method\n        pass\n\n    def load(self,\n             obj: Union[nn.Module, Optimizer, LRScheduler],\n             path_like: str,\n             plan: str = 'torch',\n             **kwargs) -> None:\n        # TODO: implement this method\n        pass\n", "path": "colossalai/booster/booster.py"}], "after_files": [{"content": "import torch\nimport torch.nn as nn\n\n__all__ = ['Accelerator']\n\n_supported_devices = [\n    'cpu',\n    'cuda',\n\n    # To be supported\n    # 'xpu',\n    # 'npu',\n    # 'tpu',\n]\n\n\nclass Accelerator:\n    \"\"\"\n    Accelerator is an abstraction for the hardware device that is used to run the model.\n\n    Args:\n        device (str): The device to be used. Currently only support 'cpu' and 'gpu'.\n    \"\"\"\n\n    def __init__(self, device: str):\n        self.device = device\n\n        assert self.device in _supported_devices, f\"Device {self.device} is not supported yet, supported devices include {_supported_devices}\"\n\n    def bind(self):\n        \"\"\"\n        Set the default device for the current process.\n        \"\"\"\n        if self.device == 'cpu':\n            pass\n        elif self.device == 'cuda':\n            # TODO(FrankLeeeee): use global environment to check if it is a dist job\n            # if is_distributed:\n            #     local_rank = EnvTable().get_local_rank()\n            #     torch.cuda.set_device(torch.device(f'cuda:{local_rank}'))\n            torch.cuda.set_device(torch.device('cuda'))\n            pass\n        else:\n            raise ValueError(f\"Device {self.device} is not supported yet\")\n\n    def configure_model(self, model: nn.Module) -> nn.Module:\n        \"\"\"\n        Move the model to the device.\n\n        Args:\n            model (nn.Module): The model to be moved.\n        \"\"\"\n        model = model.to(torch.device(self.device))\n        return model\n", "path": "colossalai/booster/accelerator.py"}, {"content": "from contextlib import contextmanager\nfrom typing import Callable, Iterable, Iterator, List, Optional, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nfrom torch import Tensor\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import _LRScheduler as LRScheduler\nfrom torch.utils.data import DataLoader\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 precisio, accelerator, and plugin.\n\n    Examples:\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        # TODO(FrankLeeeee): add plugin control logic\n        # if self.plugin is not None and self.plugin.control_accelerator:\n        #     ...\n        # create acclerator\n        self.acceleartor = Accelerator(device)\n        self.acceleartor.set_default_device()\n\n        # validate and set precision\n        if isinstance(MixedPrecision, 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    def boost(self, model: nn.Module, optimizer: Optimizer, criterion: Callable, lr_scheduler: LRScheduler,\n              dataloader: DataLoader) -> 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            lr_scheduler (LRScheduler): The lr_scheduler to be boosted.\n            dataloader (DataLoader): The dataloader to be boosted.\n        \"\"\"\n        # TODO(FrankLeeeee): add plugin control logic\n        # if self.plugin is not None and self.plugin.control_accelerator:\n        #     ...\n        model = self.acceleartor.configure_model(model)\n\n        # TODO(FrankLeeeee): consider multi-model and multi-optimizer case\n        # TODO(lsg): Add plugin control logic\n        # e.g.\n        # if self.plugin is not None and self.plugin.control_boost:\n        #    ...\n        # transform model for mixed precision\n        model, optimizer, criterion = self.mixed_precision.configure(model, optimizer, criterion)\n        return model, optimizer, criterion, lr_scheduler, dataloader\n\n    def backward(self, loss: torch.Tensor, optimizer: Optimizer) -> None:\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        # TODO: implement this method\n        pass\n\n    def save(self,\n             obj: Union[nn.Module, Optimizer, LRScheduler],\n             path_like: str,\n             plan: str = 'torch',\n             **kwargs) -> None:\n        # TODO: implement this method\n        pass\n\n    def load(self,\n             obj: Union[nn.Module, Optimizer, LRScheduler],\n             path_like: str,\n             plan: str = 'torch',\n             **kwargs) -> None:\n        # TODO: implement this method\n        pass\n", "path": "colossalai/booster/booster.py"}]}