|
| 1 | +# Copyright The PyTorch Lightning team. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | +import contextlib |
| 15 | +from typing import Any, Dict, Generator, List, Optional, Union |
| 16 | + |
| 17 | +import torch |
| 18 | +from torch import Tensor |
| 19 | +from torch.nn import Module |
| 20 | + |
| 21 | +from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment |
| 22 | +from pytorch_lightning.plugins.training_type.ddp import DDPPlugin |
| 23 | +from pytorch_lightning.utilities import _FAIRSCALE_FULLY_SHARDED_AVAILABLE |
| 24 | +from pytorch_lightning.utilities.exceptions import MisconfigurationException |
| 25 | + |
| 26 | +if _FAIRSCALE_FULLY_SHARDED_AVAILABLE: |
| 27 | + from fairscale.nn import default_auto_wrap_policy, enable_wrap |
| 28 | + from fairscale.nn.data_parallel import FullyShardedDataParallel |
| 29 | + |
| 30 | + |
| 31 | +class DDPFullyShardedPlugin(DDPPlugin): |
| 32 | + |
| 33 | + def __init__( |
| 34 | + self, |
| 35 | + cpu_offload: bool = False, |
| 36 | + flatten_parameters: bool = True, |
| 37 | + reshard_after_forward: bool = True, |
| 38 | + move_grads_to_cpu: Optional[bool] = None, |
| 39 | + fp32_reduce_scatter: Optional[bool] = None, |
| 40 | + compute_dtype: Optional[torch.dtype] = None, |
| 41 | + bucket_cap_mb: int = 25, |
| 42 | + min_num_params: int = 1e8, |
| 43 | + state_dict_to_cpu: bool = True, |
| 44 | + parallel_devices: Optional[List[torch.device]] = None, |
| 45 | + cluster_environment: ClusterEnvironment = None, |
| 46 | + ): |
| 47 | + """ |
| 48 | + Plugin for Fully Sharded Data Parallel provided by FairScale. |
| 49 | +
|
| 50 | + Full Sharded Training shards the entire model across all available GPUs, allowing you to scale model |
| 51 | + size, whilst using efficient communication to reduce overhead. In practice, this means we can remain |
| 52 | + at parity with PyTorch DDP, whilst scaling our model sizes dramatically. The technique is similar |
| 53 | + to ZeRO-Stage 3 but has been built for upstreaming to PyTorch. |
| 54 | + `For more information: https://fairscale.readthedocs.io/en/latest/api/nn/fsdp.html`. |
| 55 | + .. warning:: ``FullyShardedPlugin`` is in beta and subject to change. |
| 56 | +
|
| 57 | + Defaults have been set and options have been exposed, but may require configuration |
| 58 | + based on your level of memory/speed efficiency. We suggest having a look at this PR for more information. |
| 59 | + `https://github.com/facebookresearch/fairscale/pull/413` |
| 60 | +
|
| 61 | + Many of the helpful doc strings below came from the original FairScale documentation: |
| 62 | + `https://fairscale.readthedocs.io/en/latest/api/nn/fsdp.html` |
| 63 | +
|
| 64 | + Arguments: |
| 65 | + cpu_offload: Offload FP32 params to CPU. Only usable in precision=16 mode. |
| 66 | + (Default: False). |
| 67 | + move_grads_to_cpu: Moves gradient shards to CPU after reduction. |
| 68 | + Only disable if using CPU based optimizers |
| 69 | + (Default to ``cpu_offload``). |
| 70 | + flatten_parameters: Flattens parameter into single contiguous tensor for speed efficiency |
| 71 | + (Default: True). |
| 72 | + reshard_after_forward: Reshard parameters after the forward pass, which saves memory but slows |
| 73 | + down training. This is only relevant when resharding individual layers. |
| 74 | + (Default: True). |
| 75 | + fp32_reduce_scatter: Reduce-Scatter gradients in FP32. Only relevant in mixed precision |
| 76 | + (Default: None). |
| 77 | + compute_dtype: dtype for full parameters for computation. Default to torch.float32, |
| 78 | + unless using mixed precision, in which case defaults to torch.float16. |
| 79 | + (Default: None). |
| 80 | + bucket_cap_mb: bucket parameters so that gradient reduction |
| 81 | + can potentially overlap with backward computation. |
| 82 | + bucket_cap_mb controls the bucket size in MegaBytes (MB). |
| 83 | + Buckets are sub-divided based on world_size, |
| 84 | + so the max shard size is roughly bucket_cap_mb / world_size. |
| 85 | + Values <= 0 disable bucketing. |
| 86 | + (Default: 25). |
| 87 | + min_num_params: Number of parameters to wrap when using FairScale ``auto_wrap``. |
| 88 | + (Default: 1e8) |
| 89 | + state_dict_to_cpu: Whether to return parameters (returned by :func:`state_dict`) on CPU device. |
| 90 | + If ``False``, this will default to ``compute_device``. |
| 91 | + (Defautl: True). |
| 92 | + """ |
| 93 | + |
| 94 | + super().__init__( |
| 95 | + parallel_devices=parallel_devices, |
| 96 | + cluster_environment=cluster_environment, |
| 97 | + ) |
| 98 | + self.cpu_offload = cpu_offload |
| 99 | + self.move_grads_to_cpu = move_grads_to_cpu |
| 100 | + self.flatten_parameters = flatten_parameters |
| 101 | + self.reshard_after_forward = reshard_after_forward |
| 102 | + self.fp32_reduce_scatter = fp32_reduce_scatter |
| 103 | + self.compute_dtype = compute_dtype |
| 104 | + self.bucket_cap_mb = bucket_cap_mb |
| 105 | + self.min_num_params = min_num_params |
| 106 | + self.state_dict_device = torch.device("cpu") if state_dict_to_cpu else None |
| 107 | + self._process_group = None |
| 108 | + |
| 109 | + @property |
| 110 | + def process_group(self): |
| 111 | + if self._process_group is None: |
| 112 | + self._process_group = torch.distributed.new_group() |
| 113 | + return self._process_group |
| 114 | + |
| 115 | + def setup_distributed(self) -> None: |
| 116 | + if not self.on_gpu: |
| 117 | + raise MisconfigurationException( |
| 118 | + "You selected accelerator to be `ddp_fully_sharded`, but GPU is not available." |
| 119 | + ) |
| 120 | + super().setup_distributed() |
| 121 | + torch.cuda.set_device(self.root_device) |
| 122 | + |
| 123 | + @contextlib.contextmanager |
| 124 | + def model_sharded_context(self) -> Generator: |
| 125 | + precision = self.lightning_module.trainer.precision |
| 126 | + |
| 127 | + def wrap_policy(*args, **kwargs): |
| 128 | + return default_auto_wrap_policy(*args, **kwargs, min_num_params=self.min_num_params) |
| 129 | + |
| 130 | + with enable_wrap( |
| 131 | + wrapper_cls=FullyShardedDataParallel, |
| 132 | + auto_wrap_policy=wrap_policy, |
| 133 | + process_group=self.process_group, |
| 134 | + cpu_offload=self.cpu_offload, |
| 135 | + move_grads_to_cpu=self.move_grads_to_cpu, |
| 136 | + flatten_parameters=self.flatten_parameters, |
| 137 | + mixed_precision=precision == "mixed", |
| 138 | + reshard_after_forward=self.reshard_after_forward, |
| 139 | + fp32_reduce_scatter=self.fp32_reduce_scatter, |
| 140 | + compute_dtype=self.compute_dtype, |
| 141 | + bucket_cap_mb=self.bucket_cap_mb, |
| 142 | + state_dict_device=self.state_dict_device, |
| 143 | + ): |
| 144 | + yield |
| 145 | + |
| 146 | + def connect(self, model: Module) -> None: |
| 147 | + super().connect(model) |
| 148 | + model_call_configure_sharded_model_hook = getattr(model, "call_configure_sharded_model_hook", False) |
| 149 | + if not model_call_configure_sharded_model_hook: |
| 150 | + # if model has not called configure sharded model, we reset |
| 151 | + # the training type plugin's call_configure_sharded_model_hook |
| 152 | + # to give trainer a chance to configure. |
| 153 | + self.call_configure_sharded_model_hook = True |
| 154 | + |
| 155 | + def configure_ddp(self) -> None: |
| 156 | + if not self.cpu_offload: |
| 157 | + # When using CPU Offload, FSDP will manage the CUDA movement for us. |
| 158 | + # Note: this would be problematic for large model (which could not fit in one GPU) |
| 159 | + # as FSDP module.to(device) would first summon all parameters |
| 160 | + # (TODO: need to figure out solution) |
| 161 | + self.model_to_device() |
| 162 | + |
| 163 | + # setup optimizers after fully sharded has wrapped the lightning module |
| 164 | + self.lightning_module.trainer.accelerator.setup_optimizers(self.lightning_module.trainer) |
| 165 | + |
| 166 | + def pre_dispatch(self) -> None: |
| 167 | + if self.sync_batchnorm: |
| 168 | + self.model = self.configure_sync_batchnorm(self.model) |
| 169 | + self.configure_ddp() |
| 170 | + self.barrier() |
| 171 | + |
| 172 | + def model_to_device(self) -> None: |
| 173 | + # ensure we update the device type in the lightning module |
| 174 | + self.lightning_module.to(self.root_device) |
| 175 | + |
| 176 | + def lightning_module_state_dict(self) -> Dict[str, Union[Any, Tensor]]: |
| 177 | + # Currently it is same as default TrainingTypePlugin, i.e. return |
| 178 | + # the full state dict for FSDP, in the future, we will provide sharded |
| 179 | + # state dict. |
| 180 | + return super().lightning_module_state_dict() |
| 181 | + |
| 182 | + @property |
| 183 | + def setup_optimizers_in_pre_dispatch(self) -> bool: |
| 184 | + # Setup optimizers after the Fully Sharded Model has been made |
| 185 | + return True |
| 186 | + |
| 187 | + def training_step(self, *args, **kwargs): |
| 188 | + return self.model.training_step(*args, **kwargs) |
| 189 | + |
| 190 | + def validation_step(self, *args, **kwargs): |
| 191 | + return self.model.validation_step(*args, **kwargs) |
| 192 | + |
| 193 | + def test_step(self, *args, **kwargs): |
| 194 | + return self.model.test_step(*args, **kwargs) |
| 195 | + |
| 196 | + def predict_step(self, *args, **kwargs): |
| 197 | + return self.model.predict_step(*args, **kwargs) |
| 198 | + |
| 199 | + def post_training_step(self): |
| 200 | + pass |
| 201 | + |
| 202 | + @classmethod |
| 203 | + def register_plugins(cls, plugin_registry: Dict): |
| 204 | + plugin_registry.register( |
| 205 | + "fsdp", |
| 206 | + cls, |
| 207 | + description="Fully sharded training with checkpointing the full state dict.", |
| 208 | + ) |
![pre-commit-ci[bot]](https://avatars.githubusercontent.com/in/68672?v=4&size=40){kind=avatar}
0 commit comments