FullyShardedDataParallel
class torch.distributed.fsdp.FullyShardedDataParallel(module, process_group=None, sharding_strategy=None, cpu_offload=None, auto_wrap_policy=None, backward_prefetch=BackwardPrefetch.BACKWARD_PRE, mixed_precision=None, ignored_modules=None, param_init_fn=None, device_id=None, sync_module_states=False, forward_prefetch=False, limit_all_gathers=False, use_orig_params=False, ignored_parameters=None)
A wrapper for sharding Module parameters across data parallel
   workers. This is inspired by Xu et al. as well as the ZeRO Stage 3
   from DeepSpeed. FullyShardedDataParallel is commonly shortened to
   FSDP.
Example:
  >>> import torch
  >>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
  >>> torch.cuda.set_device(device_id)
  >>> sharded_module = FSDP(my_module)
  >>> optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
  >>> x = sharded_module(x, y=3, z=torch.Tensor([1]))
  >>> loss = x.sum()
