[c10d] [NCCL] Fix work handle for coalescing manager

[Aidyn-A]

Fixes #122807
The work handle of the coalescing job will be populated:

    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 1
    cm.wait() # actually waits

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

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/122849

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 08deacb with merge base dd8a24b ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 5, 5, linux.4xlarge.nvidia.gpu) (gh)
  test_foreach.py::TestForeachCUDA::test_binary_op_list_error_cases__foreach_add_cuda_float16

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[wconstab]

i don't understand the fix, and i'd appreciate if you also add a test case, probably to test_c10d_nccl.py.

the fix looks like its just adds some assertions and updates the coalescedDevice, but how does it impact the work obj that gets returned?

[Aidyn-A]

how does it impact the work obj that gets returned?

Here:

pytorch/torch/distributed/distributed_c10d.py

Lines 2037 to 2039 in 76a87e3

	if device:
	# Old style of letting each coll inside the context manager to call into C++ counterpart via python binding
	work = group._end_coalescing(device)

_end_coalescing was supposed to return a proper work object. However, because coalescedComm_ is nullptr it does call groupEnd, but it returns a nullptr for work object pointer:

pytorch/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp

Lines 2280 to 2286 in 76a87e3

	c10::intrusive_ptr<Work> ProcessGroupNCCL::endCoalescing(OpType optype) {
	if (coalescedComm_ == nullptr) {
	// There is no actual work being coalesced, return here
	groupEnd();
	coalescing_state_ = 0;
	return nullptr;
	}

[Aidyn-A]

cc @kwen2501

[kwen2501]

Thanks for finding the issue.
Hopefully #123170 explains the error.
I think the issue is really in _coalescing_manager so I am fixing it there.

[kwen2501]

Thanks for the fix! LGTM!

[Aidyn-A]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here