facebookresearch / facebookresearch/fairscale
[FSDP] Add store weight/grad norm option
- Dominant language
- Python
- Stars
- 3.4k
- Forks
- 293
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
FSDP to offer the possibility to compute the norms of the weights and norms of the gradients on the fly, when the weights / gradients are available with an option like `compute_weight_norm=True` in the constructor and an API to retrieve those weights.
## Motivation
Following issue https://github.com/facebookresearch/fairscale/issues/644 and the attempt in solving it in PR https://github.com/facebookresearch/fairscale/pull/692, providing a view of the parameters in the view is too low level and does not offer the best user experience as there are plenty of ways to trip and fall (see PR for the limitations: https://github.com/facebookresearch/fairscale/pull/692).
Following a discussion with @min-xu-ai, the best is to go back at the uses cases motivating https://github.com/facebookresearch/fairscale/issues/644:
1. be able to easily compute the weight and gradient norms when `flatten_parameters=True` needed for LARC like optimisers
2. be able to have separate LR, regularisation for each parameters (for instance, regularise only bias and not weight or the other way around)
This issue proposes to solve item 1.
## Interested parties
CC: @min-xu-ai @myleott @prigoyal
Contributor guide
Assessment
This issue has not been assessed yet.