• Min Xu's avatar
    [feat] FSDP: supporting multiple flatten parameter groups (#711) · ab71efb3
    Min Xu authored
    
    
    * [feat] FSDP: supporting multiple flatten parameter groups
    
    - step 2: extending FPW to support multiple flat params groups
    - FSDP still only use one group
    - unit test does this the new code paths
    - updated the changelog
    
    * first cut, mypy passed
    
    * test_flatten_params_wrapper.py::TestFlattenParams tests pass
    
    * added two more test cases and fixed a case in the code
    
    * fixed one bug with param_path_infos
    
    * fixed two more tests with hardcoded flat_param names
    
    * Update CHANGELOG.md
    Co-authored-by: default avatarMin Xu <min.xu.public@gmail.com>
    ab71efb3
test_fsdp_state_dict.py 11.7 KB