Repository: pytorch
Issue: [RFC] DeepSpeed + PT Distributed Integration
Labels: oncall: distributed,feature,triaged

## Background


[DeepSpeed](https://github.com/microsoft/DeepSpeed) reduces distributed data parallel training memory footprint by partitioning parameters, gradients, optimizer states, and recomputing activations. See [this link](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/) for more details.

It takes care of the entire training iteration instead of exposing an `nn.Module`-like API.

```python
model_engine, optimizer, trainloader, _ = deepspeed.initialize(
    args=args, 
    model=net, 
    model_parameters=parameters, 
    training_data=trainset
)

for inputs in batches:
    outputs = model_engine(inputs)
    loss = criterion(outputs, labels)
    model_engine.backward(loss)
    model_engine.step()
```


DeepSpeed1 provides optimizer states partitioning. DeepSpeed2 provides the full feature set. 


## API

To integrate DeepSpeed into PyTorch, it will be great if we can have a consistent API, i.e., `model(inputs)`, `loss.backward()`, `optimizer.step()`. In this way, we can minimize the surprise to users and can minimize the code change when existing applications would like to adopt this technique. We can try to decompose DeepSpeed into the following two concepts, open to better names :) 


* `torch.distributed.nn.ShardedDataParallel`: comply with the `torch.nn.Module` API. 
* `torch.distributed.optim.ShardedOptimizer`: comply with the `torch.optim.Optimizer` API.


And training loops can look like the following, which is almost the same as local training.


```python
import torch.optim as optim

dist.init_process_group(...)
model = torch.distributed.optim.ShardedDataParallel(model)
opt = torch.distributed.optim.ShardedOptimizer(optim.Adam, model.parameters())
# model.parameters() only returns the ones owned by the current rank

loss_fn(model(inputs), labels).backward()
opt.step()
```



## Upstream DeepSpeed to PyTorch

### `ShardedOptimizer`


