facebookresearch / facebookresearch/fairscale

[FSDP] Flatten parameters by group

Open
#697 7 comments 1 reaction 1 assignee Claimed by @min-xu-ai View on GitHub
FSDP
Dominant language
Python
Stars
3.4k
Forks
293
PR merge metrics
No merged PRs in 30d

Description

## 🚀 Feature

FSDP to offer the possibility to flatten parameters by group, for instance, to flatten all biases separately from the other 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 2.

## Workarounds

There is no workaround for the moment when `flatten_parameters=True`: having different wrappers for different modules does not offer the granularity required to flatten biases and weights of the `nn.Linear` layer separately for instance.

## Interested parties

CC: @min-xu-ai @myleott @prigoyal

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.