Repository: pytorch
Issue: comm hook error in BWD pass
Labels: oncall: distributed,module: ddp

### 🐛 Describe the bug

Dear Sir/Madam,

I have an issue with comm hook registration. The situation is the following: I am using 2 orthogonal communicators, one for model and one for data parallelism. However, some weights are shared among the model parallel GPUs, therefore I need to add an additional reduction step for those. I am trying to accomplish this using comm hooks. I have attached a small reproducer which should show the error on 4 GPUs.

```
import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(16, 16, 1)
        self.conv2 = nn.Conv2d(16, 16, 1)
        self.conv1.weight.is_shared_model = True
        self.conv1.bias.is_shared_model = True

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        return x
    

# set up global and local communicator
world_size = int(os.getenv('WORLD_SIZE', 1))
world_rank = int(os.getenv('WORLD_RANK', 0))

dist.init_process_group(backend = 'nccl',
                        rank = world_rank,
                        world_size = world_size)

local_rank = world_rank % torch.cuda.device_count()

# MP groups
tmp_group = dist.new_group(ranks = [0, 1])
if world_rank in [0, 1]:
    model_group = tmp_group
tmp_group = dist.new_group(ranks = [2, 3])
if world_rank in [2, 3]:
    model_group = tmp_group

# DP groups
tmp_group = dist.new_group(ranks = [0, 2])
if world_rank in [0, 2]:
    data_group = tmp_group
tmp_group = dist.new_group(ranks = [1, 3])
if world_rank in [1, 3]:
    data_group = tmp_group

# define comm hook:
def reduction_comm_hook(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:
    # group size
    data_groupsize = 2
    model_groupsize = 2
    
    # allreduce everything first:
    buff = bucket.buffer()
    
    # get future for allreduce
    fut = dist.all_reduce(buff, op=dist.ReduceOp.AVG, group=data_group, async_op=True).get_future()

    # get grads for shared weights
    params = bucket.parameters()
    
    # model parallel reductions
    m_grads = []
    for p in params:
        if not hasattr(p, "is_shared_model"):
            continue
        if p.is_shared_model:
            m_grads.append(p.grad.data) 

    # model part
    def model_reduction(fut):
        # allreduce the stuff
        if m_grads:
            # reduce remaining gradients
            coalesced = _flatten_dense_tensors(m_grads)
            nfut = dist.all_reduce(coalesced, op=dist.ReduceOp.SUM, group=model_group, async_op=True).get_future()
        else:
            nfut = torch.futures.Future()
            nfut.set_result(bucket.buffer())
            
        return nfut

    
    def model_copy(fut):
        if m_grads:
            # get the return value from the all_reduce
            coalesced = fut.value()[0]
            for buf, synced in zip(m_grads, _unflatten_dense_tensors(coalesced, m_grads)):
                buf.copy_(synced)

        nfut = torch.futures.Future()
        nfut.set_result(bucket.buffer())

        return nfut

    return fut.then(model_reduction).then(model_copy)

# init model and stuff
device = torch.device(f"cuda:{local_rank}")
model = Model().to(device)
ddp_model = DistributedDataParallel(model,
                                    device_ids = [device.index],
                                    output_device = [device.index],
                                    broadcast_buffers = False,
                                    find_unused_parameters = False,
                                    gradient_as_bucket_view = True,
                                    process_group = data_group)
    
ddp_model.register_comm_hook(state=None, hook=reduction_comm_hook)

inp = torch.ones((8, 16, 128, 128)).to(device)

for _ in range(10):
    out = ddp_model(inp)
    loss = torch.sum(out)
    loss.backward()
```

The error I am seeing is:

```
Traceback (most recent call last):
  File "comm_hook.py", line 117, in <module>
    loss.backward()
  File "/opt/conda/lib/python3.8/site-packages/torch/_tensor.py", line 352, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/opt/conda/lib/python3.8/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Got the following error when running the callback: TypeError: 'torch._C.Future' object is not subscriptable

At:
  comm_hook.py(88): model_copy
  comm_hook.py(97): reduction_comm_hook
```

I would like to know what I am doing wrong. If possible, can you help me to describe how to implement the comms pattern described above.

Best Regards
Thorsten Kurth

### Versions

PyTorch version: 1.11.0a0+17540c5
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.21.3
Libc version: glibc-2.31

Python version: 3.8.12 | packaged by conda-forge | (default, Oct 12 2021, 21:59:51)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-4.4.0-119-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.6.55
GPU models and configuration: 
GPU 0: Tesla V100-DGXS-16GB
GPU 1: Tesla V100-DGXS-16GB
GPU 2: Tesla V100-DGXS-16GB
GPU 3: Tesla V100-DGXS-16GB

Nvidia driver version: 470.42.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.3.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.3.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.3.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.3.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.3.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.3.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.3.2
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.2
[pip3] pytorch-quantization==2.1.2
[pip3] torch==1.11.0a0+17540c5
[pip3] torch-tensorrt==1.1.0a0
[pip3] torchtext==0.12.0a0
[pip3] torchvision==0.12.0a0
[conda] Could not collect

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang