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: `colossalai/booster/plugin/torch_fsdp_plugin.py`
Content:
```
1 from pathlib import Path
2 from typing import Callable, Iterable, Iterator, List, Optional, Tuple, Union
3 
4 import torch
5 import torch.nn as nn
6 from packaging import version
7 from torch.distributed import ProcessGroup
8 
9 
10 if version.parse(torch.__version__) >= version.parse('1.12.0'):
11     from torch.distributed.fsdp import FullStateDictConfig
12     from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
13     from torch.distributed.fsdp import StateDictType
14     from torch.distributed.fsdp.fully_sharded_data_parallel import (
15         BackwardPrefetch,
16         CPUOffload,
17         FullStateDictConfig,
18         MixedPrecision,
19         ShardingStrategy,
20     )
21 else:
22     raise RuntimeError("FSDP is not supported while torch version under 1.12.0.")
23 
24 from torch.optim import Optimizer
25 from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
26 from torch.utils.data import DataLoader
27 
28 from colossalai.checkpoint_io import CheckpointIO, GeneralCheckpointIO, utils
29 from colossalai.cluster import DistCoordinator
30 from colossalai.interface import ModelWrapper, OptimizerWrapper
31 
32 from .dp_plugin_base import DPPluginBase
33 
34 __all__ = ['TorchFSDPPlugin']
35 
36 
37 class TorchFSDPCheckpointIO(GeneralCheckpointIO):
38 
39     def __init__(self) -> None:
40         super().__init__()
41         self.coordinator = DistCoordinator()
42 
43     def load_unsharded_model(self, model: nn.Module, checkpoint: str, strict: bool):
44         checkpoint = utils.load_state_dict(checkpoint)
45         model.load_state_dict(checkpoint)
46 
47     def load_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: Path):
48         checkpoint = utils.load_state_dict(checkpoint)
49         fsdp_model = optimizer.unwrap_model()
50         sharded_osd = FSDP.scatter_full_optim_state_dict(checkpoint, fsdp_model)
51         optimizer.load_state_dict(sharded_osd)
52 
53     def save_unsharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, use_safetensors: bool):
54         """
55         Save model to checkpoint but only on master process.
56         """
57         # the model should be unwrapped in self.load_model via ModelWrapper.unwrap
58         cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
59         with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg):
60             full_model_state = model.state_dict()
61         utils.save_state_dict(full_model_state, checkpoint_file_path=checkpoint, use_safetensors=use_safetensors)
62 
63     def save_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool):
64         """
65         Save optimizer to checkpoint but only on master process.
66         """
67         assert isinstance(optimizer, FSDPOptimizerWrapper)
68         fsdp_model = optimizer.unwrap_model()
69         full_optimizer_state = FSDP.full_optim_state_dict(fsdp_model, optim=optimizer, rank0_only=True)
70         utils.save_state_dict(full_optimizer_state, checkpoint_file_path=checkpoint, use_safetensors=False)
71 
72     def save_sharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, variant: Optional[str],
73                            size_per_shard: int, use_safetensors: bool):
74         """
75         Save model to checkpoint but only on master process.
76         """
77         raise NotImplementedError("Sharded model checkpoint is not supported yet.")
78 
79     def load_sharded_model(self,
80                            model: nn.Module,
81                            checkpoint_index_file: Path,
82                            strict: bool = False,
83                            use_safetensors: bool = False,
84                            load_sub_module: bool = True):
85         """
86         Load model to checkpoint but only on master process.
87         """
88         raise NotImplementedError("Sharded model checkpoint is not supported yet.")
89 
90     def save_sharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool):
91         """
92         Save optimizer to checkpoint but only on master process.
93         """
94         raise NotImplementedError("Sharded optimizer checkpoint is not supported yet.")
95 
96     def load_sharded_optimizer(self, optimizer: Optimizer, index_file_path: str, prefix: str, size_per_shard: int):
97         """
98         Load optimizer to checkpoint but only on master process.
99         """
100         raise NotImplementedError("Sharded optimizer checkpoint is not supported yet.")
101 
102     def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):
103         """
104         Save model to checkpoint but only on master process.
105         """
106         if self.coordinator.is_master():
107             super().save_lr_scheduler(lr_scheduler, checkpoint)
108 
109 
110 class TorchFSDPModel(ModelWrapper):
111 
112     def __init__(self, module: nn.Module, *args, **kwargs) -> None:
113         super().__init__(module)
114         self.module = FSDP(module, *args, **kwargs)
115 
116     def unwrap(self):
117         return self.module
118 
119 
120 class FSDPOptimizerWrapper(OptimizerWrapper):
121 
122     def __init__(self, optimizer: Optimizer, model: nn.Module):
123         self.model = model
124         super().__init__(optimizer)
125 
126     def unwrap_model(self) -> nn.Module:
127         return self.model
128 
129 
130 class TorchFSDPPlugin(DPPluginBase):
131     """
132     Plugin for PyTorch FSDP.
133 
134     Example:
135         >>> from colossalai.booster import Booster
136         >>> from colossalai.booster.plugin import TorchFSDPPlugin
137         >>>
138         >>> model, train_dataset, optimizer, criterion = ...
139         >>> plugin = TorchFSDPPlugin()
140 
141         >>> train_dataloader = plugin.prepare_train_dataloader(train_dataset, batch_size=8)
142         >>> booster = Booster(plugin=plugin)
143         >>> model, optimizer, train_dataloader, criterion = booster.boost(model, optimizer, train_dataloader, criterion)
144 
145     Args:
146         See https://pytorch.org/docs/stable/fsdp.html for details.
147     """
148 
149     if version.parse(torch.__version__) >= version.parse('1.12.0'):
150 
151         def __init__(
152             self,
153             process_group: Optional[ProcessGroup] = None,
154             sharding_strategy: Optional[ShardingStrategy] = None,
155             cpu_offload: Optional[CPUOffload] = None,
156             auto_wrap_policy: Optional[Callable] = None,
157             backward_prefetch: Optional[BackwardPrefetch] = None,
158             mixed_precision: Optional[MixedPrecision] = None,
159             ignored_modules: Optional[Iterable[torch.nn.Module]] = None,
160             param_init_fn: Optional[Callable[[nn.Module], None]] = None,
161             sync_module_states: bool = False,
162         ):
163             super().__init__()
164             self.fsdp_kwargs = dict(process_group=process_group,
165                                     sharding_strategy=sharding_strategy,
166                                     cpu_offload=cpu_offload,
167                                     auto_wrap_policy=auto_wrap_policy,
168                                     backward_prefetch=backward_prefetch,
169                                     mixed_precision=mixed_precision,
170                                     ignored_modules=ignored_modules,
171                                     param_init_fn=param_init_fn,
172                                     sync_module_states=sync_module_states)
173     else:
174         raise RuntimeError("FSDP is not supported while torch version under 1.12.0.")
175 
176     def support_no_sync(self) -> bool:
177         False
178 
179     def no_sync(self, model: nn.Module) -> Iterator[None]:
180         raise NotImplementedError("Torch fsdp no_sync func not supported yet.")
181 
182     def control_precision(self) -> bool:
183         return True
184 
185     def supported_precisions(self) -> List[str]:
186         return ['fp16', 'bf16']
187 
188     def control_device(self) -> bool:
189         return True
190 
191     def supported_devices(self) -> List[str]:
192         return ['cuda']
193 
194     def configure(
195         self,
196         model: nn.Module,
197         optimizer: Optimizer,
198         criterion: Callable = None,
199         dataloader: DataLoader = None,
200         lr_scheduler: LRScheduler = None,
201     ) -> Tuple[Union[nn.Module, OptimizerWrapper, LRScheduler, DataLoader]]:
202 
203         # wrap the model with PyTorch FSDP
204         fsdp_model = TorchFSDPModel(model, device_id=torch.cuda.current_device(), **self.fsdp_kwargs)
205         optimizer.__init__(fsdp_model.parameters(), **optimizer.defaults)
206 
207         if not isinstance(optimizer, FSDPOptimizerWrapper):
208             optimizer = FSDPOptimizerWrapper(optimizer, fsdp_model)
209 
210         return fsdp_model, optimizer, criterion, dataloader, lr_scheduler
211 
212     def control_checkpoint_io(self) -> bool:
213         return True
214 
215     def get_checkpoint_io(self) -> CheckpointIO:
216         return TorchFSDPCheckpointIO()
217
```
--- 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/plugin/torch_fsdp_plugin.py b/colossalai/booster/plugin/torch_fsdp_plugin.py
--- a/colossalai/booster/plugin/torch_fsdp_plugin.py
+++ b/colossalai/booster/plugin/torch_fsdp_plugin.py
@@ -3,10 +3,10 @@
 
 import torch
 import torch.nn as nn
+import warnings
 from packaging import version
 from torch.distributed import ProcessGroup
 
-
 if version.parse(torch.__version__) >= version.parse('1.12.0'):
     from torch.distributed.fsdp import FullStateDictConfig
     from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
@@ -202,6 +202,11 @@
 
         # wrap the model with PyTorch FSDP
         fsdp_model = TorchFSDPModel(model, device_id=torch.cuda.current_device(), **self.fsdp_kwargs)
+
+        if len(optimizer.param_groups) > 1:
+            warnings.warn(
+                'TorchFSDPPlugin does not support optimizer that use multi param groups. The results may not be as expected if used.'
+            )
         optimizer.__init__(fsdp_model.parameters(), **optimizer.defaults)
 
         if not isinstance(optimizer, FSDPOptimizerWrapper):

{"golden_diff": "diff --git a/colossalai/booster/plugin/torch_fsdp_plugin.py b/colossalai/booster/plugin/torch_fsdp_plugin.py\n--- a/colossalai/booster/plugin/torch_fsdp_plugin.py\n+++ b/colossalai/booster/plugin/torch_fsdp_plugin.py\n@@ -3,10 +3,10 @@\n \n import torch\n import torch.nn as nn\n+import warnings\n from packaging import version\n from torch.distributed import ProcessGroup\n \n-\n if version.parse(torch.__version__) >= version.parse('1.12.0'):\n     from torch.distributed.fsdp import FullStateDictConfig\n     from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n@@ -202,6 +202,11 @@\n \n         # wrap the model with PyTorch FSDP\n         fsdp_model = TorchFSDPModel(model, device_id=torch.cuda.current_device(), **self.fsdp_kwargs)\n+\n+        if len(optimizer.param_groups) > 1:\n+            warnings.warn(\n+                'TorchFSDPPlugin does not support optimizer that use multi param groups. The results may not be as expected if used.'\n+            )\n         optimizer.__init__(fsdp_model.parameters(), **optimizer.defaults)\n \n         if not isinstance(optimizer, FSDPOptimizerWrapper):\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "from pathlib import Path\nfrom typing import Callable, Iterable, Iterator, List, Optional, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nfrom packaging import version\nfrom torch.distributed import ProcessGroup\n\n\nif version.parse(torch.__version__) >= version.parse('1.12.0'):\n    from torch.distributed.fsdp import FullStateDictConfig\n    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n    from torch.distributed.fsdp import StateDictType\n    from torch.distributed.fsdp.fully_sharded_data_parallel import (\n        BackwardPrefetch,\n        CPUOffload,\n        FullStateDictConfig,\n        MixedPrecision,\n        ShardingStrategy,\n    )\nelse:\n    raise RuntimeError(\"FSDP is not supported while torch version under 1.12.0.\")\n\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 CheckpointIO, GeneralCheckpointIO, utils\nfrom colossalai.cluster import DistCoordinator\nfrom colossalai.interface import ModelWrapper, OptimizerWrapper\n\nfrom .dp_plugin_base import DPPluginBase\n\n__all__ = ['TorchFSDPPlugin']\n\n\nclass TorchFSDPCheckpointIO(GeneralCheckpointIO):\n\n    def __init__(self) -> None:\n        super().__init__()\n        self.coordinator = DistCoordinator()\n\n    def load_unsharded_model(self, model: nn.Module, checkpoint: str, strict: bool):\n        checkpoint = utils.load_state_dict(checkpoint)\n        model.load_state_dict(checkpoint)\n\n    def load_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: Path):\n        checkpoint = utils.load_state_dict(checkpoint)\n        fsdp_model = optimizer.unwrap_model()\n        sharded_osd = FSDP.scatter_full_optim_state_dict(checkpoint, fsdp_model)\n        optimizer.load_state_dict(sharded_osd)\n\n    def save_unsharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, use_safetensors: bool):\n        \"\"\"\n        Save model to checkpoint but only on master process.\n        \"\"\"\n        # the model should be unwrapped in self.load_model via ModelWrapper.unwrap\n        cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)\n        with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg):\n            full_model_state = model.state_dict()\n        utils.save_state_dict(full_model_state, checkpoint_file_path=checkpoint, use_safetensors=use_safetensors)\n\n    def save_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool):\n        \"\"\"\n        Save optimizer to checkpoint but only on master process.\n        \"\"\"\n        assert isinstance(optimizer, FSDPOptimizerWrapper)\n        fsdp_model = optimizer.unwrap_model()\n        full_optimizer_state = FSDP.full_optim_state_dict(fsdp_model, optim=optimizer, rank0_only=True)\n        utils.save_state_dict(full_optimizer_state, checkpoint_file_path=checkpoint, use_safetensors=False)\n\n    def save_sharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, variant: Optional[str],\n                           size_per_shard: int, use_safetensors: bool):\n        \"\"\"\n        Save model to checkpoint but only on master process.\n        \"\"\"\n        raise NotImplementedError(\"Sharded model checkpoint is not supported yet.\")\n\n    def load_sharded_model(self,\n                           model: nn.Module,\n                           checkpoint_index_file: Path,\n                           strict: bool = False,\n                           use_safetensors: bool = False,\n                           load_sub_module: bool = True):\n        \"\"\"\n        Load model to checkpoint but only on master process.\n        \"\"\"\n        raise NotImplementedError(\"Sharded model checkpoint is not supported yet.\")\n\n    def save_sharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool):\n        \"\"\"\n        Save optimizer to checkpoint but only on master process.\n        \"\"\"\n        raise NotImplementedError(\"Sharded optimizer checkpoint is not supported yet.\")\n\n    def load_sharded_optimizer(self, optimizer: Optimizer, index_file_path: str, prefix: str, size_per_shard: int):\n        \"\"\"\n        Load optimizer to checkpoint but only on master process.\n        \"\"\"\n        raise NotImplementedError(\"Sharded optimizer checkpoint is not supported yet.\")\n\n    def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):\n        \"\"\"\n        Save model to checkpoint but only on master process.\n        \"\"\"\n        if self.coordinator.is_master():\n            super().save_lr_scheduler(lr_scheduler, checkpoint)\n\n\nclass TorchFSDPModel(ModelWrapper):\n\n    def __init__(self, module: nn.Module, *args, **kwargs) -> None:\n        super().__init__(module)\n        self.module = FSDP(module, *args, **kwargs)\n\n    def unwrap(self):\n        return self.module\n\n\nclass FSDPOptimizerWrapper(OptimizerWrapper):\n\n    def __init__(self, optimizer: Optimizer, model: nn.Module):\n        self.model = model\n        super().__init__(optimizer)\n\n    def unwrap_model(self) -> nn.Module:\n        return self.model\n\n\nclass TorchFSDPPlugin(DPPluginBase):\n    \"\"\"\n    Plugin for PyTorch FSDP.\n\n    Example:\n        >>> from colossalai.booster import Booster\n        >>> from colossalai.booster.plugin import TorchFSDPPlugin\n        >>>\n        >>> model, train_dataset, optimizer, criterion = ...\n        >>> plugin = TorchFSDPPlugin()\n\n        >>> train_dataloader = plugin.prepare_train_dataloader(train_dataset, batch_size=8)\n        >>> booster = Booster(plugin=plugin)\n        >>> model, optimizer, train_dataloader, criterion = booster.boost(model, optimizer, train_dataloader, criterion)\n\n    Args:\n        See https://pytorch.org/docs/stable/fsdp.html for details.\n    \"\"\"\n\n    if version.parse(torch.__version__) >= version.parse('1.12.0'):\n\n        def __init__(\n            self,\n            process_group: Optional[ProcessGroup] = None,\n            sharding_strategy: Optional[ShardingStrategy] = None,\n            cpu_offload: Optional[CPUOffload] = None,\n            auto_wrap_policy: Optional[Callable] = None,\n            backward_prefetch: Optional[BackwardPrefetch] = None,\n            mixed_precision: Optional[MixedPrecision] = None,\n            ignored_modules: Optional[Iterable[torch.nn.Module]] = None,\n            param_init_fn: Optional[Callable[[nn.Module], None]] = None,\n            sync_module_states: bool = False,\n        ):\n            super().__init__()\n            self.fsdp_kwargs = dict(process_group=process_group,\n                                    sharding_strategy=sharding_strategy,\n                                    cpu_offload=cpu_offload,\n                                    auto_wrap_policy=auto_wrap_policy,\n                                    backward_prefetch=backward_prefetch,\n                                    mixed_precision=mixed_precision,\n                                    ignored_modules=ignored_modules,\n                                    param_init_fn=param_init_fn,\n                                    sync_module_states=sync_module_states)\n    else:\n        raise RuntimeError(\"FSDP is not supported while torch version under 1.12.0.\")\n\n    def support_no_sync(self) -> bool:\n        False\n\n    def no_sync(self, model: nn.Module) -> Iterator[None]:\n        raise NotImplementedError(\"Torch fsdp no_sync func not supported yet.\")\n\n    def control_precision(self) -> bool:\n        return True\n\n    def supported_precisions(self) -> List[str]:\n        return ['fp16', 'bf16']\n\n    def control_device(self) -> bool:\n        return True\n\n    def supported_devices(self) -> List[str]:\n        return ['cuda']\n\n    def configure(\n        self,\n        model: nn.Module,\n        optimizer: Optimizer,\n        criterion: Callable = None,\n        dataloader: DataLoader = None,\n        lr_scheduler: LRScheduler = None,\n    ) -> Tuple[Union[nn.Module, OptimizerWrapper, LRScheduler, DataLoader]]:\n\n        # wrap the model with PyTorch FSDP\n        fsdp_model = TorchFSDPModel(model, device_id=torch.cuda.current_device(), **self.fsdp_kwargs)\n        optimizer.__init__(fsdp_model.parameters(), **optimizer.defaults)\n\n        if not isinstance(optimizer, FSDPOptimizerWrapper):\n            optimizer = FSDPOptimizerWrapper(optimizer, fsdp_model)\n\n        return fsdp_model, optimizer, criterion, dataloader, lr_scheduler\n\n    def control_checkpoint_io(self) -> bool:\n        return True\n\n    def get_checkpoint_io(self) -> CheckpointIO:\n        return TorchFSDPCheckpointIO()\n", "path": "colossalai/booster/plugin/torch_fsdp_plugin.py"}], "after_files": [{"content": "from pathlib import Path\nfrom typing import Callable, Iterable, Iterator, List, Optional, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nimport warnings\nfrom packaging import version\nfrom torch.distributed import ProcessGroup\n\nif version.parse(torch.__version__) >= version.parse('1.12.0'):\n    from torch.distributed.fsdp import FullStateDictConfig\n    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n    from torch.distributed.fsdp import StateDictType\n    from torch.distributed.fsdp.fully_sharded_data_parallel import (\n        BackwardPrefetch,\n        CPUOffload,\n        FullStateDictConfig,\n        MixedPrecision,\n        ShardingStrategy,\n    )\nelse:\n    raise RuntimeError(\"FSDP is not supported while torch version under 1.12.0.\")\n\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 CheckpointIO, GeneralCheckpointIO, utils\nfrom colossalai.cluster import DistCoordinator\nfrom colossalai.interface import ModelWrapper, OptimizerWrapper\n\nfrom .dp_plugin_base import DPPluginBase\n\n__all__ = ['TorchFSDPPlugin']\n\n\nclass TorchFSDPCheckpointIO(GeneralCheckpointIO):\n\n    def __init__(self) -> None:\n        super().__init__()\n        self.coordinator = DistCoordinator()\n\n    def load_unsharded_model(self, model: nn.Module, checkpoint: str, strict: bool):\n        checkpoint = utils.load_state_dict(checkpoint)\n        model.load_state_dict(checkpoint)\n\n    def load_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: Path):\n        checkpoint = utils.load_state_dict(checkpoint)\n        fsdp_model = optimizer.unwrap_model()\n        sharded_osd = FSDP.scatter_full_optim_state_dict(checkpoint, fsdp_model)\n        optimizer.load_state_dict(sharded_osd)\n\n    def save_unsharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, use_safetensors: bool):\n        \"\"\"\n        Save model to checkpoint but only on master process.\n        \"\"\"\n        # the model should be unwrapped in self.load_model via ModelWrapper.unwrap\n        cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)\n        with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg):\n            full_model_state = model.state_dict()\n        utils.save_state_dict(full_model_state, checkpoint_file_path=checkpoint, use_safetensors=use_safetensors)\n\n    def save_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool):\n        \"\"\"\n        Save optimizer to checkpoint but only on master process.\n        \"\"\"\n        assert isinstance(optimizer, FSDPOptimizerWrapper)\n        fsdp_model = optimizer.unwrap_model()\n        full_optimizer_state = FSDP.full_optim_state_dict(fsdp_model, optim=optimizer, rank0_only=True)\n        utils.save_state_dict(full_optimizer_state, checkpoint_file_path=checkpoint, use_safetensors=False)\n\n    def save_sharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, variant: Optional[str],\n                           size_per_shard: int, use_safetensors: bool):\n        \"\"\"\n        Save model to checkpoint but only on master process.\n        \"\"\"\n        raise NotImplementedError(\"Sharded model checkpoint is not supported yet.\")\n\n    def load_sharded_model(self,\n                           model: nn.Module,\n                           checkpoint_index_file: Path,\n                           strict: bool = False,\n                           use_safetensors: bool = False,\n                           load_sub_module: bool = True):\n        \"\"\"\n        Load model to checkpoint but only on master process.\n        \"\"\"\n        raise NotImplementedError(\"Sharded model checkpoint is not supported yet.\")\n\n    def save_sharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool):\n        \"\"\"\n        Save optimizer to checkpoint but only on master process.\n        \"\"\"\n        raise NotImplementedError(\"Sharded optimizer checkpoint is not supported yet.\")\n\n    def load_sharded_optimizer(self, optimizer: Optimizer, index_file_path: str, prefix: str, size_per_shard: int):\n        \"\"\"\n        Load optimizer to checkpoint but only on master process.\n        \"\"\"\n        raise NotImplementedError(\"Sharded optimizer checkpoint is not supported yet.\")\n\n    def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):\n        \"\"\"\n        Save model to checkpoint but only on master process.\n        \"\"\"\n        if self.coordinator.is_master():\n            super().save_lr_scheduler(lr_scheduler, checkpoint)\n\n\nclass TorchFSDPModel(ModelWrapper):\n\n    def __init__(self, module: nn.Module, *args, **kwargs) -> None:\n        super().__init__(module)\n        self.module = FSDP(module, *args, **kwargs)\n\n    def unwrap(self):\n        return self.module\n\n\nclass FSDPOptimizerWrapper(OptimizerWrapper):\n\n    def __init__(self, optimizer: Optimizer, model: nn.Module):\n        self.model = model\n        super().__init__(optimizer)\n\n    def unwrap_model(self) -> nn.Module:\n        return self.model\n\n\nclass TorchFSDPPlugin(DPPluginBase):\n    \"\"\"\n    Plugin for PyTorch FSDP.\n\n    Example:\n        >>> from colossalai.booster import Booster\n        >>> from colossalai.booster.plugin import TorchFSDPPlugin\n        >>>\n        >>> model, train_dataset, optimizer, criterion = ...\n        >>> plugin = TorchFSDPPlugin()\n\n        >>> train_dataloader = plugin.prepare_train_dataloader(train_dataset, batch_size=8)\n        >>> booster = Booster(plugin=plugin)\n        >>> model, optimizer, train_dataloader, criterion = booster.boost(model, optimizer, train_dataloader, criterion)\n\n    Args:\n        See https://pytorch.org/docs/stable/fsdp.html for details.\n    \"\"\"\n\n    if version.parse(torch.__version__) >= version.parse('1.12.0'):\n\n        def __init__(\n            self,\n            process_group: Optional[ProcessGroup] = None,\n            sharding_strategy: Optional[ShardingStrategy] = None,\n            cpu_offload: Optional[CPUOffload] = None,\n            auto_wrap_policy: Optional[Callable] = None,\n            backward_prefetch: Optional[BackwardPrefetch] = None,\n            mixed_precision: Optional[MixedPrecision] = None,\n            ignored_modules: Optional[Iterable[torch.nn.Module]] = None,\n            param_init_fn: Optional[Callable[[nn.Module], None]] = None,\n            sync_module_states: bool = False,\n        ):\n            super().__init__()\n            self.fsdp_kwargs = dict(process_group=process_group,\n                                    sharding_strategy=sharding_strategy,\n                                    cpu_offload=cpu_offload,\n                                    auto_wrap_policy=auto_wrap_policy,\n                                    backward_prefetch=backward_prefetch,\n                                    mixed_precision=mixed_precision,\n                                    ignored_modules=ignored_modules,\n                                    param_init_fn=param_init_fn,\n                                    sync_module_states=sync_module_states)\n    else:\n        raise RuntimeError(\"FSDP is not supported while torch version under 1.12.0.\")\n\n    def support_no_sync(self) -> bool:\n        False\n\n    def no_sync(self, model: nn.Module) -> Iterator[None]:\n        raise NotImplementedError(\"Torch fsdp no_sync func not supported yet.\")\n\n    def control_precision(self) -> bool:\n        return True\n\n    def supported_precisions(self) -> List[str]:\n        return ['fp16', 'bf16']\n\n    def control_device(self) -> bool:\n        return True\n\n    def supported_devices(self) -> List[str]:\n        return ['cuda']\n\n    def configure(\n        self,\n        model: nn.Module,\n        optimizer: Optimizer,\n        criterion: Callable = None,\n        dataloader: DataLoader = None,\n        lr_scheduler: LRScheduler = None,\n    ) -> Tuple[Union[nn.Module, OptimizerWrapper, LRScheduler, DataLoader]]:\n\n        # wrap the model with PyTorch FSDP\n        fsdp_model = TorchFSDPModel(model, device_id=torch.cuda.current_device(), **self.fsdp_kwargs)\n\n        if len(optimizer.param_groups) > 1:\n            warnings.warn(\n                'TorchFSDPPlugin does not support optimizer that use multi param groups. The results may not be as expected if used.'\n            )\n        optimizer.__init__(fsdp_model.parameters(), **optimizer.defaults)\n\n        if not isinstance(optimizer, FSDPOptimizerWrapper):\n            optimizer = FSDPOptimizerWrapper(optimizer, fsdp_model)\n\n        return fsdp_model, optimizer, criterion, dataloader, lr_scheduler\n\n    def control_checkpoint_io(self) -> bool:\n        return True\n\n    def get_checkpoint_io(self) -> CheckpointIO:\n        return TorchFSDPCheckpointIO()\n", "path": "colossalai/booster/plugin/torch_fsdp_plugin.py"}]}