The optimizer part is relatively independent, and [fairscale](https://github.com/facebookresearch/fairscale) has already implemented one using PyTorch optimizer API. We can start the upstream effort from there. See [`oss.py`](https://github.com/facebookresearch/fairscale/blob/525e709bcb851db9eeb5d09a41501e6ef478d714/fairscale/optim/oss.py#L92-L98)

```python
class ShardedOptimizer(torch.optim.Optimizer):
    def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:
        loss = self.optim.step(closure=closure)
        for rank, param_groups in enumerate(self.partition_parameters()):
            for param_group in param_groups:
                for param in param_group["params"]:
                    dist.broadcast(param, rank, group=self.group)
        return loss
```


### `ShardedDataParallel`


I haven’t got into all details, but we might be able to wrap DeepSpeed algorithm as an `nn.Module` using custom autograd functions. The high-level idea is that we can split the input model into shards and then insert phony layers between shards to load and drop model parameters/grads accordingly. 


```python
# representing one shard of the model, which expose APIs to load/drop parameters
# and gather gradients
class ModelShard(nn.Module):

    def __init__(self, cpu_model_shard, owner_rank, pg):
        ...
        
    def forward_load(self):
        # materialize local GPU parameters, can enhance with bucketing
        futures = [pg.broadcast(p, owner_rank, async_op=True) 
                   for p in self.parameters()]
        # NB: this requires consolidating c10d work with torch.futures.Future
        torch.futures.wait_all(futures)

    ...
    
    def reduce_grads(self):
        futures = [pg.reduce(p, owner_rank, async_op=True)
                   for p in self.parameters()]
        torch.futures.wait_all(futures)

# The phony layer is a synchronization point between model shards.
# In the forward pass, it drops parameters in the previous shard and
# loads parameters for the next shard. In the backward pass, it does 
# the reverse and also gathers gradients to the owner.
# It does not change or create any outputs at all, instead it just
# forward the input as the output.
class PhonyLayer(Function)
    
    @staticmethod
    def forward(ctx, prev_shard, next_shard, *inputs)
        # check None accordingly
        prev_shard.forward_drop()
        next_shard.forward_load()
        ctx.prev_shard = prev_shard
        ctx.next_shard = next_shard
        return *inputs
        
    @staticmethod
    def backward(ctx, *grad_outputs):
        ctx.next_shard.reduce_grads()
        ctx.next_shard.backward_drop()
        ctx.prev_shard.backward_load()
        return *grad_outputs
```

With the above building blocks `ShardedDataParallel` can assemble the `PhonyLayer` and `ModelShard` objects in its forward pass. 

```python
class ShardedDataParallel(nn.Module):
    def __init__(self, cpu_model: nn.Sequential, pg: ProcessGroup):
        # split the input model into shards based on config/profiling/etc.
        self.model_shards = split(cpu_model, pg)
        
    def forward(self, *inputs):
        for prev, next in zip([None, *shards], [*shards, None]):
            inputs = prev(inputs) if prev else inputs
            inputs = PhonyLayer.apply(prev, next, *inputs)
        return inputs

```

## Integrate DeepSpeed into DDP


Existing `torch.distributed.DistributedDataParallel` provides gradient bucketing and comm/comp overlapping features through the `Reducer`, which would be useful for DeepSpeed too, especially as DeepSpeed targets large models. Instead of implementing everything from scratch for DeepSpeed, we can try to integrate DeepSpeed into DDP. The DeepSpeed optimizer can stay independent, and the rest of DeepSpeed algorithms can be decomposed to fit into the `Reducer` API. However, today’s `Reducer` API is not sufficient though, as it would materialize all buckets at construction time, install all autograd hooks at construction time, and always uses `AllReduce` to communicate gradients. We will need the [modularized DDP API](https://github.com/pytorch/pytorch/issues/37002) .

The current DDP algorithm can be decomposed and briefly represented using the following Python pseudo-code (although the current implementation is in C++). The four concepts (Grad Reader, Grad Bucketer, Comm Scheduler, and Grad Writer) in #37002 map to four functions below. The `grad_ready` function triggers DDP backward logics. 

```python
class ReducerBase:
    def __init__(self, model, pg, *bucket_cap_mb* = 25):
        self.model, self.*bucket_cap_mb*, self.pg = model, *bucket_cap_mb*, pg
        self.grad2bucket = self.build_buckets()
        # let param.grad points to bucket offset (view)
        self.register_grad_hooks()
        
    def register_grad_hooks(self):
        for p in self.model.parameters():
            # register self.grad_ready to p's GradAccumulator as callback
            
    # Grad Bucketer
    def build_buckets(self):
        # map self.model.parameters() to a list of GradBucket objects
        
    # Grad Reader
    def grad_ready(self, grad, idx):
        bucket = self.grad2bucket[idx]
        bucket.pending_grad -= 1
        if self.bucket.pending_grad == 0:
            # only launch one bucket when all its preceding ones are 
            # launched, so that we won't mismatch allreduce comms
            for i in range(min_pending_bucket, bucket.idx + 1):
                if self.buckets[i].pending_grad == 0:
                    self.bucket_ready(self.buckets[i])
       
    # Comm Scheduler
    def bucket_ready(self, bucket):
        self.futures.append(
            self.pg.allreduce(bucket, async_op=True).then(self.comm_ready)
        )
        self.pending_bucket -= 1
        if self.pending_buckets == 0
            # the last bucket should block to make sure when backward()
            # returns, all grads are indeed ready
            torch.futures.wait_all(self.futures)
        
    # Grad Writer
    def comm_ready(self, comm_future):
        pass
```


To integrate DeepSpeed with DDP, we can override the `build_buckets` and `bucket_ready` functions accordingly to take care of the gradient communication in the backward pass. Besides, as not all parameters will be materialized at construction time, it should leave out the invocation of `register_grad_hooks` in ctor as well.


```python
class ShardedReducer(RecuderBase):
    ...
    def build_buckets(self):
        # map self.model.parameters() to a list of GradBucket objects
        # only materialize the ones this rank owns. 
        
    # Comm Scheduler
    def bucket_ready(self, bucket):
        self.futures.append(
            self.pg.reduce(bucket, bucket_owner_map(bucket.idx), async_op=True)
        )
        self.model_shard.pending_bucket -= 1
        if self.model_shard.pending_buckets == 0
            torch.futures.wait_all(self.futures)
```



The `PhonyLayer` would still be necessary to load and drop model shards. However, instead of gathering grads, it only needs to run `grad_ready` accordingly to trigger `Reducer` code.


```python
class ModelShard(nn.Module):
    def __init__(self, cpu_model_shard, owner_rank, pg, reducer):
        ...
        self.reducer = reducer
        
    def register_grad_reader(self):
        # this shard must have been materialized when reaching here
        for p in self.parameters():
            # register self.reducer.grad_ready to p's GradAccumulator as callback

# this is the same as the above PhonyLayer except it no longer
# communicates grad on its own, instead it registers grad hooks 
# at dynamically loaded parameters in the backward pass.
class PhonyLayer(Function)
    
    @staticmethod
    def forward(ctx, prev_shard, next_shard, *inputs)
        # check None accordingly
        prev_shard.forward_drop()
        next_shard.forward_load()
        ctx.prev_shard = prev_shard
        ctx.next_shard = next_shard
        return *inputs
        
    @staticmethod
    def backward(ctx, *grad_outputs):
        # ctx.next_shard.reduce_grads()
        # check None accordingly
        ctx.next_shard.backward_drop()
        ctx.prev_shard.backward_load()
        ctx.prev_shard.register_grad_reader()
        return *grad_outputs
```


Then, the `SharedDataparallel` `nn.Module` becomes:

```python
class ShardedDataParallel(nn.Module):

    def __init__(self, cpu_model: nn.Sequential, pg: ProcessGroup):
        # Split the input model into shards based on config/profiling/etc.
        # Let each model shard keep a copy of the reducer, so that it can 
        # call reducer.grad_ready accordingly
        self.reducer = ShardedReducer(cpu_model, pg)
        self.shards = split(cpu_model, pg, self.redcuer)
        
    def forward(self, *inputs):
        # same as the above one
```

## Implementation Plan

* Phase 1: we should start from the `ShardedOptimizer`, as there are no road blockers as of today and fairscale already has an elegant implementation. 
* Phase 2: when the composable `Reducer` API is ready, we can move on to add parameter + gradients sharding. 
* Phase 3: add activation re-computation. 



cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @xush6528 @osalpekar @jiayisuse @agolynski