GitLab

* seemingly fix flakyness for gloo by checking all coms handles

* smal fix, remove unused flags

* remove usused flag

* add back max_batch flag

* adding back lazy_construction

* adding back lazy_construction

* add missing device arg
Co-authored-by: Anjali Sridhar <anj@devfair0443.h2.fair>

* [chore] 0.3.1 release

- mainly because vissl needs the new version
- added a doc on release steps

* Update CHANGELOG.md
Co-authored-by: anj-s <32556631+anj-s@users.noreply.github.com>

* review comments
Co-authored-by: anj-s <32556631+anj-s@users.noreply.github.com>

* Document FSDP tips and tricks in a separate file

* Fix packed sequence apply

* Update fairscale/utils/containers.py
Co-authored-by: Min Xu <24926999+min-xu-ai@users.noreply.github.com>

* [fix]: handle inputs with containers

- this is an issue surfaces by vissl as well
- fix seems to be super simple
- also cleaned up two tests with respect to multiple such tests
  running back to back (they don't do that presently)

* cleanup

* fix

* lint

Update example as style arg was removed in https://github.com/facebookresearch/fairscale/pull/345

:(

* [refactor] enhance wrap and auto_wrap

- Two things were done in this PR
  1. We don't need to import FSDP in wrap.py since the wrapper class
     type is stored in the context now.
  2. We can use a `auto_wrap_policy` function to customize wrapping policy
     for auto_wrap, including size of module, blacklist, exclude list
- The auto_wrap function got simplified a bit as a minor side effect.

* Update fairscale/nn/wrap/auto_wrap.py
Co-authored-by: Sean Naren <sean@grid.ai>

* addressed comments

* addressed more comments
Co-authored-by: Sean Naren <sean@grid.ai>

* [perf][minor] cache the rank lookups, small shardedddp perf fix
* tiny improvement, code quality

* change empty shard handling for OSS, do not rely on asserts
* code review

* [hotfix]: fix a bug in CI command

* debug

* debug

* bump cache ver

* fix

* eq

* check

* bump

* addressed comment

* [feat]: checkpoint and normalization

- added special handling of BN for track_running_stats and checkpointing
- we test BN/LN and checkpointing
- we test them with mixed precision

* Fix ampnet unit test by adding delegate object

* Remove comments

- cover them in terms of code path only
- numerically, AdaScale is different on SDP/FSDP than DDP, mainly
  due to partial view of the gradients.
- this doesn't mean it is definitely not useful but it is yet to
  be validated.
- not going to spend too much time until we have a real use case.

* [chore] move a test script

* add a shortcut for installing

* more skipping

* keep apt-get part

[refactor] Use logging in place of print statements, remove unused functions and other minor refactoring changes. (#461)

* fix pipe logging and other cleanups

* more log/debug changes

This adds a context manager that assists in making child modules with similar defaults.
Usage:
```
from fairscale.nn.misc import enable_wrap, wrap

with enable_wrap(**handleful_of_important_params):
    layer_1 = wrap(torch.nn.Linear(5, 5))
    layer_2 = wrap(torch.nn.Linear(5, 5), flatten_parameters=True) # Override parameters if you'd like

# without the context manager, creates Linear layer
layer_1 = wrap(torch.nn.Linear(5, 5))
```
If not within the FSDP context, this would be a no-op. This makes it easier to annotate layers without having to copy any changes in parameters.

* [chores]: CI py39 on GPU and more efficiency

* add test list files

* fix

* add test list files

* split benchmark run into 2 runs

* fix 1.8 version and balance benchmarks

* fix

* fix

* fix

* fix

* recording tests

* py39 install fix

* test again

* move tests

* reorg tests

* skip tests for torch 1.8 due to an upstream bug

* removed __init__.py from tests since it confuses pytest

* Revert "removed __init__.py from tests since it confuses pytest"

This reverts commit 7e156ba33dfaa5ed052031780613ec0cb57a45b0.

* don't include __init__ in file list

* notes on __init__.py and added missing ones

* fixed mypy in a test file

* balance test runtime

* better pip install

* balance more

* pip fix

* balance

* balance more, all test should finish within 20m now

* minor license update

* trying cu102

* more doc and addressed Ben's comments

* debugging

* debugging...

* [test] FSDP: add the failing test for #421

* skip on 1.5

* better skipping

* Update tests/nn/data_parallel/test_fsdp_grad_scaler.py
Co-authored-by: Sam Shleifer <sshleifer@gmail.com>
Co-authored-by: Sam Shleifer <sshleifer@gmail.com>

* [fix] FSDP corner case of all params at in the children

* lint

* fix

* tradeoff

* fix doc build

* review comments