Repository: pytorch
Issue: Incorrect Reduce collective result with `_coalescing_manager` 
Labels: oncall: distributed

Hi,

I am trying to execute a series of Reduce operations. I can execute a single Reduce operation using the `torch.distributed.reduce(...)` API. To better utilize the resources (network/compute), NCCL also provides a grouped collective execution API called `ncclGroupStart/End`, when a series of collective operations are surrounded by `ncclGroupStart/End` calls, NCCL will coalesce the individual calls into an optimized collective kernel. PyTorch seems to have the `_coalescing_manager` in distributed_c10d. It seems to do exactly what I'm trying to achieve. However, upon closer inspection, I see that the reduced values are incorrect.

Here is the exact minified python script which fails to compute the correct value. In this script, I reduce tensors with && without group mode && observe a difference in the result. On rank 0, `grad_slice2` and `grad_slice2_validation` operate on identical buffers, yet reduce to different values. 

```python
# cat > application.py
import os
import torch
import torch.distributed as dist

device = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(device)
default_group = dist.init_process_group('nccl')
world_size = dist.get_world_size()
rank = dist.get_rank()

numel = 2
grad_slice1 = torch.rand((numel,), dtype=torch.bfloat16, device=rank%8)
grad_slice2 = torch.rand((numel,), dtype=torch.bfloat16, device=rank%8)
root_ranks = [1, 0]
grad_slices = [grad_slice1, grad_slice2]
grad_slice1_validation = grad_slice1.clone().detach()
grad_slice2_validation = grad_slice2.clone().detach()
validation_slices = [grad_slice1_validation, grad_slice2_validation]

# Sequential reduces.
async_handles = []
for dst_rank, grad_slice in zip(root_ranks, grad_slices):
  async_handles.append(dist.reduce(grad_slice, dst=dst_rank, async_op=True))

for h in async_handles: h.wait()
torch.cuda.synchronize()
torch.distributed.barrier()

# Coalesced reduces.
async_handles = []
device = grad_slice1.device
with torch.distributed.distributed_c10d._coalescing_manager(
  group=torch.distributed.group.WORLD, device=device, reqs=async_handles):
  for dst_rank, grad_slice in zip(root_ranks, validation_slices):
    async_handles.append(dist.reduce(grad_slice, dst=dst_rank, async_op=True))

for h in async_handles: h.wait()
torch.cuda.synchronize()
torch.distributed.barrier()

if not torch.all(torch.eq(grad_slice1, grad_slice1_validation) == True).item() or \
  not torch.all(torch.eq(grad_slice2, grad_slice2_validation) == True).item():
  print(f"{rank} Incorrect result grad_slice1={grad_slice1} grad_slice1_validation={grad_slice1_validation}")
  print(f"{rank} Incorrect result grad_slice2={grad_slice2} grad_slice2_validation={grad_slice2_validation}")
  print("")
else:
  ...
```

Here is the output:

```bash
[1,0]<stdout>:NCCL version 2.16.2+cuda11.8
[1,0]<stdout>:0 Incorrect result grad_slice1=tensor([0.0101, 0.0140], device='cuda:0', dtype=torch.bfloat16) grad_slice1_validation=tensor([0.0101, 0.0140], device='cuda:0', dtype=torch.bfloat16)
[1,0]<stdout>:0 Incorrect result grad_slice2=tensor([15.9375, 14.5000], device='cuda:0', dtype=torch.bfloat16) grad_slice2_validation=tensor([16.0000, 14.4375], device='cuda:0', dtype=torch.bfloat16)
[1,0]<stdout>:
```

`15.9375`->`16.0000` and `14.5000`->`14.4375`.

One may argue that the difference here is minimal, but the computations should be exact here. So I'm trying to understand if the coalesce semantic is incorrect used OR there is a bug in ProcessGroupNCCL.

I've also written a C++ application to validate if this is actually an issue with NCCL group mode execution. That is available here: https://gist.github.com/0x6b64/6b52d912c98e5c55ac826dec86a11744

I'm working on 4 Amazon EC2 P4D instances, here is my run command.

```bash
mpirun -np $((8*4)) --hostfile ~/hostfile python application.py
```

I've not been able to reproduce the same incorrect computation directly working with NCCL. 

Note: the computation isn't always wrong, its only wrong is "sometimes" (whatever that means, but atleast it appears for the example I've attached here).

Any feedback here will help greatly!

### Versions
Collecting environment information...
PyTorch version: 2.0.0
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.26.3
Libc version: glibc-2.31

Python version: 3.10.8 | packaged by conda-forge | (main, Nov 22 2022, 08:26:04) [GCC 10.4.0] (64-bit runtime)
Python platform: Linux-5.10.173-154.642.amzn2.x86_64-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY

Nvidia driver version: 525.85.12
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu