Repository: pytorch
Issue: [RFC] PyTorch Sharder for distributed training
Labels: oncall: distributed,triaged,sharded_tensor

### 🚀 The feature, motivation and pitch

# 🚀 Feature

Provide a detailed API design for high-level PyTorch Sharding architecture. These generic sharding interfaces are for PyTorch users to shard models easily for distributed training.

# Motivation

There is a need to provide a standardized sharding mechanism in PyTorch. There are several types of model parallelism paradigms which require sharding (ex: pipeline parallelism, intra-layer parallelism etc.). We have proposed a standardized sharding spec and a few building blocks for users to deal with sharding PyTorch models in https://github.com/pytorch/pytorch/issues/55207. This is a follow-up to the aforementioned proposal after we have looked into the several recent use cases such as Megatron-LM.

# Pitch
Concrete goals for this proposal:

1. Provide a detailed design for how to extend the current sharding spec and sharding plan to accommodate further customizations.
2. Provide a detailed design for high-level APIs for module-level sharding.
3. Provide a way to optimize parameters when they are `ShardedTensor`s.
4. Provide APIs designs for additional low-level primitives:
     a. Provide a way to store the temporary partial result (`PartialTensor`) and replicate result (`ReplicatedTensor`) to avoid extra communication.
     b. Define the rule of how `ShardedTensor`, `PartialTensor` and `ReplicatedTensor` interacts with each other and how they can be extended to DDP and **FSDP.**
     c. Provide a way to reshard a `ShardedTensor` 


