pytorch

21a0e809 - [Composable API] Match `fully_shard()` comm. schedule with wrapper FSDP (#90387)

[Composable API] Match `fully_shard()` comm. schedule with wrapper FSDP (#90387) - This PR introduces a new concept, the _communication module_ (denoted `comm_module`), that represents the module responsible for the unshard/reshard pair for a `FlatParamHandle`. This is well-defined because the current design assumes that each `FlatParamHandle` only has _one_ unshard/reshard pair for either the forward or backward pass. - For the wrapper code path, the `comm_module` is exactly the module already being passed to the `FlatParamHandle` constructor. - For the composable code path, the `comm_module` is not necessarily the module already being passed to the `FlatParamHandle`. This is because the module already being passed is always the local FSDP root module to give complete FQNs, instead of local FQNs. Distinguishing the communication module from the local FSDP root module can provide more flexibility for non-recursive wrapping designs in the future. - This PR adds a unit test `test_unshard_reshard_order` that explicitly checks that `_unshard` and `_reshard` are called in the exactly the same order across the two code paths. - This PR does not fix `test_checkpoint_fsdp_submodules_use_reentrant`. However, the error message changes, so this PR accommodates that. - The error is now the same as if we used the equivalent wrapper FSDP: ``` test_model.u1 = FSDP(test_model.u1, use_orig_params=True) test_model.u2 = FSDP(test_model.u2, use_orig_params=True) ``` - The error is also the same as if we used wrapper FSDP with `use_orig_params=False`, so it is not unique to `use_orig_params=True`. --- **`comm_module` Example** ``` model = Model( seq1: nn.Sequential( nn.Linear nn.ReLU nn.Linear nn.ReLU ) seq2: nn.Sequential( nn.Linear nn.ReLU nn.Linear nn.ReLU ) ) policy = ModuleWrapPolicy({nn.Sequential}) fully_shard(model, policy=policy) FullyShardedDataParallel(model, auto_wrap_policy=policy) ``` - This policy constructs two `FlatParamHandle`s, one for `seq1` and one for `seq2`. - `FullyShardedDataParallel` will pass `seq1` and `seq2` as the `module` argument to the two `FlatParamHandle`s, respectively. - `fully_shard()` will pass `model` as the `module` argument to every `FlatParamHandle`. - `FullyShardedDataParallel` will pass `seq1` and `seq2` as the `comm_module` argument to the two `FlatParamHandle`s, respectively. - `fully_shard()` will pass `seq1` and `seq2` as the `comm_module` argument to the two `FlatParamHandle`s, respectively. Pull Request resolved: https://github.com/pytorch/pytorch/pull/90387 Approved by: https://github.com/mrshenli