pytorch

afcc7c7f - [FSDP] Generalize prefetching; lower unshard/reshard to handle (#83665)

[FSDP] Generalize prefetching; lower unshard/reshard to handle (#83665) ### Additional Constructor Changes - `self.sharding_strategy` - If the world size is 1, I clamp the sharding strategy to `NO_SHARD`, regardless of the passed-in sharding strategy, since the behavior is fully equivalent. This absolves the need for `p._is_sharded or self.world_size == 1` checks in the core code. Once we fully shift the paradigm to using handles, this should result in a clear net positive. However, for now, we still have some places where we interface directly with the `FlatParameter`, in which case we have some temporary hacky code. - `HandleConfig` - As a part of the new design abstraction, much logic is lowered to the `FlatParamHandle`. This requires the handle be aware of mixed precision, CPU offloading, sharding strategy, and the process group (for world size > 1). To be less error-prone, I re-defined the `dataclass`s and `enum`s for the handle. These can be removed and coalesced with the existing ones. - The drawback is that the `FlattenParamsWrapper` constructor now takes in the `HandleConfig` to forward it to the `FlatParamHandle` constructor. I tolerate this since we plan to retire the FPW. For now, the handle's process group attributes are set later when we call `handle.shard()`. - We will dive into this logic lowering later. For now, the idea is we need to pass some extra info to the handle, which must go through the FPW. - `FullyShardedDataParallel._shard_parameters()` -> `FlatParamHandle.shard()` - [Important] Generalizing attributes to remove the 1 `FullyShardedDataParallel` : 1 `FlatParameter` assumption - **Before:** `_fsdp_graph_order`, `_pre_backward_hook_full_params_prefetched`, `_forward_full_params_prefetched`, `reshard_after_forward` are with respect to 1 `FullyShardedDataParallel` - **After:** (1) We use `FlatParamHandle` in place of `FullyShardedDataParallel`. (2) The atomic unit for forward and pre-backward is a _group_ of handles involved in the same module's forward/pre-backward. This is represented as `Tuple[FlatParamHandle, ...]`. For now, this is **always a singleton tuple**, but this shift enables a module having multiple FSDP parameters (which we have use cases for). - `_reset_lazy_init()` attributes - The prefetched flags are merged into `self._handles_prefetched`, which is directly defined in the constructor. `reshard_after_forward` is retired since it can be fully determined by other attributes (`_is_root` and `sharding_strategy`). ## FSDP Runtime: Unshard The first step is to read the existing `_rebuild_full_params()`. A few notable observations: - It returns `Tuple[Tensor, bool]`. The first element is the _padded unsharded flattened parameter_, and the second element is whether we can free it upon exiting `summon_full_params()`. This return value is **only used in `summon_full_params()`**. - If parameter mixed precision is enabled and the `FlatParameter` is already unsharded, then the low precision shard (`_mp_shard`) is still re-allocated on GPU. (It is freed at the end of the method.) - If CPU offloading is enabled and the `FlatParameter` is already unsharded, then there is a no-op `p.data = p.data.to(self.compute_device, non_blocking=True)`. - Inside `summon_full_params()`, `mixed_precision_cast_ran` is always `False`. Therefore, the return value for the `not p._is_sharded and mixed_precision_cast_ran` branch is unused. -`summon_full_params()` can only be called (before forward or after backward) or (between forward and backward). Given this, I cannot think of a case where we call `summon_full_params()`, the `FlatParameter` is already unsharded, but `reshard_after_forward` is `True`. The `FlatParameter` should be sharded (before forward or after backward), and the `FlatParameter` may only be unsharded (between forward and backward) if `reshard_after_forward` is `False`. - If parameter mixed precision is enabled and the sharding strategy is a sharded one, then inside `summon_full_params()`, the `FlatParameter` is unsharded in full precision. This involves allocating a new padded unsharded flattened parameter on GPU in full precision since `_full_param_padded` is in the low precision. Some comments: - Ideally, we reduce the complexity of the core code path: i.e. unshard for forward and pre-backward. If the return value is only used for `summon_full_params()`, we should consider if we can compartmentalize that logic. - The branching is complex, and some return values are never used, where this fact is not immediately obvious. We should see if we can reduce the branch complexity. Disclaimer: The difference in attribute semantics between `NO_SHARD` and the sharded strategies makes it challenging to unify the cases. This PR does not attempt to address that since it requires more design thought. However, it does attempt to reduce the complexity for the sharded strategies. ### Unshard: Core Code Path Let us trace through the new logical unshard. 1. `FullyShardedDataParallel._unshard(self, handles: List[FlatParamHandle], prepare_gradient: bool)` - This iterates over the handles and calls `handle.pre_unshard()`, `handle.unshard()`, and `handle.post_unshard(prepare_gradient)` in the all-gather stream. 2. `FlatParamHandle.needs_unshard(self)` - We take an aside to look at this key subroutine. - For `NO_SHARD`, this returns `False`. - For sharded strategies, this checks if the padded unsharded flattened parameter is allocated. The padded unsharded flattened parameter is the base tensor for the unpadded unsharded flattened parameter, which is a view into the padded one. Thus, the padded one's allocation fully determines if the `FlatParameter` is unsharded. - For sharded strategies, to accommodate the parameter mixed precision + `summon_full_params()` case, we introduce `_full_prec_full_param_padded`, which is the padded unsharded flattened parameter in full precision. The helper `_get_padded_unsharded_flat_param()` takes care of this casing and returns the padded unsharded flattened parameter. Instead of allocating a new tensor each time, we manually manage `_full_prec_full_param_padded`'s storage just like for `_full_param_padded`. 3. `FlatParamHandle.pre_unshard(self)` - For sharded strategies, the postcondition is that the handle's `FlatParameter` points to the tensor to all-gather. This should be on the communication device and in the desired precision. The allocation and usage of the low precision shard for parameter mixed precision and the CPU -> GPU copy for CPU offloading both classify naturally in the pre-unshard. - For sharded strategies, if the `FlatParameter` does not need to be unsharded, `pre_unshard()` is a no-op. This avoids unnecessarily allocating and freeing the low precision shard. - For `NO_SHARD`, we simply preserve the existing semantics. 4. `FlatParamHandle.unshard(self)` - If the handle was resharded without freeing the padded unsharded flattened parameter (e.g. `summon_full_params()` between forward and backward when `reshard_after_forward=False`), then the `FlatParameter` points to the sharded flattened parameter. We need to switch to using the unsharded parameter. This is a design choice. Alternatively, we may not switch to using the sharded flattened parameter in `reshard()` if we do not free the padded unsharded flattened parameter. However, the postcondition that the `FlatParameter` points to the sharded flattened parameter after `reshard()` is helpful logically, so I prefer this approach. - Otherwise, this allocates the padded unsharded flattened parameter, all-gathers, and switches to using the unpadded unsharded flattened parameter. - In the future, we may add an option to `unshard()` that additionally all-gathers the gradient. 5. `FlatParamHandle.post_unshard(self, prepare_gradient: bool)` - For sharded strategies, if using parameter mixed precision, this frees the low precision shard. More generally, this should free any sharded allocations made in `pre_unshard()` since the all-gather has been launched. If using CPU offloading, the GPU copy of the local shard goes out of scope after `unshard()` and is able to be garbage collected. **We should understand if there is any performance difference between manually freeing versus deferring to garbage collection since our usage is inconsistent.** For now, I preserve the existing semantics here. - `prepare_gradient` is meant to be set to `True` for the pre-backward unshard and `False` for the forward unshard. This runs the equivalent logic of `_prep_grads_for_backward()`. - This post-unshard logic (notably the gradient preparation) now runs in the all-gather stream, which is fine because we always have the current stream wait for the all-gather stream immediately after `FullyShardedDataParallel._unshard()`. IIUC, we do not need to call `_mp_shard.record_stream(current_stream)` (where `current_stream` is the default stream) because `_mp_shard` is allocated and freed in the same (all-gather) stream. - A postcondition is that the `FlatParameter` is on the compute device. It should also have the unpadded unsharded size (though I do not have a check for this at the moment). ### Unshard: `summon_full_params()` Now that we see how the logical unshard has been reorganized for the core code path, let us dive into `summon_full_params()`. The two constraints are: 1. If using parameter mixed precision, we should unshard in full precision. 2. We must determine if we should free the padded unsharded flattened parameter upon exiting. The first constraint is addressed as described before in the core unshard code path, so it remains to explore the second constraint. I propose a simple rule: **We free iff we actually unshard the `FlatParameter` in `summon_full_params()`** (i.e. it was not already unsharded). We perform a case analysis: **Parameter mixed precision enabled:** * `NO_SHARD`: `flat_param.data` points to `flat_param._local_shard`, which is the full precision unsharded flattened parameter. This is **not safe to free**. * `FULL_SHARD` / `SHARD_GRAD_OP`: We force full precision and all-gather to `_full_prec_full_param_padded`. We do not support `nested summon_full_params()`, so `_full_prec_full_param_padded` must be unallocated. We unshard, and it is **safe to free**. **Parameter mixed precision disabled:** * `NO_SHARD`: This is the same as with mixed precision enabled. This is **not safe to free**. * `FULL_SHARD` / `SHARD_GRAD_OP`: We all-gather to `_full_param_padded`. It may already be unsharded. * Already unsharded: The unshard is a no-op. This is **not safe to free**. * For `FULL_SHARD`, this can happen for the root FSDP instance after `forward()` but before backward. * For `SHARD_GRAD_OP`, this can happen for all FSDP instances after `forward()` but before backward. * Needs unshard: We unshard. This is **safe to free**. Therefore, we see that it is not safe to free when using `NO_SHARD` and when using a sharded strategy but the `FlatParameter` is already unsharded. This is precisely the proposed rule. There were two notable edge cases that the existing code did not address. 1. The existing code tests if the `FlatParameter` is already unsharded by checking the allocation status of `_full_param_padded`. When using parameter mixed precision, this is the incorrect tensor to check. If `_full_param_padded` is allocated (e.g. when `reshard_after_forward=False` and calling `summon_full_params()` between forward and backward), the already-unsharded check is a false positive, and `summon_full_params()` does not correctly force full precision. https://github.com/pytorch/pytorch/issues/83068 - This PR's `needs_unshard()` check correctly routes to the appropriate padded unsharded flattened parameter depending on the calling context (i.e. if it needs to force full precision or not). 2. The existing code does not free the GPU copy of the padded unsharded flattened parameter when calling `summon_full_params(offload_to_cpu=True)`. It unshards the `FlatParameter`, moves the padded unsharded flattened parameter to CPU, and sets the `FlatParameter` data to be the appropriate unpadded view into the padded unsharded flattened parameter on CPU. However, `_full_param_padded` still points to the all-gathered padded unsharded flattened parameter on GPU, which is kept in memory. https://github.com/pytorch/pytorch/issues/83076 - This PR frees the GPU copy and reallocates it upon exiting `summon_full_params()`. This is essential for avoiding peak GPU memory usage from increasing as we recurse through the module tree. There may be some cases where we can avoid reallocation altogether, but that can be addressed in a follow-up PR. - This PR offloads the *unpadded* unsharded flattened parameter to CPU directly instead of the *padded* one. As far as I can tell, there is no need to include the padding since unflattening the original parameters does not require the padding. - The relevant code is in the context manager `FlatParamHandle.to_cpu()`. ### Unshard: Mixed-Precision Stream This PR removes the mixed precision stream usage. As is, I do not think there is any extra overlap being achieved by the stream usage. The low precision shard is allocated and copied to in the mixed precision stream ([code](https://github.com/pytorch/pytorch/blob/1f99bdfcc4a3f97d28471a531d2b69def762f6ba/torch/distributed/fsdp/fully_sharded_data_parallel.py#L1401-L1412)), and the current stream (in this case the all-gather stream) waits for the mixed precision stream ([code](https://github.com/pytorch/pytorch/blob/1f99bdfcc4a3f97d28471a531d2b69def762f6ba/torch/distributed/fsdp/fully_sharded_data_parallel.py#L1414)). However, we immediately schedule an all-gather that communicates that exact low precision shard ([code](https://github.com/pytorch/pytorch/blob/1f99bdfcc4a3f97d28471a531d2b69def762f6ba/torch/distributed/fsdp/fully_sharded_data_parallel.py#L3338)) with no other meaningful computation between. If we remove the mixed precision stream, the low precision shard is allocated and copied to in the all-gather stream (including the non-blocking CPU -> GPU copy if using CPU offloading). Under this PR's design, we may consider a "pre-unshard" stream for all logical pre-unshard data transfers if we want to overlap in the future. IIUC, the overlap opportunity exists if there are multiple `FlatParameter`s per module, and we only have the all-gather stream wait for the data transfer corresponding to the local shard it communicates, not the others. If we agree on removing the mixed-precision stream for now, I will remember to delete it from `_init_streams()`. ## FSDP Runtime: Reshard Like with unshard, the first step is the look at the existing `_free_full_params()` and `_use_param_local_shard()`. A few notable observations: - For only `NO_SHARD`, `_free_full_params()` includes a call to `_free_mp_shard()`. - For `summon_full_params()`, there is a separate `_free_full_params_and_use_local_shard()` that duplicates the main logic of `_free_full_params()` and calls `_use_param_local_shard()`. - In `forward()`, if `reshard_after_forward=True`, we call `_free_full_params()` and then `_free_mp_shard()`. Hence, for `NO_SHARD`, the `_free_mp_shard()` is a no-op. - In the post-backward hook, we typically call `_free_full_params()` and `_free_mp_shard()`. The `_free_mp_shard()` is a no-op for `NO_SHARD` and if `reshard_after_forward=True`. Some comments: - The code certainly works, but some of the no-ops are subtle. When possible, we should make it clear when calls are no-ops or not. It is good that the existing code documents that `_free_mp_shard()` is a no-op in the post-backward hook when `reshard_after_forward=True`. However, there are still some non-obvious no-ops (around `NO_SHARD`). - We should see if we can avoid the duplicate `_free_full_params_and_use_local_shard()`. Let us trace through the logical reshard: 1. `FullyShardedDataParallel._reshard(self, handles: List[FlatParamHandle], free_unsharded_flat_params: List[bool])` - The two args should have the same length since they are to be zipped. - The goal of having `free_unsharded_flat_params` is that the caller should be explicit about whether the (padded) unsharded flattened parameter should be freed. The low precision shard is always meant to be freed (as early as possible), so there is no corresponding `List[bool]`. 2. `FlatParamHandle.reshard(self, free_unsharded_flat_param: bool)` - This frees the (padded) unsharded flattened parameter if `free_unsharded_flat_param` and switches to using the sharded flattened parameter. - Echoing back to forcing full precision in `summon_full_params()`, `_free_unsharded_flat_param()` frees the correct tensor by using `_get_padded_unsharded_flat_parameter()`. 3. `FlatParamHandle.post_reshard(self)` - I am not fully content with the existence of this method, but this seems to be an unavoidable consequence of `NO_SHARD`. Perhaps, this may be useful in the future for other reasons though. - Right now, this method is only meaningful for `NO_SHARD` + parameter mixed precision + outside `summon_full_params()`. `_mp_shard` is not freed in the post-unshard since it is also the low precision _unsharded_ flattened parameter, so we must delay the free until the the post-reshard. Below the `FlatParamHandle.reshard()` and `post_reshard()` layer, there should not be any no-ops. One final comment I will mention is that I like the `pre_unshard()`, `unshard()`, `post_unshard()`, and `reshard()`, `post_reshard()` organization because it makes it clear what the boundaries are and their temporal relationship. Through that, we can set pre- and post-conditions. Furthermore, we can eventually convert logic to hooks that may be registered on the `FlatParamHandle` (for `pre_unshard()`, `post_unshard()`, and `post_reshard()`). This may improve the customizability of FSDP. ## FSDP Runtime: `forward()` - This PR reorganizes `forward()` in preparation for non-recursive wrapping, which uses pre-forward and post-forward hooks that expect the signature `hook(module, input)`. For FSDP, the `module` and `input` arguments are not used. - This PR creates a new method `_fsdp_root_pre_forward()` to handle the logic only the root FSDP should run. ## FSDP Prefetching Finally, we dive into the prefetching changes. Some highlights: 1. This PR unifies the execution order validation and prefetching implementations. - Both involve the execution order and can be unified to share some boilerplate. 2. Execution order validation only runs when the distributed debug level is `INFO`. - We have yet to have one success case where we actually catch an unintended source of dynamism. The warning is also too verbose. Hence, we are gating it by the `INFO` level. 3. This PR moves prefetching to be with respect to groups of handles (as mentioned in the constructor comment). - This is essential for supporting prefetching with non-recursive wrapping. 4. This PR does not include "bubbles", i.e. modules with no handles, in the recorded execution order(s). This deviates from the existing implementation. - This makes prefetching possibly more aggressive (when there are such bubbles), but it should not have significant performance implications either way. 5. This PR changes backward prefetching to reset the post-forward order each iteration (as intended). 6. This PR changes forward prefetching to use the first iteration's pre-forward order instead of the first iteration's post-forward order. (We can discuss whether we want this in this PR or not. Otherwise, I can keep it as using the post-forward order to preserve the existing semantics.) This PR also removes the `all_gather_stream.wait_stream(current_stream)` before forward prefetching because it does not help with high GPU reserved memory. We can add that back if desired. ### Appendix #### Reverse Post-Forward Order Is Not Always the Pre-Backward Order The existing PT-D FSDP pre-backward prefetching uses the reverse post-forward order. <details> <summary>Model Code</summary> ``` class Model(nn.Module): def __init__(self): super().__init__() self.block1 = nn.Sequential( nn.Conv2d(3, 4, kernel_size=3), nn.BatchNorm2d(4), nn.ReLU(inplace=True), ) self.block2 = nn.Sequential( nn.Conv2d(4, 4, kernel_size=3), nn.BatchNorm2d(4), nn.ReLU(inplace=False), ) self.block3 = nn.Linear(12, 8) self.head = nn.Sequential( nn.AdaptiveAvgPool2d(output_size=(1, 1)), nn.Flatten(), nn.Linear(4, 10), ) def forward(self, x): x = self.block1(x) x = self.block2(x) x = self.block3(x) return self.head(x) model = Model().cuda() fsdp_kwargs = {} model.block1[1] = FSDP(model.block1[1], **fsdp_kwargs) # BN2d model.block2[1] = FSDP(model.block2[1], **fsdp_kwargs) # BN2d model.block1 = FSDP(model.block1, **fsdp_kwargs) model.block2 = FSDP(model.block2, **fsdp_kwargs) model.block3 = FSDP(model.block3, **fsdp_kwargs) model = FSDP(model, **fsdp_kwargs) ``` </details> <details> <summary>Execution Orders </summary> ``` Pre-backward hook for ('head.2.weight', 'head.2.bias') 140339520587136 (model) Pre-backward hook for ('weight', 'bias') 140339461194656 (block3) Pre-backward hook for ('0.weight', '0.bias') 140339520589776 (block2) Pre-backward hook for ('weight', 'bias') 140339520587664 (block2 BN) Pre-backward hook for ('weight', 'bias') 140339520586656 (block1 BN) Pre-backward hook for ('0.weight', '0.bias') 140339520588768 (block1) Pre-forward order: ('head.2.weight', 'head.2.bias') 140339520587136 (model) ('0.weight', '0.bias') 140339520588768 (block1) ('weight', 'bias') 140339520586656 (block1 BN) ('0.weight', '0.bias') 140339520589776 (block2) ('weight', 'bias') 140339520587664 (block2 BN) ('weight', 'bias') 140339461194656 (block3) Reverse post-forward order: ('head.2.weight', 'head.2.bias') 140339520587136 (model) ('weight', 'bias') 140339461194656 (block3) ('0.weight', '0.bias') 140339520589776 (block2) ('weight', 'bias') 140339520587664 (block2 BN) ('0.weight', '0.bias') 140339520588768 (block1) ('weight', 'bias') 140339520586656 (block1 BN) ``` </details> Differential Revision: [D39293429](https://our.internmc.facebook.com/intern/diff/D39293429) Pull Request resolved: https://github.com/pytorch/pytorch/pull/83665 Approved by: https://github.com/zhaojuanmao