Repository: pytorch
Issue: [RFC] Per-Parameter-Sharding FSDP
Labels: triaged,module: fsdp

# Per-Parameter-Sharding FSDP

## Motivation
As we looked toward next-generation training, we found limitations in our existing FSDP, mainly from the _flat parameter_ construct. To address these, we propose a simpler FSDP design that shards each parameter on dim-0 (hence, "per-parameter sharding") that enables:
1. Flexible fp8 all-gather: fp8 weights and other non-fp8 parameters can be flexibly mixed in the same all-gather
2. Flexible frozen parameters: frozen and non-frozen parameters can be flexibly mixed in the same communication group without using extra memory
3. Communication-free sharded state dicts: matching the training and state dict representation simplifies and speeds up checkpointing
4. Future communication optimization in the compiler: a partial graph compiler like `torch.compile` can change the communication groups for all-gather/reduce-scatter

Overall, we strive for simplicity and composability in the new design. 

We have validated the prototype on Llama-like models, achieving on-par throughput while using less memory. The only performance downside to this approach are extra copies before/after collectives; however, we plan to mitigate this with custom copy kernels/fast-paths.

We welcome any and all feedback on this effort!

## API
The API is _not_ backward compatible with the existing FSDP (though model/optimizer checkpoints should be). We use this rewrite as a chance to simplify the constructor and expose new capabilities. A tentative API looks like:
```
@contract(state_cls=FSDPState)
def fully_shard(
  module: nn.Module,
  *,
  mesh: Optional[DeviceMesh] = None,
  reshard_after_forward: Union[bool, int] = True,
  init_policy: InitPolicy = InitPolicy(),
  mp_policy: MixedPrecisionPolicy = MixedPrecisionPolicy(),
  offload_policy: OffloadPolicy = OffloadPolicy(),
  comm_policy: PrefetchPolicy = CommPolicy(),
) -> nn.Module:  # return value only used by `contract` for checks
```

```
@dataclassf
class InitPolicy:
  module_init_fn: Optional[Callable[[nn.Module], None]] = None
  sync_module_states: bool = False
@dataclass
class MixedPrecisionPolicy:
  param_dtype: Optional[torch.dtype] = None
  reduce_dtype: Optional[torch.dtype] = None
  output_dtype: Optional[torch.dtype] = None
  cast_forward_inputs: bool = True
@dataclass
class OffloadPolicy:
  offload_type: Optional[str] = None  # "cpu"
@dataclass
class CommPolicy:
  forward_prefetch_limit: int = 1
  backward_prefetch_limit: int = 1
```
Naming and organization are open to feedback!

