Repository: pytorch
Issue: [Prototype][RFC] PyTorch FullyShardedDataParallel(FSDP) API Proposal
Labels: oncall: distributed,triaged,module: fsdp

FairScale FullyShardedDataParallel (FSDP) API supports large model training and is being quickly adopted by internal and external users. The long term goal of upstreaming the API to PyTorch is to release a stable, more general and performant version of FSDP with improved usability, new production features (like checkpointing APIs and etc) and newly integrated tools (compression, logging, debugging) for various production use cases. 

### Goal of Proposing a new API in PyTorch
New API proposed here is mainly to improve FSDP usability
1. minimal friction to migrate current FairScale FSDP use cases
2. simplified, stable and extensible configs
3. new feature: support to wrap a module that is constructed on meta device
    a. when model can not fit in the memory, users have to change model construction codes and manually wrap/shard layers by layers;   
    b. use meta device can potentially let users to use FSDP without changing model construction codes. e.g. users can pass a meta module to FSDP API, then the FSDP API can recursively wrap its submodules and materialize these modules. In this way, users can simply replace DDP API with FSDP API for larger models without changing any model codes. 

### About Fairscale FSDP
https://engineering.fb.com/2021/07/15/open-source/fsdp/
https://github.com/facebookresearch/fairscale/blob/master/fairscale/nn/data_parallel/fully_sharded_data_parallel.py

**Example of using FairScale FSDP API**

```
class model:
    def _init(self):
        self.layer1 = FSDP(nn.linear())
        self.layer2 = FSDP(nn.linear())
        self.layer3 = nn.linear()
fsdp_model = FSDP(model())
```

### PyTorch FSDP API Proposal

```
// Right now only support float32 and float16
class DataType(str, Enum):
   Float32 = "float32"
   Float16 = "float16"
   // future: "bfloat16", ...
   
// Many combinations, it is better to use data class to simplify
// if enabled, default to float16 types 
class MixPrecision:
    InputType: DataType = "float16"
    ParamType: DataType = "float16"
    ComputeType: DataType = "float16"
    BufferType: DataType = "float16"
    ReduceType: DataType = "float16"
    clear_autocast_cache: bool = false
    
// Right now it depends on MixPrecision is enabled or not, 
// but in theory it does not have to depend on mix precision
class CPUOffload:
     Grad: bool = false
     Param: bool = false 
     StateDict: bool = false
 

FullyShardedDataParallel(
        self,
        // Module could be materialized or on meta device, the goal of supporting 
        // meta device is that users do not have to manuall wrap submodules 
        // during mode construction, especially when models can not fit in the CPU memory. 
        module: nn.Module, 
        process_group: Optional[ProcessGroup] = None,
        mix_precision: Optional[MixPrecision] = None,
        cpu_offload: Optional[CPUOffload] = None,
        // If auto_wrap_policy is specified, submodules should not have been
        // wrapped yet, otherwise it will throw errors; 
        // Those submodules will be wrapped as 
        // FSDP instances based on this policy automatically. 
        // module could be materialized or is on meta device,
        auto_wrap_policy: Optional[callable] = None
 ): 
        // Leaf and reminder layers will not have auto_wrap_policy specified
        if auto_wrap_policy: 
            // Recursively wrap submodules base on user provided policy
            wrapped_module, num_wrapped_params = _recursive_wrap(module, process_group, bucket_cap_mb, mix_precision, cpu_offload, module_init_policy)
            if num_wrapped_params == len(module.parameters()): 
                return wrapped_module
        
        // unsharded modules 
        self.unsharded_modules = [m if not isinstance(m, fsdp) for m in self.modules()]
        // materialize and initialize modules on meta device
        for m in self.unsharded_modules:
            if m.device == 'meta' and hasattr(m, "reset_parameters")
                m.to_empty(device=current_device)
                m.reset_parameters()  
        
        // flatten needs torch.cat(); sharding needs flatten().chunk(), so at this 
        // point tensors have already been materialized
        flatten()
        shard()
       
        ....      
```

**Example of using PyTorch FSDP API**
```
class model:
    def _init(self):
        self.layer1 = nn.linear()
        self.layer2 = nn.linear()
        self.layer3 = nn.linear()
fsdp_model = FSDP(model(), 
                  auto_wrap_policy=functools.partial(default_auto_wrap_policy, min_num_params=1e7))
```


cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd @cbalioglu @gcramer23