|
|
""" |
|
|
This file includes public APIs for FSDP such as the classes used for the |
|
|
constructor arguments. |
|
|
""" |
|
|
|
|
|
from dataclasses import dataclass |
|
|
from enum import auto, Enum |
|
|
|
|
|
from typing import Optional, Sequence, Type |
|
|
|
|
|
import torch |
|
|
from torch.nn.modules.batchnorm import _BatchNorm |
|
|
|
|
|
__all__ = [ |
|
|
"ShardingStrategy", |
|
|
"BackwardPrefetch", |
|
|
"MixedPrecision", |
|
|
"CPUOffload", |
|
|
"StateDictType", |
|
|
"StateDictConfig", |
|
|
"FullStateDictConfig", |
|
|
"LocalStateDictConfig", |
|
|
"ShardedStateDictConfig", |
|
|
"OptimStateDictConfig", |
|
|
"FullOptimStateDictConfig", |
|
|
"LocalOptimStateDictConfig", |
|
|
"ShardedOptimStateDictConfig", |
|
|
"StateDictSettings", |
|
|
] |
|
|
|
|
|
|
|
|
class ShardingStrategy(Enum): |
|
|
""" |
|
|
This specifies the sharding strategy to be used for distributed training by |
|
|
:class:`FullyShardedDataParallel`. |
|
|
|
|
|
- ``FULL_SHARD``: Parameters, gradients, and optimizer states are sharded. |
|
|
For the parameters, this strategy unshards (via all-gather) before the |
|
|
forward, reshards after the forward, unshards before the backward |
|
|
computation, and reshards after the backward computation. For gradients, |
|
|
it synchronizes and shards them (via reduce-scatter) after the backward |
|
|
computation. The sharded optimizer states are updated locally per rank. |
|
|
- ``SHARD_GRAD_OP``: Gradients and optimizer states are sharded during |
|
|
computation, and additionally, parameters are sharded outside |
|
|
computation. For the parameters, this strategy unshards before the |
|
|
forward, does not reshard them after the forward, and only reshards them |
|
|
after the backward computation. The sharded optimizer states are updated |
|
|
locally per rank. Inside ``no_sync()``, the parameters are not resharded |
|
|
after the backward computation. |
|
|
- ``NO_SHARD``: Parameters, gradients, and optimizer states are not sharded |
|
|
but instead replicated across ranks similar to PyTorch's |
|
|
:class:`DistributedDataParallel` API. For gradients, this strategy |
|
|
synchronizes them (via all-reduce) after the backward computation. The |
|
|
unsharded optimizer states are updated locally per rank. |
|
|
- ``HYBRID_SHARD``: Apply ``FULL_SHARD`` within a node, and replicate parameters across |
|
|
nodes. This results in reduced communication volume as expensive all-gathers and |
|
|
reduce-scatters are only done within a node, which can be more performant for medium |
|
|
-sized models. |
|
|
- ``_HYBRID_SHARD_ZERO2``: Apply ``SHARD_GRAD_OP`` within a node, and replicate parameters across |
|
|
nodes. This is like ``HYBRID_SHARD``, except this may provide even higher throughput |
|
|
since the unsharded parameters are not freed after the forward pass, saving the |
|
|
all-gathers in the pre-backward. |
|
|
""" |
|
|
|
|
|
FULL_SHARD = auto() |
|
|
SHARD_GRAD_OP = auto() |
|
|
NO_SHARD = auto() |
|
|
HYBRID_SHARD = auto() |
|
|
_HYBRID_SHARD_ZERO2 = auto() |
|
|
|
|
|
|
|
|
class BackwardPrefetch(Enum): |
|
|
""" |
|
|
This configures explicit backward prefetching, which improves throughput by |
|
|
enabling communication and computation overlap in the backward pass at the |
|
|
cost of slightly increased memory usage. |
|
|
|
|
|
- ``BACKWARD_PRE``: This enables the most overlap but increases memory |
|
|
usage the most. This prefetches the next set of parameters *before* the |
|
|
current set of parameters' gradient computation. This overlaps the *next |
|
|
all-gather* and the *current gradient computation*, and at the peak, it |
|
|
holds the current set of parameters, next set of parameters, and current |
|
|
set of gradients in memory. |
|
|
- ``BACKWARD_POST``: This enables less overlap but requires less memory |
|
|
usage. This prefetches the next set of parameters *after* the current |
|
|
set of parameters' gradient computation. This overlaps the *current |
|
|
reduce-scatter* and the *next gradient computation*, and it frees the |
|
|
current set of parameters before allocating memory for the next set of |
|
|
parameters, only holding the next set of parameters and current set of |
|
|
gradients in memory at the peak. |
|
|
- FSDP's ``backward_prefetch`` argument accepts ``None``, which disables |
|
|
the backward prefetching altogether. This has no overlap and does not |
|
|
increase memory usage. In general, we do not recommend this setting since |
|
|
it may degrade throughput significantly. |
|
|
|
|
|
For more technical context: For a single process group using NCCL backend, |
|
|
any collectives, even if issued from different streams, contend for the |
|
|
same per-device NCCL stream, which implies that the relative order in which |
|
|
the collectives are issued matters for overlapping. The two backward |
|
|
prefetching values correspond to different issue orders. |
|
|
""" |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
BACKWARD_PRE = auto() |
|
|
BACKWARD_POST = auto() |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class MixedPrecision: |
|
|
""" |
|
|
This configures FSDP-native mixed precision training. |
|
|
|
|
|
Attributes: |
|
|
param_dtype (Optional[torch.dtype]): This specifies the dtype for model |
|
|
parameters during forward and backward and thus the dtype for |
|
|
forward and backward computation. Outside forward and backward, the |
|
|
*sharded* parameters are kept in full precision (e.g. for the |
|
|
optimizer step), and for model checkpointing, the parameters are |
|
|
always saved in full precision. (Default: ``None``) |
|
|
reduce_dtype (Optional[torch.dtype]): This specifies the dtype for |
|
|
gradient reduction (i.e. reduce-scatter or all-reduce). If this is |
|
|
``None`` but ``param_dtype`` is not ``None``, then this takes on |
|
|
the ``param_dtype`` value, still running gradient reduction in low |
|
|
precision. This is permitted to differ from ``param_dtype``, e.g. |
|
|
to force gradient reduction to run in full precision. (Default: |
|
|
``None``) |
|
|
buffer_dtype (Optional[torch.dtype]): This specifies the dtype for |
|
|
buffers. FSDP does not shard buffers. Rather, FSDP casts them to |
|
|
``buffer_dtype`` in the first forward pass and keeps them in that |
|
|
dtype thereafter. For model checkpointing, the buffers are saved |
|
|
in full precision except for ``LOCAL_STATE_DICT``. (Default: |
|
|
``None``) |
|
|
keep_low_precision_grads (bool): If ``False``, then FSDP upcasts |
|
|
gradients to full precision after the backward pass in preparation |
|
|
for the optimizer step. If ``True``, then FSDP keeps the gradients |
|
|
in the dtype used for gradient reduction, which can save memory if |
|
|
using a custom optimizer that supports running in low precision. |
|
|
(Default: ``False``) |
|
|
cast_forward_inputs (bool): If ``True``, then this FSDP module casts |
|
|
its forward args and kwargs to ``param_dtype``. This is to ensure |
|
|
that parameter and input dtypes match for forward computation, as |
|
|
required by many ops. This may need to be set to ``True`` when only |
|
|
applying mixed precision to some but not all FSDP modules, in which |
|
|
case a mixed-precision FSDP submodule needs to recast its inputs. |
|
|
(Default: ``False``) |
|
|
cast_root_forward_inputs (bool): If ``True``, then the root FSDP module |
|
|
casts its forward args and kwargs to ``param_dtype``, overriding |
|
|
the value of ``cast_forward_inputs``. For non-root FSDP modules, |
|
|
this does not do anything. (Default: ``True``) |
|
|
_module_classes_to_ignore: (Sequence[Type[nn.Module]]): This specifies |
|
|
module classes to ignore for mixed precision when using an |
|
|
``auto_wrap_policy``: Modules of these classes will have FSDP |
|
|
applied to them separately with mixed precision disabled (meaning |
|
|
that the final FSDP construction would deviate from the specified |
|
|
policy). If ``auto_wrap_policy`` is not specified, then this does |
|
|
not do anything. This API is experimental and subject to change. |
|
|
(Default: ``(_BatchNorm,)``) |
|
|
|
|
|
.. note:: This API is experimental and subject to change. |
|
|
|
|
|
.. note:: Only floating point tensors are cast to their specified dtypes. |
|
|
|
|
|
.. note:: In ``summon_full_params``, parameters are forced to full |
|
|
precision, but buffers are not. |
|
|
|
|
|
.. note:: Layer norm and batch norm accumulate in ``float32`` even when |
|
|
their inputs are in a low precision like ``float16`` or ``bfloat16``. |
|
|
Disabling FSDP's mixed precision for those norm modules only means that |
|
|
the affine parameters are kept in ``float32``. However, this incurs |
|
|
separate all-gathers and reduce-scatters for those norm modules, which |
|
|
may be inefficient, so if the workload permits, the user should prefer |
|
|
to still apply mixed precision to those modules. |
|
|
|
|
|
.. note:: By default, if the user passes a model with any ``_BatchNorm`` |
|
|
modules and specifies an ``auto_wrap_policy``, then the batch norm |
|
|
modules will have FSDP applied to them separately with mixed precision |
|
|
disabled. See the ``_module_classes_to_ignore`` argument. |
|
|
|
|
|
.. note:: ``MixedPrecision`` has ``cast_root_forward_inputs=True`` and |
|
|
``cast_forward_inputs=False`` by default. For the root FSDP instance, |
|
|
its ``cast_root_forward_inputs`` takes precedence over its |
|
|
``cast_forward_inputs``. For non-root FSDP instances, their |
|
|
``cast_root_forward_inputs`` values are ignored. The default setting is |
|
|
sufficient for the typical case where each FSDP instance has the same |
|
|
``MixedPrecision`` configuration and only needs to cast inputs to the |
|
|
``param_dtype`` at the beginning of the model's forward pass. |
|
|
|
|
|
.. note:: For nested FSDP instances with different ``MixedPrecision`` |
|
|
configurations, we recommend setting individual ``cast_forward_inputs`` |
|
|
values to configure casting inputs or not before each instance's |
|
|
forward. In such a case, since the casts happen before each FSDP |
|
|
instance's forward, a parent FSDP instance should have its non-FSDP |
|
|
submodules run before its FSDP submodules to avoid the activation dtype |
|
|
being changed due to a different ``MixedPrecision`` configuration. |
|
|
|
|
|
Example:: |
|
|
|
|
|
>>> # xdoctest: +SKIP("undefined variables") |
|
|
>>> model = nn.Sequential(nn.Linear(3, 3), nn.Linear(3, 3)) |
|
|
>>> model[1] = FSDP( |
|
|
>>> model[1], |
|
|
>>> mixed_precision=MixedPrecision(param_dtype=torch.float16, cast_forward_inputs=True), |
|
|
>>> ) |
|
|
>>> model = FSDP( |
|
|
>>> model, |
|
|
>>> mixed_precision=MixedPrecision(param_dtype=torch.bfloat16, cast_forward_inputs=True), |
|
|
>>> ) |
|
|
|
|
|
The above shows a working example. On the other hand, if ``model[1]`` |
|
|
were replaced with ``model[0]``, meaning that the submodule using |
|
|
different ``MixedPrecision`` ran its forward first, then ``model[1]`` |
|
|
would incorrectly see ``float16`` activations instead of ``bfloat16`` |
|
|
ones. |
|
|
|
|
|
""" |
|
|
|
|
|
param_dtype: Optional[torch.dtype] = None |
|
|
reduce_dtype: Optional[torch.dtype] = None |
|
|
buffer_dtype: Optional[torch.dtype] = None |
|
|
keep_low_precision_grads: bool = False |
|
|
cast_forward_inputs: bool = False |
|
|
cast_root_forward_inputs: bool = True |
|
|
_module_classes_to_ignore: Sequence[Type[torch.nn.Module]] = (_BatchNorm,) |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class CPUOffload: |
|
|
""" |
|
|
This configures CPU offloading. |
|
|
|
|
|
Attributes: |
|
|
offload_params (bool): This specifies whether to offload parameters to |
|
|
CPU when not involved in computation. If ``True``, then this |
|
|
offloads gradients to CPU as well, meaning that the optimizer step |
|
|
runs on CPU. |
|
|
""" |
|
|
|
|
|
offload_params: bool = False |
|
|
|
|
|
|
|
|
class StateDictType(Enum): |
|
|
""" |
|
|
This enum indicates that which type of ``state_dict`` the FSDP module is |
|
|
currently processing (returning or loading). |
|
|
The default value is FULL_STATE_DICT to comply the PyTorch convention. |
|
|
..note:: |
|
|
FSDP currently supports three types of ``state_dict``: |
|
|
1. ``state_dict/load_state_dict`: this pair of APIs return and load |
|
|
the non-sharded, unflattened parameters. The semantics is the |
|
|
same as using DDP. |
|
|
2. ``_local_state_dict/_load_local_state_dict``: this pair of APIs return |
|
|
and load local sharded, flattened parameters. The values returned |
|
|
by ``_local_state_dict`` can be directly used by FSDP and is only |
|
|
meaningful to FSDP (because parameters are flattened). Note that |
|
|
these APIs are meant for use via the :func:`state_dict_type` |
|
|
context manager as follows: |
|
|
>>> # xdoctest: +SKIP("undefined variables") |
|
|
>>> with fsdp.state_dict_type(StateDictType.LOCAL_STATE_DICT): |
|
|
... state = fsdp.state_dict() # loads local state dict |
|
|
3. ``_sharded_state_dict/_load_sharded_state_dict``: this pair of APIs |
|
|
return and load sharded, unflattened parameters. The ``state_dict`` |
|
|
return by ``sharded_state_dict`` can be used by all other parallel |
|
|
schemes (resharding may be required). |
|
|
""" |
|
|
|
|
|
FULL_STATE_DICT = auto() |
|
|
LOCAL_STATE_DICT = auto() |
|
|
SHARDED_STATE_DICT = auto() |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class StateDictConfig: |
|
|
""" |
|
|
``StateDictConfig`` is the base class for all ``state_dict`` configuration |
|
|
classes. Users should instantiate a child class (e.g. |
|
|
``FullStateDictConfig``) in order to configure settings for the |
|
|
corresponding ``state_dict`` type supported by FSDP. |
|
|
|
|
|
Attributes: |
|
|
offload_to_cpu (bool): If ``True``, then FSDP offloads the state dict |
|
|
values to CPU, and if ``False``, then FSDP keeps them on GPU. |
|
|
(Default: ``False``) |
|
|
""" |
|
|
|
|
|
offload_to_cpu: bool = False |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class FullStateDictConfig(StateDictConfig): |
|
|
""" |
|
|
``FullStateDictConfig`` is a config class meant to be used with |
|
|
``StateDictType.FULL_STATE_DICT``. We recommend enabling both |
|
|
``offload_to_cpu=True`` and ``rank0_only=True`` when saving full state |
|
|
dicts to save GPU memory and CPU memory, respectively. This config class |
|
|
is meant to be used via the :func:`state_dict_type` context manager as |
|
|
follows: |
|
|
|
|
|
>>> # xdoctest: +SKIP("undefined variables") |
|
|
>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP |
|
|
>>> fsdp = FSDP(model, auto_wrap_policy=...) |
|
|
>>> cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) |
|
|
>>> with FSDP.state_dict_type(fsdp, StateDictType.FULL_STATE_DICT, cfg): |
|
|
>>> state = fsdp.state_dict() |
|
|
>>> # `state` will be empty on non rank 0 and contain CPU tensors on rank 0. |
|
|
>>> # To reload checkpoint for inference, finetuning, transfer learning, etc: |
|
|
>>> model = model_fn() # Initialize model in preparation for wrapping with FSDP |
|
|
>>> if dist.get_rank() == 0: |
|
|
>>> # Load checkpoint only on rank 0 to avoid memory redundancy |
|
|
>>> state_dict = torch.load("my_checkpoint.pt") |
|
|
>>> model.load_state_dict(state_dict) |
|
|
>>> # All ranks initialize FSDP module as usual. `sync_module_states` argument |
|
|
>>> # communicates loaded checkpoint states from rank 0 to rest of the world. |
|
|
>>> fsdp = FSDP(model, device_id=torch.cuda.current_device(), auto_wrap_policy=..., sync_module_states=True) |
|
|
>>> # After this point, all ranks have FSDP model with loaded checkpoint. |
|
|
|
|
|
Attributes: |
|
|
rank0_only (bool): If ``True``, then only rank 0 saves the full state |
|
|
dict, and nonzero ranks save an empty dict. If ``False``, then all |
|
|
ranks save the full state dict. (Default: ``False``) |
|
|
""" |
|
|
|
|
|
rank0_only: bool = False |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class LocalStateDictConfig(StateDictConfig): |
|
|
pass |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class ShardedStateDictConfig(StateDictConfig): |
|
|
""" |
|
|
``ShardedStateDictConfig`` is a config class meant to be used with |
|
|
``StateDictType.SHARDED_STATE_DICT``. |
|
|
|
|
|
Attributes: |
|
|
_use_dtensor (bool): If ``True``, then FSDP saves the state dict values |
|
|
as ``DTensor``, and if ``False``, then FSDP saves them as |
|
|
``ShardedTensor``. (Default: ``False``) |
|
|
|
|
|
.. warning:: ``_use_dtensor`` is a private field of :class:`ShardedStateDictConfig` |
|
|
and it is used by FSDP to determine the type of state dict values. Users should not |
|
|
manually modify ``_use_dtensor``. |
|
|
""" |
|
|
|
|
|
_use_dtensor: bool = False |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class OptimStateDictConfig: |
|
|
""" |
|
|
``OptimStateDictConfig`` is the base class for all ``optim_state_dict`` |
|
|
configuration classes. Users should instantiate a child class (e.g. |
|
|
``FullOptimStateDictConfig``) in order to configure settings for the |
|
|
corresponding ``optim_state_dict`` type supported by FSDP. |
|
|
|
|
|
Attributes: |
|
|
offload_to_cpu (bool): If ``True``, then FSDP offloads the state dict's |
|
|
tensor values to CPU, and if ``False``, then FSDP keeps them on the |
|
|
original device (which is GPU unless parameter CPU offloading is |
|
|
enabled). (Default: ``True``) |
|
|
""" |
|
|
|
|
|
offload_to_cpu: bool = True |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class FullOptimStateDictConfig(OptimStateDictConfig): |
|
|
""" |
|
|
Attributes: |
|
|
rank0_only (bool): If ``True``, then only rank 0 saves the full state |
|
|
dict, and nonzero ranks save an empty dict. If ``False``, then all |
|
|
ranks save the full state dict. (Default: ``False``) |
|
|
""" |
|
|
|
|
|
rank0_only: bool = False |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class LocalOptimStateDictConfig(OptimStateDictConfig): |
|
|
offload_to_cpu: bool = False |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class ShardedOptimStateDictConfig(OptimStateDictConfig): |
|
|
""" |
|
|
``ShardedOptimStateDictConfig`` is a config class meant to be used with |
|
|
``StateDictType.SHARDED_STATE_DICT``. |
|
|
|
|
|
Attributes: |
|
|
_use_dtensor (bool): If ``True``, then FSDP saves the state dict values |
|
|
as ``DTensor``, and if ``False``, then FSDP saves them as |
|
|
``ShardedTensor``. (Default: ``False``) |
|
|
|
|
|
.. warning:: ``_use_dtensor`` is a private field of :class:`ShardedOptimStateDictConfig` |
|
|
and it is used by FSDP to determine the type of state dict values. Users should not |
|
|
manually modify ``_use_dtensor``. |
|
|
""" |
|
|
|
|
|
_use_dtensor: bool = False |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class StateDictSettings: |
|
|
state_dict_type: StateDictType |
|
|
state_dict_config: StateDictConfig |
|
|
optim_state_dict_config: OptimStateDictConfig |
|
|
|