## ShardingPlan
PyTorch Distributed already has a well defined [ShardingSpec](https://github.com/pytorch/pytorch/issues/55207) that is flexible enough to express different sharding configurations. The ShardingSpec can be applied individually to tensors, however we need a way to specify Sharding for multiple Tensors (ex: apply a sharding configuration to a nn.Module). For this purpose, we have ShardingPlan:

```python
@dataclass
class ShardingPlan:
    """
    Representation of sharding plan.
    Dict keyed by FQN of parameter to PlacementSpec.
    """

    plan: Dict[str, PlacementSpec]
```

## Extensibility of PlacementSpec/ShardingSpec

PyTorch Distributed supports a bunch of PlacementSpecs like `DevicePlacementSpec`, `ChunkShardingSpec` etc. These are essentially based on the types of Tensors that can be created in OSS PyTorch. However, there are use cases internally where Tensors need to be placed on UVM or SSD and users can do more customization on top of that. For this purpose, we need to ensure that PlacementSpec/ShardingSpec is an extensible interface for customized use cases:
```python
class PlacementSpec(ABC):
    """
    Base class representing the placement of an entity. Subclasses of this
    class can be used to specify customized placements which might not be
    covered by existing APIs.
    """
  @abstractmethod
  def apply(self, tensor: torch.Tensor) -> torch.Tensor:
    """
    Retrieves a tensor and places it on the appropriate device.
    """


@dataclass
class DevicePlacementSpec(PlacementSpec):
    """
    Associates placement of an entity with a single device.

    Args:
        device(:class:`torch.distributed._remote_device`): The device to place the entity on.
    """

    device: torch.distributed._remote_device

    def __post_init__(self):
        if not isinstance(self.device, torch.distributed._remote_device):
            self.device = torch.distributed._remote_device(self.device)

    def apply(self, tensor: torch.Tensor) -> torch.Tensor:
        return tensor.to(self.device)

class ShardingSpec(object):
    """
    Base class representing sharding specifications.
    """
    @abstractmethod
    def build_metadata(self,
                       tensor_sizes: torch.Size,
                       tensor_properties: TensorProperties,
                       process_group=None) -> ShardedTensorMetadata:
        """
        Given a global tensor size list, define how to shard a tensor like this shape
        across ranks, return a list of ShardMetadata.
        """

    @abstractmethod
    def shard(self, tensor: torch.Tensor, src_rank: int = 0, process_group=None) -> "ShardedTensor":
        """
        Given a global tensor on src_rank, shard this tensor
        across ranks, return a ShardedTensor.
        """
```

With an interface like this, PlacementSpec and `ShardingSpec` can be extended by users to build their own customizable specifications. The `Sharder` can then use implementations of these interfaces to appropriately shard the model with these custom specs. In particular:
* `ShardedTensor` construction: `ShardedTensor.__init__` with different constructors (empty, ones, full, etc.) can use `ShardingSpec.build_metadata` to calculate the layout of shards and initialize the local tensors correspondingly
* `shard_tensor/parameter` API can just adapt to accept any sharding spec by calling into `ShardingSpec.shard` API, to shard the global tensor into a `ShardedTensor`.

## Customizable ShardingPlanner
As a first step, the user would manually build a `ShardingPlan` and provide that to the Sharder. However, long term we could provide a `ShardingPlanner` interface with some default planners in torch.distributed itself

```python
class ShardingPlanner(abc.ABC):

  @abstractmethod
  def build_plan(self, module: nn.Module) -> ShardingPlan:
```

## Sharder (Module-level sharding)
PyTorch Sharding Primitives mostly deal with tensor level operations whereas some new use cases does sharding at the nn.Module level. Although, `ShardedTensor` proposal provides a way for custom sharded operators, things like EmbeddingBagCollection cannot be implemented solely at the `ShardedTensor` level. As a result, we need some sort of higher module level sharding.

The proposal for PyTorch Distributed for module level sharding consists of the following key components:

* `shard_module`: High level `nn.Module` level API to shard an `nn.Module`. This replaces Tensors in the module with `ShardedTensor`s based on the provided `sharding_plan`.
* `sharding_plan`: Dict from parameter name to `PlacementSpec` indicating how a particular parameter should be sharded/placed across the available set of devices.
* `sharded_module_swapper`: Once sharding is done, this component is an extension that provides a way for users to swap in their custom sharded modules with the regular ones in the module.
* `reshard_output`: A way for users to specify how the output of a particular module should be resharded. For example, `ShardedTensor` based ops might produce a particular sharding configuration as output by default and this specification can reshard the output as needed by the user. For example, in megatron-lm we have a col-wise linear which produces sharded output, then point-wise ops like relu are run and finally a row-wise linear is run on the sharded output. Now the column-wise linear needs to produce non-sharded output and this is where an option like this can be used.

The API specification and pseudocode below elaborate on what this would look like.


```python
class ShardedModuleSwapper(ABC)

  @abstractmethod
  def process(module: nn.Module) -> nn.Module:
    """
    Processes a module and if needed swaps it with a custom sharded   
    Implementation. Should return ``None`` if no swapping should be 
    performed.

    The passed in module would consist of ShardedTensors and a common way 
    to perform module swapping would be to use the state_dict of the passed 
    in module and apply it to the new sharded module via its 
    load_state_dict method.
    """


def shard_module(
  module: nn.Module, 
  plan: ShardingPlan,
  sharded_module_swapper: ShardedModuleSwapper = None,
  reshard_output: Dict[str, ShardingSpec] = None,
):
"""
Shards a given module according to the provided sharding_plan. This method 
first shards all the parameters according to the given sharding_plan. Next,
If sharded_module_swapper is specified, it recursively traverses the module 
tree and calls ``sharded_module_swapper.process()`` for each submodule. If ``sharded_module_swapper.process()`` returns ``None``, the method continues to recurse further down the tree. If ``sharded_module_swapper.process()`` returns an ``nn.Module``, then the current module is replaced with the return value and we don't recurse further.

If ``reshard_output`` is specified, the module's output is resharded 
according to the provided spec.

Args:
  module: The module to apply sharding to
  sharding_plan: Dict from param name to PlacementSpec/ShardingSpec 
     to apply to each parameter.
  sharded_module_swapper: Implementation of ShardedModuleSwapper for module 
     swapping.
  reshard_output: Dict from submodule to the ShardingSpec in terms of how
     we would like to reshard the output of a module which produces a 
     ShardedTensor. It can also allow gathering the local output for each 
     rank as a Tensor to ensure further processing in a data parallel 
     fashion. 
"""
```

Example
Let’s demonstrate how these APIs can be used to replicate a simplified version of [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) like sharding.

```python
class SimpleMegatron(nn.Module):
    def __init__(self):
        self.fc1 = nn.Linear(...)
        self.relu = nn.Relu()
        self.fc2 = nn.Linear(...)
        
    def forward(self, inp):
        return self.fc2(self.relu(self.fc1(inp)))
```

This can be expressed as follows:

```python
# See https://pytorch.org/tutorials/prototype/skip_param_init.html for 
# 'meta' device
model = SimpleMegatron(..., device='meta')
# Default process group
dist.init_process_group(..., world_size = 8)
shard_module(
  module=model,
  plan = {
    # Shard on columns
    'fc1': ChunkShardingSpec(
       dim=1, placements=['rank0', ...,'rank7']
    ),
    # Shard on rows
    'fc2': ChunkShardingSpec(
       dim=0, placements=['rank0', ...,'rank7']
    ),
  },
  reshard_output = {
    # Reshard output of fc2 back to data parallel.
    'fc2': ChunkShardingSpec(
      dim=0, placements=['rank0', ...,'rank15']
    ),
  }
)
```

## Sharding + Data Parallel

One very common deployment of Sharding + Data Parallel is the [original Megatron-LM paper](https://arxiv.org/pdf/1909.08053.pdf). In this setting typically we do Tensor Parallelism across 8 GPUs of a single host and Data Parallelism across hosts. We plan to reuse PyTorch DDP to express Replication instead of having a separate higher level API for this.

This can be achieved by running a process per GPU such that each process within a host is part of a Tensor Parallel group and processes across hosts are part of a data parallel group.

The following pseudo code illustrates this for 2 hosts with 8 GPUs each:
```python
# As an example, for ShardedTensor across 8 GPUs being replicated across 
# two hosts.
# We would like the local shards of rank 0 and rank 8 to be replicated, 
# rank 1 and rank 9 to be replicated and so on.
class MyModel(nn.Module):
   def __init__(self):
       self.fc1 = nn.Linear(...)
       self.fc2 = nn.Linear(...)
       self.fc3 = nn.Linear(...)
  
   def forward(self, x):
       return self.fc2(self.fc2(self.fc1(x))))
      
m = MyModel()

# Shard model
self_rank = dist.get_rank(...)
local_rank = self_rank % 8
spec = ChunkShardingSpec(
   dim=0,
   placements=[
       "rank:{local_rank}/cuda:0",
       "rank:{local_rank + 1}/cuda:1",
       "rank:{local_rank + 2}/cuda:2",
       "rank:{local_rank + 3}/cuda:3",
       "rank:{local_rank + 4}/cuda:4",
       "rank:{local_rank + 5}/cuda:5",
       "rank:{local_rank + 6}/cuda:6",
       "rank:{local_rank + 7}/cuda:7",
   ],
)

shard_module(
  module=m,
  plan = {
    'fc1': spec,
    'fc2': spec,
    'fc3': spec,    
  },
)

# Create appropriate DDP process groups for each rank (need to call 
# new_group even on ranks which do not participate due to current API 
# semantics).
self_pg = None
for i in range(8):
  pg = dist.new_group([self_rank, (self_rank + 8) % 16])
  if i == local_rank:
    self_pg = pg

ddp = DDP(m, process_group=self_pg)

# Run DDP
res = ddp(torch.rand(...))
```

## Sharded Optimizers
With `ShardedTensor`, we also need an optimizer that works with module parameters that are `ShardedTensor`s. A ShardedOptimizer can be used as a wrapper to handle `ShardedTensor`.

```python
class ShardedOptimizer(object):
  def __init__(
    self, 
    named_params: Mapping[str, Union[Tensor, ShardedTensor]], 
    optimizer_class, 
    *optimizer_args, 
    **optimizer_kwargs):
    """
    Collects all Tensors and local shards of ShardedTensor, uses these 
    Tensors as ``params`` for the optimizer while building the optimizer 
    using ``optimizer_class``, ``*optimizer_args`` and 
    ``*optimizer_kwargs``.
    """

  def step(self, closure=None):
    """
    Regular optimizer step
    """
  
  # Functional optimizer API.
  def step_param(self, param: Tensor, grad: Optional[Tensor]):
    """
    Can be used to fuse backward and optimizer where ``step_param`` is 
    called as part of a backward hook on parameters.
    """

  def state_dict(self) -> Dict[str, Any]:
    """
    Returned state and param_groups will contain parameter keys
    instead of parameter indices like torch.optim.Optimizer.
    Example:
      param_group : 
      [
        {'lr': 0.1, 'momentum': 0, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': ["p1", "p2"]},
        {'lr': 0.01, 'momentum': 0, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': ["sharded_p3"]}
      ]
      state: {
        "p1": {'momentum_buffer': None},
        "p2": {'momentum_buffer': None},
        "sharded_p3": {'momentum_buffer': sharded_momentum, 'step': [1, 2]}
      }
    """

  def load_state_dict(self, state_dict: Mapping[str, Any]):
```

Until, https://github.com/pytorch/pytorch/issues/65018 is resolved `torch.nn.Module` doesn’t recognize `ShardedTensor` as a parameter and as a result, `module.parameters()` and `module.named_parameters()` won’t work to retrieve the appropriate `ShardedTensor`s.

We can have the following APIs in torch.distributed to work around this in the meantime:
```python
def module_named_params_with_sharded_tensor(
  module: torch.nn.Module, 
  prefix: str = '', 
  recurse: bool = True,
) -> Iterator[Tuple[str, Union[torch.nn.Parameter, ShardedTensor]]]:
```


## Additional Low Level Primitives
ShardedTensor alone is not sufficient and we need two additional primitives ReplicatedTensor and PartialTensor.

![Sharded Tensor (1)](https://user-images.githubusercontent.com/6937752/152415159-196f9786-4aa4-415a-93c5-501fec3ba543.jpg)
![Partial Tensor (1)](https://user-images.githubusercontent.com/6937752/152415261-23674903-ae10-4d95-82d2-2f1023b00cda.jpg)
![Replicated Tensor (1)](https://user-images.githubusercontent.com/6937752/152415280-31f91f22-b10a-4065-983a-59a77ba374f8.jpg)




## PartialTensor for storing temporary result
For some sharded op performed on a sharded tensor, the returned local results can be represented as a temporary Partial Tensor. If we want to consolidate partial results for next-step processing, we then first do the aggregation and convert results to a `shardedTensor`. For example, in Megatron LM, the second linear layer of MLP will return a PartialTensor because this is a row-wise sharding case. We then need to consolidate all local results to get the input for the next component. Below is what we proposed about the API design:

```python
class PartialTensor(object):
   """
   PartialTensor is an abstraction to represent Tensors that need 
   aggregation across multiple devices and multiple processes.

   PartialTensor is initialized in an SPMD like fashion where each rank
   initializes the PartialTensor. The PartialTensor object on each rank
   then only stores the local partial shard and the aggregation way to 
   get a full tensor.

   PartialTensor doesn't provide any Tensor like operations but is a 
   wrapper providing the Tensor representing the local partial shard.

   Using these, users can apply custom distributed sharded computations
   on top of this primitive.
   """

  def __init__(self, local_shard: Tensor, pg: ProcessGroup, aggr_op: ReduceOp):
      self.local_partial_shard = local_partial_shard # Tensor
      self.pg = pg
      self.aggr_op = aggr_op # Enum ReduceOp

  def reshard(self, resharding_spec: ShardingSpec) -> ShardedTensor:
  """
  Consolidate all partial tensors across multiple ranks and covert result
  to a shardedTensor. We achieve this in two logical steps:
     1. First aggregate the partial results across different ranks with
        the aggregation op defined in aggr_op. We use `reduce_scatter` as 
        an optimization method. 
     2. Transform the aggr result into a ShardedTensor based on the given
        shardingspec.
  """

# Now code for sharded linear with row-wise sharding (sharding_dim = 1).
# The Linear op returns a PartialTensor in this situation, we need to 
# consolidate the local results to get a complete result. 
Sharding_dim = 1
spec = ChunkShardingSpec(
  dim=sharding_dim,
  placements=[
    "rank:0/cuda:0",
    "rank:1/cuda:1",
    "rank:2/cuda:2",
    "rank:3/cuda:3",
  ],
)

Sharding_dim = 0
reshard_spec = ChunkShardingSpec(
  dim=sharding_dim,
  placements=[
    "rank:0/cuda:0",
    "rank:1/cuda:1",
    "rank:2/cuda:2",
    "rank:3/cuda:3",
  ],
)

# Perform sharded linear op.
fc = nn.Linear(10, 16)
shard_parameter(fc, "weight", spec)
fc = collect_local_shards(reshard_output(fc, reshard_spec))
result = fc(torch.rand(5, 10))

# Simple Megatron-LM example
Sharding_dim = 1
spec_rowwise = ChunkShardingSpec(
  dim=sharding_dim,
  placements=[
    "rank:0/cuda:0",
    "rank:1/cuda:1",
    "rank:2/cuda:2",
    "rank:3/cuda:3",
  ],
)

Sharding_dim = 0
spec_colwise = ChunkShardingSpec(
  dim=sharding_dim,
  placements=[
    "rank:0/cuda:0",
    "rank:1/cuda:1",
    "rank:2/cuda:2",
    "rank:3/cuda:3",
  ],
)
reshard_spec = spec_colwise

class SimpleMegatronLM(torch.nn.Module):
  def __init__(self, linear_size, rank=None):
     super().__init__()
     self.fc1 = torch.nn.Linear(*linear_size[0])
     self.gelu = torch.nn.GELU()
     self.fc2 = torch.nn.Linear(*linear_size[1])

  def forward(self, inp):
     return self.fc2(self.gelu(self.fc1(inp)))


sharded_megatron_lm = self.SimpleMegatronLM([[17, 12], [12, 29]])
shard_parameter(sharded_megatron_lm.fc1, "weight", spec_colwise)
shard_parameter(sharded_megatron_lm.fc2, "weight", spec_rowwise)
sharded_megatron_lm = collect_local_shards(
  reshard_output(sharded_megatron_lm, reshard_spec))
sharded_output = sharded_megatron_lm(torch.rand(15, 17))
```


## ReplicatedTensor 

ReplicatedTensor represents a tensor that is replicated across the `world_size` and has the same value on each rank. This concept is introduced to help us reduce the communication overhead when necessary. For example, if we perform element-wise operations on tensors that are exactly the same on each rank, we don’t need to perform any communication. The ReplicatedTensor allows us to do many more advanced optimizations in `ShardedTensor` ops. Below we propose the API design:

```python
class ReplicatedTensor(torch.Tensor):
    """
    Denotes a tensor which is replicated across the world_size and has the
    same value on each rank.
    """
    @classmethod
    def __torch_function__(self, func, types, args=(), kwargs=None):
        new_args = []
        new_kwargs = {}
        for arg in args:
            if isinstance(arg, ShardedTensor):
                # redispatch to ShardedTensor
                return arg._internal_torch_function(func, types, args, kwargs)
            if isinstance(arg, ReplicatedTensor):
                new_args.append(arg.tensor)
            else:
                new_args.append(arg)

        if kwargs is not None:
            for k, v in kwargs.items():
                if isinstance(v, ShardedTensor):
                    # redispatch to ShardedTensor
                    return v._internal_torch_function(func, types, args, kwargs)
                if isinstance(v, ReplicatedTensor):
                    new_kwargs[k] = v.tensor
                else:
                    new_kwargs[k] = v

        return ReplicatedTensor(func(*new_args, **new_kwargs))
```

## ReplicatedTensor/ShardedTensor/PartialTensor inter-op

For deterministic operations, the inter-op among those three tensors describes as follows:

```python
# ReplicatedTensor have same content on each rank, so deterministically
# the result should also be ReplicatedTensor
ReplicatedTensor + ReplicatedTensor = ReplicatedTensor

# Tensor might be different, so result should be tensor instead
ReplicatedTensor + Tensor = Tensor

# ShardedTensor/PartialTensor have different local shards on each rank,    # result should be ShardedTensor/PartialTensor
ReplicatedTensor + ShardedTensor = ShardedTensor

ReplicatedTensor + PartialTensor = PartialTensor
```

## Extend ReplicatedTensor to DDP/FSDP

We should extend the ReplicatedTensor concepts to DDP/FSDP as we are already using this concept in those projects. There are benefits for both DDP/FSDP to apply ReplicatedTensor, one is to consolidate the tensor representations and make sure we know ReplicatedTensor be exactly the same across ranks, the other is to enable more optimization opportunities by reducing the communication overheads. 

To apply ReplicatedTensor to DDP and FSDP, some rules are defined as follows

* DDP: 
Any parameters wrapped by PyTorch DDP, will be replaced with ReplicatedTensor.
Parameter updates with gradients produce torch.Tensors according to the rule we specified above, as gradients are different by different batches.

* FSDP: 
For parameters wrapped by PyTorch FSDP, they will be `nn.Parameter` and internally be sharded by FSDP(possibly integrating `ShardedTensor` with it in the future?). 
As FSDP materializing a layer during execution, it makes all the parameters to ReplicatedTensor, this involves `ShardedTensor` -> local tensor -> ReplicatedTensor.

## ShardedTensor Resharding
We want to reshard the temporary results from the sharded op performed on a sharded tensor. This is mentioned in the previous section. Below is what we proposed about the API design:

```python
class ShardedTensor(...):
  ...... 
  def reshard(
        self, 
        resharding_spec: ShardingSpec) -> ShardedTensor:
  """
  Reshards a tensor given the ``resharding_spec``.
  """
  

  def local_tensor(self) -> Tensor:
  """
  Return the local tensor on the current rank.
  """

# APIs to hook a module with reshard/collect local shards in the forward pass. These will be public APIs exposed to users.
def reshard_output(
      module: nn.Module, 
      resharding_spec: ShardingSpec) -> nn.Module:
"""
Hook a module with resharding in the forward pass according to the given      
``resharding_spec``.
"""
    def hook_func(_module, _input, output):
        return output.reshard(resharding_spec)
    module.register_forward_hook(hook_func)
    return module


def collect_local_shard(module: nn.Module) -> nn.Module:
"""
Hook a module with local shards collection in the forward pass.
"""
    def hook_func(_module, _input, output):
        return output.collect_local_shards()
    module.register_forward_hook(hook_func)
    return module

# Now code for sharded linear with col-wise sharding (sharding_dim = 0).
# The Linear Op returns a shardedTensor in this case with sharding_dim = 2.
# To get the final result, we will need to reshard it to dim = 0.
Sharding_dim = 0
spec = ChunkShardingSpec(
  dim=sharding_dim,
  placements=[
    "rank:0/cuda:0",
    "rank:1/cuda:1",
    "rank:2/cuda:2",
    "rank:3/cuda:3",
  ],
)

# Perform sharded linear op.
fc = nn.Linear(10, 16)
shard_parameter(fc, "weight", spec)
fc = collect_local_shards(reshard_output(fc, spec))
result = fc(torch.rand(5, 10))
```

## Checkpoint based resharding
While checkpointing Sharded modules, many use cases require loading the model in a different sharding configuration and possibly on a larger or smaller number of devices. For this purpose, we need to have support for resharding checkpoints at load time. 








### Alternatives

_No response_

### Additional context

_No response_

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @pietern @SciPioneer