[C10D] Fix coalescedCollective op Flight Recording

[wconstab]

Stack from ghstack (oldest at bottom):

• -> [C10D] Fix coalescedCollective op Flight Recording #120430

Also noticed and filed #120516 during this work. May land this as is and then test/fix the other varieties of coalesced collectives later.

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 b54bdbc with merge base 92a2b21 ():

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

• trunk / linux-focal-cuda12.1-py3.10-gcc9 / test (nogpu_AVX512, 1, 1, linux.2xlarge) (gh)
  Process completed with exit code 1.

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

[kwen2501]

Yes, I think you would need to put back the capture_status == c10::cuda::CaptureStatus::None check and the else path.

Would be also nice to pack everything:

inc_pending_event_queries()
check capture_status
dec_pending_event_queries()

inside workEnqueue so that we don't have these checks everywhere. (can be in a different PR)

[wconstab]

I think my reasoning here was that coalescing_state_ should NEVER be active inside this function. Becuase this function is part of the 'coalescing manager' implementation, which would call its own startCoalescing function later.

If that is correct, i'll add an assertion here to clarify it.

[wconstab]

i'm not sure how the cudagraph capture path works. Is it exclusive with coalescing or is it additive?

should the new logic be like this? I don't understand why the cuda event increment is unconditional and the decrement is conditional.

  // Notify graphs before we check the capture status preemptively
  at::cuda::CUDAGraph::inc_pending_event_queries();

  workEnqueue(work);
  if (capture_status != c10::cuda::CaptureStatus::None) {
    at::cuda::CUDAGraph::dec_pending_event_queries();
  }

[kwen2501]

Sorry for the confusion.

What I meant is like this:

  at::cuda::CUDAGraph::inc_pending_event_queries();

  if (capture_status == c10::cuda::CaptureStatus::None) {
    workEnqueue(work);
  } else {
    at::cuda::CUDAGraph::dec_pending_event_queries();
  }

It preemptively increases the pending work number, if this work is actually enqueued, then this is fine.
If this work ends up not being enqueued because CUDA Graph capture is active, then it winds back the pending work number.

[kwen2501]

You can check elsewhere in the code how workEnqueue(work) is conditioned. I guess they are more or less done the same way. Hence I propose putting the above CUDA Graph logic into the workEnqueue function itself, so that we don't have the if-else at every call site.

[wconstab]

hmm.

is it necessary to preemptively increment and then decrement? or can we just increment when we workEnqueue nad never decrement?

[kwen2501]

There was a reason. See here:
#110665 (comment)
Cc @eqy

[kwen2501]

Thanks for the last fix about workEnqueue. LGTM.

[eqy]

Yes, I think this explanation is accurate; one nit would maybe be editing "(a) capturing is already in progress, so we can not enqueue" to "(a) capturing is already in progress (we cannot enqueue in this case)."

[wconstab]

@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

[pytorchmergebot]

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

• Lint / lintrunner-clang / linux-job

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

[wconstab]

@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