[c10d][NCCL] _coalescing_manager does not produce proper work handle

[Aidyn-A]

🐛 Describe the bug

The list of work objects (works) is empty:

    with dist._coalescing_manager(group=pg_nccl, device=device, async_ops=True) as cm:
        dist.all_reduce(a)
        dist.all_reduce(b)
    print(len(cm.works)) # prints 0
    cm.wait()

In this situation cm.wait() becomes a no-op, which causes red-before-write issues whenever _coalescing_manager is used.

Versions

PyTorch version: 2.4.0.dev20240326+cu121

### Tasks
- [ ] https://github.com/pytorch/pytorch/pull/122651
- [ ] https://github.com/pytorch/pytorch/pull/122849

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang

[Aidyn-A]

#123170 does not fix main issue. @kwen2501, could you please try:

git clone https://github.com/NVIDIA/apex.git
cd apex
python setup.py develop --deprecated_fused_adam --distributed_adam --cuda_ext
python apex/contrib/test/optimizers/test_dist_adam.py -f

It fails with numerical mismatches because it starts to do ops before coalesced communications are over. The PR #122651 effectively fixes the issue, and all the numerical mismatches in the test are gone.

[kwen2501]

Hi @Aidyn-A, which issue is the main issue?

I tried the program you paste with #123170 :

import torch
import torch.distributed as dist

dist.init_process_group(backend="nccl", timeout=timedelta(seconds=10))
rank = dist.get_rank()
device = torch.device(f"cuda:{rank}")

# Create tensors
s = 100 * 1024 * 1024
with dist._coalescing_manager(device=device, async_ops=True) as cm:
    dist.all_reduce(torch.ones(s, device=device))
    dist.all_reduce(torch.ones(s, device=device))

print(len(cm.works))
cm.wait()
torch.cuda.synchronize()
print(f"{dist.get_rank()} done")

The is output I got:

NCCL version 2.20.5+cuda12.1
1
1
1 done
0 done

[Aidyn-A]

The main Issue we are seeing is in apex.

#123170 does not fix main issue. @kwen2501, could you please try:

git clone https://github.com/NVIDIA/apex.git
cd apex
python setup.py develop --deprecated_fused_adam --distributed_adam --cuda_ext
python apex/contrib/test/optimizers/test_dist_adam.py -f

It fails with numerical mismatches because it starts to do ops before coalesced communications are over. The PR #122651 effectively fixes the issue, and all the numerical mismatches in the test are gone.

[kwen2501]

Oh you are right. Even though len(cm.works) is 1, the point it captures the stream is before nccl kernel is enqueued. The reason is that we have multiple level of ncclGroupStart()|End() in this case:

ncclGroupStart();    # _start_coalescing
ncclGroupStart();    # collectiveCoalesced
ncclAllreduce();
ncclAllreduce();
ncclGroupEnd();    # collectiveCoalesced
ncclGroupEnd();    # _end_coalescing

NCCL won't enqueue kernel until hitting the last GroupEnd().

[kwen2501]

In the doc string, there is the following line:

pytorch/torch/distributed/distributed_c10d.py

Lines 1970 to 1971 in 64d7430

	device (`torch.device`, optional): Default is None, set to a device if
	there isn't a `**_coalesced` implementation by the backend.

Since all_reduce has a "Fast-Path" implementation for coalescing, we previously did not expect users to set the device. But I understand why libraries (like Apex) may want to set it in general. We will think of ways to improve this UI.

[kwen2501]

I closed #123170, let's go with #122849 .
Thanks for helping me understand the problem better. And sorry about the noise!