GitLab

* tmp

* test again

* test again

* add new test

* clean up

* add test file to the testlist

* more comments

* add changelog
Co-authored-by: Min Xu <min.xu.public@gmail.com>

Co-authored-by: Min Xu <min.xu.public@gmail.com>

- Fixes from Benjamin.

Original commit msg:
  - Fixes #1041. I just had a minute or two, hoping that it's enough :)
Co-authored-by: Min Xu <min.xu.public@gmail.com>

Co-authored-by: Min Xu <min.xu.public@gmail.com>

Co-authored-by: Min Xu <min.xu.public@gmail.com>

* [fix] better handling non-flatten in FSDP

- see the detailed comment about that backward firing case
- also minor debugging help in FSDP
- also minor fix in FPW's state dict

* [feat] disallow reset_parameters by default

* [feat] adding fsdp_instances API - useful in check wrapping by user code

* [fix] one line fix but more than a day of debugging

* fixed the case of loading combined check with empty fsdp instances

* fixed another bug around state loading the root/nonroot module full param caching due to not resharding after forward

* [feat] support .half and .float better

* fixed a bug in gather optim state losses extra keys from the original state_dict

* fixed a test failure in mixed precision

* fixed another bug affecting no_sync grad acc

* fixed a bug and a test in fsdp optim state

* fixed another corner case

* added a comment

* skip ssd offload tests

* skip fsdp one for ssd overload
Co-authored-by: Min Xu <min.xu.public@gmail.com>

* [fix] don't import ProcessGroup eagerly

- move the import into typing since it is only used for type checking
- fixes #1057

* more fixes

* one more

* tested at least
Co-authored-by: Min Xu <min.xu.public@gmail.com>

* [bug] fix optim state gather when there is empty FSDP instances

* fixes an anssert and a test bug

* [minor] [FSDP] add a better for pure fp16

* [minor] [wrap] add a flag to help fsdp pure fp16 wrapping

* [feat] add random_sparse_mask api

* correct test skip
Co-authored-by: Min Xu <min.xu.public@gmail.com>

* [fix]: add a context for supporting state_dict from a non-FSDP parent module

* formatting
Co-authored-by: Min Xu <min.xu.public@gmail.com>

* [feat] support sst disabled and dst disabled cases

* added tests
Co-authored-by: Min Xu <min.xu.public@gmail.com>

* update nightly version

* update wgit to use numpy for load/store

- this is introduced with new nightly torch version, which made torch.save() not
  producing deterministic bytes
- this make tensor<->numpy conversion and then do the save/load to avoid that issues.

* fixed tests
Co-authored-by: Min Xu <min.xu.public@gmail.com>

* added a profiling class

* no more type ignore after merging main

* fixed a int/round bug

* add unit tests

* skip if no cuda for a test
Co-authored-by: Min Xu <min.xu.public@gmail.com>

* update examples and comment

* fixed issue with fft/ifft only doing the last dim

* fixed a int/round bug; fixed tests

* add cuda tests

* add atol and rtol

* skip cuda test correctly
Co-authored-by: Min Xu <min.xu.public@gmail.com>

* [Feat] implements lossy_compress with tests

1. Implements a method lossy_compress that takes in a dense tensor and returns a reconstruction with sst and dst, and optionally with sparsity.

[Feat] Implements dense_sst_to_dst and sst_dst_to_dense methods and adds tests

1. Implements the dense_sst_to_dst and sst_dst_to_dense method.
2. Adds tests for perfect reconstruction with all top-k across different dims.
3. Adds tests for the two new methods.

* [fix]: experimental import fix

* [test]: catch issue #1042 in the future

* revert trigger for failure

* add numpy dep for users
Co-authored-by: Min Xu <min.xu.public@gmail.com>

Co-authored-by: Min Xu <min.xu.public@gmail.com>

* [Feat] dense to sst implementation
1. Implementation of dense_to_sst function.
2. calculating the threshold for both the cases of top-k-element and top-k-percentage (fraction)
3. assertions to verify that the top_k_elements is smaller than the numel along the same dim
4. top_k_percent to top-k conversion
5. When calculating SST, now the real part of the complex dense_freq is used instead of the magnitudes.

* [Feat, Tests] transform method addition, handling of top_k_element None case
1. Addition of a transform method
2. Adds code to handle the dim=None case for top_k_element

* [Feat, Refactor] Reorganizations, new assertions and fixes.
1. XOR for validation that both of topk percent and element are not set, or both simultaneously unset. One and only one is set.
3. Distills topk and percent both to topk using unified helper function .
5. Adds a scatter topk values function to scatter values for SST and in future DST.
6. Validation for percentage range, and ensures k is never 0.
7. Uses config validation, adds config validation for top_k_element > 0 if not None.

* [fix] handle EMA in the state_dict

* better fix

* [fix]: experimental import fix

* Update fairscale/experimental/__init__.py

* Update fairscale/experimental/__init__.py
Co-authored-by: Min Xu <min.xu.public@gmail.com>

* [fix] unclose FD and not load/store metadata many times

* one more stat

* Update fairscale/experimental/wgit/sha1_store.py

* add name to the objects when added

* dict key can be int from a state_dict

* removed top_level_objects key; it should be added into repo, not sha1_store
Co-authored-by: Min Xu <min.xu.public@gmail.com>

Co-authored-by: Min Xu <min.xu.public@gmail.com>

* flip per_tensor's default

* fixed original size computation
Co-authored-by: Min Xu <min.xu.public@gmail.com>

* additional metadata, step 1

* add gzip option to repo::add

* add repo:add's return value and some refactoring and todo

* added size metadata to sha1_store

* added names metadata to sha1_store
Co-authored-by: Min Xu <min.xu.public@gmail.com>

* formatting change, no logical change

* formatting and name change, no logical change

* [refactor] sha1_store's path arg

- make sha1_store's path arg directly the path, not its parent
- this is because sha1_store is not like a .git or a .wgit dir, which is
  nested inside another "working" dir. It is simply a store, which
  is using a given dir.
- updated repo and tests as well.

* remove a test warning due to deprecated API from torch

* [refactor] change how dot_wgit_dir_path is used

- it should only be assigned in __init__.
- we use it in error checking in the rest APIs.

* simplify the init a bit

* refactor the sanity check

* moved some functions, no code change

* [feat] added per-tensor add to the repo

* enabled gzip compression on add

* fix a unit test

* add a note

* make sha1 store work on general dict

* handle general state_dict from a model, not just a module's one-level OrderedDict

* formatting
Co-authored-by: Min Xu <min.xu.public@gmail.com>

Co-authored-by: Min Xu <min.xu.public@gmail.com>

Co-authored-by: Min Xu <min.xu.public@gmail.com>

Co-authored-by: Min Xu <min.xu.public@gmail.com>

Co-authored-by: Min Xu <min.xu.public@gmail.com>

Co-authored-by: Min Xu <min.xu.public@gmail.com>