[FSDP] `use_orig_params=True` Follow-Ups & Known Issues
Labels: oncall: distributed, triaged, module: fsdp

This issue tracks some follow-ups and known issues with `use_orig_params=True`.

For classification, we define a few tags:
- [<==] Supported in existing FSDP / with `use_orig_params=False` but unsupported with `use_orig_params=True`
- [===] Supported in existing FSDP / with `use_orig_params=False` and supported with `use_orig_params=True`
- [!!!] Unsupported in existing FSDP / with `use_orig_params=False` and unsupported with `use_orig_params=True`

If a [To Do] tag is combined with [===] then that means the "supposed with `use_orig_params=True`" is not yet true but is expected to become true.

---

**Activation Checkpointing Compatibility**
- [To Do] Non-reentrant activation checkpointing is untested.

**`sync_module_states=True`**
- [To Do]
https://github.com/pytorch/pytorch/blob/1655b47a384d5e6ba31420b5daee1c029a821387/torch/distributed/fsdp/fully_sharded_data_parallel.py#L1498-L1499

**[<==] `named_parameters()` between Pre-Forward and Post-Forward**
- ~For `use_orig_params=True`, calling `fsdp_model.named_parameters()` between the pre-forward and post-forward (specifically, after the unshard and before the reshard) will result in **nothing being returned**. This is because the original parameters are switched to being _plain_ `Tensor` views into the underlying `FlatParameter`, not `nn.Parameter` views. Since they are plain `Tensor` views, they are no longer registered and hence not returned by `named_parameters()`.~
- For `use_orig_params=False`, doing so returns the `FlatParameter`s because they are always registered.
- ~Given the existing design, this is a hard blocker. However, we are currently unaware of any users requesting for this, and this only affects `named_parameters()` -- the parameters are still directly accessible.~
- To support TorchDynamo + FSDP, we actually register the plain `Tensor`s by manually modifying the module's `_parameters`. This means that `named_parameters()` returns the expected generator, only the parameters are `Tensor`s instead of `nn.Parameter`s, and the references should not be relied upon across iterations.

**[!!!] Parameter Changes**
- The user may change the parameter variable, i.e. `module.param = new_param`. In that case, FSDP needs to update its internal reference to the parameter variable and related metadata.
- Since from the experience with DDP, parameter change is not an important ask, we **do not plan to formally support this**. We can provide some warning to prevent silent errors.
- If we were to support this, then FSDP needs to update its internal reference to the parameter variable and related metadata. Moreover, FSDP should keep `weakref`s to the parameters to enable immediate garbage collection of the old parameter, assuming that there are no issues when using `weakref`s. Finally, FSDP should handle shared parameters becoming non-shared and vice versa.
```
# Shared becoming non-shared
lin = nn.Linear(5, 5)
lin.weight = lin.bias
fsdp_model = FSDP(lin, use_orig_params=True)
lin.weight = nn.Parameter(torch.randn((5, 5)))
# Non-shared becoming shared
lin = nn.Linear(5, 5)
fsdp_model = FSDP(lin, use_orig_params=True)
lin.weight = lin.bias
```

**[===] Gradient Changes**
- The user may change the gradient variable, i.e. `module.param.grad = new_grad`. Unlike for parameters, `new_grad` maybe `None`. This typically encodes the intention to zero the gradient while saving memory and can happen via `optim.zero_grad(set_to_none=True)` or `module.zero_grad(set_to_none=True)`.
- A `None` gradient and a zero gradient are identical with respect to gradient computation, but they are not identical with respect to optimizer computation. The optimizer still steps on parameters with zero gradient, which means that optimizer states may update, while the optimizer does not step on parameters with `None` gradient.
- [Done] To address the difference between `None` and zero gradient, FSDP should maintain a mask over the original parameters that indicates whether it is `None` or not. Since the `flat_param.grad` interfaces with gradient computation, it is fine for it to treat `None` and zero identically, but when pointing the `orig_param.grad`s into `flat_param.grad`, FSDP must use the mask to set `orig_param.grad = None` as needed.
- [<==] FSDP can only propagate changes to the original parameters and their gradients to the underlying `FlatParameter`s and their gradients during FSDP-owned logic. This exposes windows during which those changes are not reflected. For example, if the user calls `optim.zero_grad(set_to_none=True)`, the memory frees cannot happen until the pre-forward, when execution enters FSDP-owned logic.

**[===] Optimizer Checkpointing**
- [To Do] We need to implement optimizer state checkpointing for `use_orig_params=True`, where the optimizer states are now per original parameter instead of per `FlatParameter`. Notably, some parameters may be unused and hence not have state. The design for implementing this is still under discussion.

**Performance Benchmarking**
- ~[To Do] Compare performance of `use_orig_params=True` vs. `False` on some representative workloads.~
- Benchmarking using the `transformer_framework` repo shows parity for T5-11B, RegNet-9B, and DeepViT-8B.

**`.data`**
- The entire approach to `use_orig_params=True` relies on setting the original parameters' `.data` in a controlled way. Ultimately, we want to migrate from this to a more robust approach such as a tensor subclass. We are investing in this `.data` approach to unblock in the near-term and since it is still a good exercise to uncover edge cases.

---

Some engineering follow-ups:
- Formalize contract for parameter and gradient writeback. Then, document it for users, and harden unit tests following that contract.

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @kwen2501