• tmarkstrum's avatar
    [FSDP] Enable FSDP reduce scatter overlap (#897) · 0a526bcb
    tmarkstrum authored
    * enable reduce scatter overlap with other operations
    
    * fixed unit tests and added docstrings for the new parameters for fsdp
    
    * fixed more unit tests
    
    * fixed unit tests
    
    * avoided the pickle error on process_group_reduce_scatter
    
    * removed an unnecessary parameter in unit tests
    
    * remove unnecessary prints
    
    * fixed the docstring
    
    * skipped the test_offload unit test because this unit test failed in the main branch
    
    * removed the enable_reduce_scatter_overlap API parameter
    
    * added doc string for the defualt value of process_group_reduce_scatter parameter
    
    * fixed a syntax bug
    
    * fixed a bug which cause unitest failure
    
    * removed the all_gather in the ProcessGroupName enum
    
    * added more comment
    
    * changed the default value of process_group_reduce_scatter from None to ProcessGroupName.reduce_scatter
    0a526bcb
test_fsdp.py 35.2 KB