- [`@contract`](https://github.com/pytorch/pytorch/blob/b5dd37f23efecdd27d1ff8862dcd59613c36211d/torch/distributed/_composable/contract.py#L26) checks that the API (`fully_shard`) [preserves `named_parameters()`/`named_buffers()`/`named_modules()`](https://github.com/pytorch/pytorch/blob/b5dd37f23efecdd27d1ff8862dcd59613c36211d/torch/distributed/_composable/contract.py#L149-L163) and [embeds a state and registry](https://github.com/pytorch/pytorch/blob/b5dd37f23efecdd27d1ff8862dcd59613c36211d/torch/distributed/_composable/contract.py#L94-L97) on the module. This is a broader composability effort with other distributed training APIs (e.g. DDP, AC, and TorchRec).
- `module` gives the module to shard with FSDP. The call to `fully_shard` constructs _one_ communication group, where the group's parameters include those in `module.parameters()` except those already assigned to another API (whether `fully_shard` or other). The group's parameters and gradients are all-gathered and reduce-scattered together, respectively.
- `mesh` gives the `DeviceMesh` over which FSDP shards. It can be a 1D mesh for normal FSDP or a 2D mesh for HSDP, where sharding is on the inner dimension. If `mesh is None`, then FSDP constructs a global 1D mesh by default. The mesh's `device_type` gives the device used for initialization and training.
- `reshard_after_forward` controls the parameter behavior after forward and can trade off memory and communication.
    - If `True`, then this reshards parameters after forward and all-gathers in backward.
    - If `False`, then this keeps the unsharded parameters in memory after forward and avoids the all-gather in backward.
    - If an `int`, then this represents the world size to reshard to after forward. It should be in `(1, shard world size)`, e.g. the intra-node size.
    - With respect to [DeepSpeed ZeRO](https://www.deepspeed.ai/tutorials/zero/), `True` is analogous to ZeRO-3; `False` is analogous to ZeRO-2; and an `int` is analogous to [ZeRO++](https://www.deepspeed.ai/tutorials/zeropp/) hpZ. 
- `init_policy` gives the `module_init_fn` and `sync_module_states` fields for initializing a meta-device module and for broadcasting managed module states from rank 0, respectively.
- `mp_policy` controls various mixed precision options.
    - `param_dtype` specifies the unsharded parameter dtype, which controls the all-gather dtype and the dtype for forward/backward computation.
    - `reduce_dtype` specifies the reduce-scatter dtype. If `None`, this inherits from `param_dtype`. A common example is `param_dtype=torch.bfloat16` and `reduce_dtype=torch.float32`.
    - `output_dtype` specifies the dtype to cast floating-point forward outputs. This is useful for cases where different submodules have different mixed precision policies.
    - `cast_forward_inputs` specifies whether to cast forward floating-point input tensors to `param_dtype`. This is `True` by default to ensure correctness for typical use cases. However, this can be set to `False` if the user wants to handle cast (e.g. if one input tensor should _not_ be cast).
- `offload_policy` gives whether to enable CPU offloading. This may be extended in the future for [NVMe offloading](https://github.com/pytorch/pytorch/issues/76924), but that is not planned as of today.
- `prefetch_policy` gives option to prefetch all-gathers more aggressively in forward and backward. This an advanced option and most likely should be justified with a profiler trace.

Functionally, applying `fully_shard` to a module:
-  Replaces its parameters with `DTensor`s with `(Shard(0),)` placement (meaning that they are sharded over a 1D mesh on dim-0)
- Registers forward hooks to (1) all-gather/free parameters and (2) register backward hooks that similarly all-gather/free parameters and also reduce-scatter gradients
- Overrides some existing `nn.Module` methods and exposes some new methods on the module

To achieve communication and computation overlap, the user must apply `fully_shard` to different modules (hence constructing multiple communication groups) plus the root module. For transformer architectures, this conventionally means applying `fully_shard` to each transformer block and then to the overall model.
```
for module in model.modules();
  if isinstance(module, TransformerBlock):
    fully_shard(module, **fsdp_kwargs)
fully_shard(model, **fsdp_kwargs)
```

## System Design
The design is mostly inspired from existing FSDP with some new considerations for composability with other distributed training APIs, where preserving fully-qualified names (FQNs) is greatly preferred.

### Overview
- `fully_shard`: frontend API; applied to a module
- `FSDPModule`: lightweight stateless class for overriding existing `nn.Module` methods and adding new methods; used to dynamically construct the subclass (e.g. `FSDPLinear` from `Linear` and `FSDPModule`)
- `FSDPState`: class for FSDP-related state corresponding to the module; embedded on the module using `@contract`
- `FSDPParamGroup`: class for one parameter communication group; at most 1 per `FSDPState` (0 if no managed parameters)
- `FSDPParam`: class for one parameter in the communication group

![sys_overview](https://github.com/pytorch/pytorch/assets/31054793/104e2df0-6856-4bae-97d8-ba0c9c5e7242)


Compared with existing FSDP:
- `fully_shard` + `FSDPState` together are analogous to `FullyShardedDataParallel`.
- `FSDPParamGroup` is analogous to `FlatParamHandle`, and `FSDPParam` is _roughly_ analogous to `FlatParameter`.

We can generally partition the FSDP logic into three components: (1) initialization, (2) runtime (i.e. forward/backward), and (3) checkpointing. The runtime has the greatest complexity.

### FSDP Initialization
The initialization generally logic consists of:
1. Validating and normalizing user arguments
2. Determining the managed modules and their parameters/buffers
3. Materializing meta-device modules if needed
5. Moving states to `device`
6. Sharding parameters

(1) to (4) live in `fully_shard.py` and `_fsdp_init.py`. (5) goes through `_fsdp_state.py` to `_fsdp_param_group.py` to `_fsdp_param.py`

We compute managed modules with a graph traversal (e.g. DFS) from the passed-in module, stopping at and excluding any module with a non-composable API (e.g. `replicate` for DDP) or `fully_shard` already applied.

### FSDP Runtime
The runtime goals are:
1. All-gather parameters before they are used and free them after they are used (in both forward/backward)
2. Reduce-scatter gradients after they are computed (in backward)

In forward, we define the parameters' _live_ interval by `nn.Module.forward()`. In backward, we define the live interval starting from when the gradients wrt. the module outputs have been computed to when the gradients wrt. the module inputs and wrt. the parameters have been computed. These definitions map to points in the execution accessible by `nn.Module` and autograd hooks.

The forward definition can be unsound when parameters are used outside their owning module's `forward()`. For example, this can happen when users define a custom method that runs forward logic like `generate()` or when a parent module uses a child module's parameter without calling `forward()`. We discuss handling these cases later.

#### `FSDPParamGroup` & `FSDPState`
Each `fully_shard` call embeds an `FSDPState` object on the module. This `FSDPState` can have one `FSDPParamGroup` if it manages nonzero parameters. Then, this `FSDPParamGroup` manages a group of parameters, represented as `FSDPParam`s. 

Comparing `FSDPState` and `FSDPParamGroup`:
- `FSDPParamGroup` is responsible for logic that only happens if there are managed parameters.
- `FSDPState` is responsible for logic that happens regardless of if there are managed parameters.

Each `FSDPParamGroup` defines two base operations around parameters:
- Unshard: allocates and all-gathers parameters
- Reshard: frees all-gathered parameters


#### FSDP Runtime Schedule
The runtime schedule looks roughly like:
- Root pre-forward (`register_forward_pre_hook()` in `_fsdp_state.py`): lazy initialization if needed; stream synchronization; move inputs to GPU if needed
- For each state/parameter group:
	- State pre-forward (`register_forward_pre_hook()` in `_fsdp_state.py`): cast forward inputs if needed
	- Parameter group pre-forward (in `_fsdp_param_group.py`): unshard parameters
	- Parameter group post-forward (in `_fsdp_param_group.py`): reshard parameters
	- State post-forward (`register_forward_hook()` in `_fsdp_state.py`): cast forward outputs if needed
- Root post-forward (`register_forward_hook()` in `_fsdp_state.py`)
- Root pre-backward (`register_hook()` on module output tensors in `_fsdp_state.py`): queue callback
- For each parameter group:
	- Pre-backward (`register_hook()` on module output tensors in `_fsdp_param_group.py`): unshard parameters; prefetch parameters if needed
	- Post-backward (`autograd.Function.backward` in `_fsdp_param_group.py`): reshard parameters; reduce-scatter gradients
- Root post-backward final callback (`_execution_engine.queue_callback()` in `_fsdp_state.py`): finalize parameter group backwards; stream/event synchronization; reset data structures

In practice, the nesting of parameter groups' modules can lead to some of the (pre-forward, post-forward) and (pre-backward, post-backward) intervals to be nested. Some care must be taken to ensure each pre-backward hook only runs once per (pre-backward, post-backward) interval; there may be multiple such intervals in the case a module participates multiple times per forward.

For the post-backward hook, the custom autograd function forward is an identity function on the input tensors that require gradients, and the backward is an identity function that additionally calls the post-backward hook.

#### `FSDPParam`
To understand how the unshard and reshard operations are implemented, we examine the `FSDPParam` class.

As a general framework, each `FSDPParam` considers the following tensors:
- Original parameter:
    - This is the parameter on the module when passed to FSDP.
- Sharded parameter:
    - This derives by sharding the original parameter on dim-0 via `DTensor`'s `Shard(0)` placement.
    - This is always in GPU memory (unless CPU offloading is enabled).
    - If the original parameter is a `torch.Tensor`, then this is a 1D `DTensor` with `(Shard(0),)` placements.
    - If the original parameter is a `DTensor`, then this is a 2D `DTensor` with an additional `Shard(0)` placement along the FSDP mesh dim.
    - This implies that the optimizer step runs through `DTensor`'s operator dispatch, making reducing CPU overhead critical. This can be achieved mostly via foreach or fused optimizer implementations, which amortize the overheads over tensor lists.
- All-gather input:
    - This is the `torch.Tensor` passed to the all-gather as input.
    - For the default `torch.Tensor` original parameter case, this has the same data as the sharded parameter.
- All-gather output:
    - This is the `torch.Tensor` resulting from all-gathering the all-gather input.
    - For the default `torch.Tensor` original parameter case, this has the same data as the unsharded parameter.
- Unsharded parameter:
    - This is the parameter used for forward/backward computation derived from the all-gather output.
    - This is the leaf for autograd.

We distinguish between sharded parameter vs. all-gather input and between unsharded parameter vs. all-gather output to allow for pre-all-gather and post-all-gather transformations:
- All-gather input = pre-all-gather transform(sharded parameter)
- Unsharded parameter = post-all-gather transform(all-gather output)

For example, for a `Float8Linear` weight, the pre-all-gather transform can be to cast the fp32 weight to fp8 and return its `torch.Tensor` data.

When `reshard_after_forward` is an `int`, there is a another tensor for the temporary data constructed from the `unsharded_param` by chunking to the post-forward world size. Overall, this means that each `FSDPParam` can be in one of three states: `SHARDED`, `SHARDED_POST_FORWARD`, or `UNSHARDED`. The only disallowed state transitions are between `SHARDED` and `SHARDED_POST_FORWARD`.

#### FSDP & Autograd
To implement freeing the unsharded parameter after forward and restoring it before backward, FSDP uses `untyped_storage().resize_(0)` and `resize_(orig_storage_size)`. This is a _hacky_ trick to make autograd work, even in the presence of aliases. Autograd packs a reference to `unsharded_param` (or possibly a view of it) in forward; FSDP frees the storage unbeknownst to autograd on the promise that it will restore it before the gradient computation in backward; and FSDP does so in the pre-backward hook.

Replacing this with saved tensor hooks is tricky due to the problem of how to reconstruct a packed tensor from its base tensor. This is especially tricky when the unsharded parameter is a  tensor subclass like `DTensor`.

#### FSDP Collectives
FSDP relies on batched all-gather and reduce-scatter for communication efficiency. (HSDP additionally relies on all-reduce, but we omit that for brevity.) When using per-parameter sharding, the options include (1) copy-in/copy-out and (2) NCCL group coalescing. Our internal benchmarks show that the former is achieves higher bandwidth utilization than the latter (and aside from the copies, is the same as the existing FSDP collectives, making them more trustable), so we choose that approach. (Adding group coalescing as an option may be future work.)

For all-gather:
- Allocate a contiguous tensor for the all-gather output
- Set the all-gather input as a view into the all-gather output
- Copy each `sharded_param` into the rank's all-gather input at the appropriate offset
	- This gives an interleaved layout like `(p1 shard1, p2 shard1, p3 shard1, p1 shard2, p2 shard2, p3 shard2)` for parameters `p1`, `p2`, `p3` and 2 ranks
- All-gather
- Copy each parameter's shards out to make them contiguous
	- This is currently implemented as viewing to get `(p1 shard1, p1 shard2)`, ..., `(p3 shard1, p3 shard2)` and calling `torch.cat` on the shards in a for loop

Note: For this copy-in/out approach, the unsharded parameter does not need to be padded, but for NCCL group coalescing it does.

For reduce-scatter:
- Allocate a contiguous tensor for the reduce-scatter input
- Copy each `unsharded_param.grad` into the reduce-scatter input
    - This similarly gives an interleaved layout like `(g1 shard1, g2 shard1, g3 shard1, g1 shard2, g2 shard2, g3 shard2)` for parameters `g1`, `g2`, `g3` and 2 ranks
    - This is currently implemented using `torch._foreach_copy_`
- Allocate a separate tensor for reduce-scatter output
- Reduce-scatter
- View each new sharded parameter gradient from the reduce-scatter output and accumulate with `sharded_param.grad` as needed

We find that for loop + `torch.cat` achieves better memory bandwidth than `torch._foreach_copy_` for Llama models (however more extensive microbenchmarking may be valuable). We use `_foreach_copy_` for reduce-scatter (1) for simplicity when handling the possible dtype cast (e.g. when the source is bf16 but destination is fp32) and (2) because autograd does not compute the gradients with padding (and we do not know how to make it do so). For (2), we want to avoid worst-case doubling the number of tensors to concatenate due to the padding per gradient.

To implement "fp8 all-gather", we check if the parameter group has any fp8 weights. If so, then we use some additional logic to view tensors using a 1-byte dtype before all-gathering and restoring the original dtype after.

#### FSDP, Streams, & Avoiding `recordStream`
Per-parameter-sharding FSDP uses a few different streams for overlapping:
- All-gather copy-in stream: overlap copy-in with forward compute and reduce-scatter
- All-gather stream: overlap all-gather with forward/backward computation
- Reduce-scatter stream: overlap with backward computation; convenient to queue additional work like division
- (For HSDP) all-reduce stream: overlap with all-gather/reduce-scatter/backward computation ([details](https://github.com/pytorch/pytorch/pull/106080))

When a tensor is allocated in one stream and used in another, PyTorch requires extra synchronization to ensure correctness since streams do not synchronize with each other by default. One option is to use `recordStream`, which offloads the synchronization burden to the CUDA caching allocator but means that memory is not freed until _GPU_ uses finish (which could be much later than the CPU free). Another option is to manually manage the stream dependencies using explicit synchronization. Refer to [@janeyx99 's note](https://dev-discuss.pytorch.org/t/fsdp-cudacachingallocator-an-outsider-newb-perspective/1486) for more details.

Fairscale FSDP, PyTorch FSDP, and DeepSpeed ZeRO use `recordStream`. Per-parameter-sharding FSDP uses manual synchronization instead since doing so can deterministically achieve the desired memory usage (often saving memory in real workloads) without any blocking of the CPU thread.

For all-gather, we allocate the all-gather output tensor in the all-gather stream, but we copy out its contents in the default stream so that the unsharded parameters are allocated in the default stream. This ensures that autograd will not insert any of its own `recordStream` calls in backward. For reduce-scatter, we allocate the reduce-scatter input and output in the reduce-scatter stream, so since we view into the output, the sharded gradients are allocated in the reduce-scatter stream.

For all-gather, the synchronization is as follows:
- We record an event after the all-gather copy-in for the all-gather stream to wait on.
- We record an event after the all-gather for the default stream to wait on.
- We record an event after the all-gather copy-out for the all-gather copy-in/all-gather streams to wait on.

For reduce-scatter, the synchronization is as follows:
- We record an event after the reduce-scatter copy-in for the default stream to wait on before freeing the autograd-computed gradients.
- We do not allocate the reduce-scatter input in the default stream because then we must have the default stream wait for the reduce-scatter (or else after, the reduce-scatter input will be freed with its memory incorrectly reused in the default stream).
- We do not allocate the reduce-scatter output in the default stream even though its lifetime should contain the reduce-scatter itself for simplicity. I.e., this one could go either way.
- There does not need to be any synchronization between the reduce-scatter stream and the all-gather stream since they do not share tensors. The all-gather stream operates on the all-gather output (and input, which is a view into it), and the reduce-scatter stream operates on the reduce-scatter input (with gradients copied in) and the reduce-scatter output.

Other synchronization:
- When CPU offloading is enabled, we record an event after the non-blocking gradient D2H offload for the CPU to block on at the end of backward (to ensure that the copy finishes).
- We record an event after resharding after forward to a different world size for the all-gather copy-in/all-gather streams to wait on since that data is allocated in the default stream and used in the all-gather copy-in stream.

### FSDP Checkpointing
Existing FSDP supports three kinds of state dicts: full, sharded, and local. We define a "clean" fully-qualified name (FQN) to be one without any `nn.Module` wrapper prefixes (e.g. `_fsdp_wrapped_module.` for FSDP or `module.` for DDP).
- A full state dict maps clean FQN to unsharded `torch.Tensor`. Since this may use too much GPU memory, it offers `rank0_only: bool` and `offload_to_cpu: bool` options. This requires all-gathering parameters.
- A sharded state dict maps clean FQN to sharded `DTensor`. For existing FSDP, this requires all-gathering parameters and resharding to per-parameter sharding.
- A local state dict maps clean FQN to the sharded `FlatParameter` as a `ShardedTensor`. This does not require communication but cannot be loaded into a different distributed training setup.

For per-parameter-sharding FSDP, we only support sharded state dict.
- Getting a full state dict can be implemented as a post-processing step to saving the sharded state dict.
- Sharded and local state dicts are equivalent since the training and checkpointing representation match.

## FSDP Dynamic Subclass
To avoid modifying fully-qualified names but still be able to override and add new methods, we dynamically construct a subclass and change the passed-in module's class:
```
@contract(state_cls=FSDPState)  
def fully_shard(module: nn.Module, ...) -> nn.Module:  
  ...
  cls = module.__class__  
  new_cls =  type(f"FSDP{cls.__name__}", (FSDPModule, cls), {})  
  module.__class__ = new_cls  
  return module  
  
class FSDPModule:  
  # Override `nn.Module` methods as desired  
  # Add new methods as desired
```

## fp8
See https://github.com/pytorch/pytorch/pull/114733#issuecomment-1832374773 for some description of (1) the plan for fp8 and FSDP and (2) what we require from the fp8 library (e.g. float8-experimental or TransformerEngine).


## Discussion

**What about auto wrapping?**
For context, auto wrapping refers to "automatically" traversing the module tree and applying FSDP via the `auto_wrap_policy` argument. It _does not_ refer to automatically finding the right modules to apply FSDP to. For example, for `fully_shard`, the existing `ModuleWrapPolicy` is effectively equivalent to the following code:
```
for module in model.modules():
  if isinstance(module, module_classes):
    fully_shard(module, **kwargs)
fully_shard(model, **kwargs)
```
More precisely, auto wrapping follows a post-order traversal to determine which modules to apply the API to and with which kwargs, performs some checks, and then actually makes the function calls. This post-order traversal applies to other APIs as well like activation checkpointing, for which there exists an `apply_activation_checkpointing()` API that calls into original FSDP `_recursive_wrap()` subroutine for the traversal.

One option then is to simply expose an API like:
```
def post_order_apply(
  fn: Callable, module: nn.Module, policy, ignored_modules: Iterable[nn.Module], **kwargs
):
```
This is similar to `nn.Module.apply()` with some more functionality. `fully_shard` will need to be special-cased to always apply to the root module (e.g. as indicated by an attribute on the function itself like `fully_shard._force_root = True`).

The majority of open source users use auto wrapping, whereas internally only a few do. If people feel strongly, we can once again include an argument like `apply_policy` to `fully_shard` to mimic `auto_wrap_policy`. (Since `fully_shard` is not a "wrapper" in the existing sense, so "wrap" may be a misnomer.) In that case, activation checkpointing will have a standalone function for the post-order apply as well.


**Should we support `buffer_dtype`?**
Existing FSDP supports a `buffer_dtype` argument for enabling buffer mixed precision. However, the current semantics are not exactly aligned with the semantics for parameters. For buffers, existing FSDP does not maintain both the low and high precision copies. Instead, when saving a state dict, it in-place casts the buffers back to the high precision (losing the numeric precision).

For now, we defer implementing `buffer_dtype` for simplicity. If there is a future motivating use case, then we may consider maintaining both low and high precision copies (at the cost of higher memory usage) to align the semantics with parameters.

**How does FSDP work with shared parameters?**
For any API like `fully_shard` that (1) swaps existing parameters with new ones and (2) can be applied bottom-up, there is no way to preserve sharedness after the parameter swap. We need something like `torch.swap_tensor()` that preserves the references even after swapping (to-do: check if it supports swapping a `torch.Tensor` with a subclass like `DTensor`).

Thus, for now, we impose the same requirement as existing FSDP: shared parameters must be assigned to the same `fully_shard` instance. This requires some care with how `fully_shard` is applied.

**What about existing FSDP?**
One option is to add another argument to `FullyShardedDataParallel` like `use_per_param: bool = False` that if set to `True` calls into `fully_shard`. We can remap arguments or warn/error on invalid ones. This might help some users migrate. However, if users want the full power of `fully_shard` (e.g. `reshard_after_forward` as an `int`), then they should use `fully_shard` directly. We are preferring to not do this.

## Progress Tracker
The main code has landed in PyTorch. See https://github.com/pytorch/pytorch/issues/120003 for the remaining gaps, which may take more time to fill.

Some notes:
- The current non-test code sits at ~3k lines, whereas the existing FSDP non-test code sits at ~14k lines.
- We integrate fp8 all-gather first with https://github.com/pytorch-labs/float8_experimental.


cc @zhaojuanmao @mrshenli @rohan-varma @fegin @penguinwu @albanD @wanchaol @wconstab