ProcessGroupNCCL: use _TimeoutManager to provide sane NCCL abort semantics

[d4l3k]

This disables normal ProcessGroupNCCL timeouts and instead uses _TimeoutManager to provide user space NCCL abort semantics using a background timer thread and a CUDA event to track completion.

Notable changes:

• Every time we wait on ProcessGroupNCCL collective we register an abort callback that fires if it doesn't complete in time.
• Long timeouts will result in many extra timeout entries that aren't cancelled since we don't proactively cleanup the timer events.
• This depends on NCCL aborting which is known to be buggy in NCCL 2.25.1-1
• We use _opts_hook and _wrap_work to ProcessGroupWrapper to provide wrappers around the base ProcessGroupNCCL to add in this logic.
• We always disable the timeout passed to base ProcessGroupNCCL to avoid watchdog/heartbeat from failing. We may also need to set other options on creation to avoid having background NCCL errors crash the process.

This also updates the process_group_test.py resiliency tests since they were broken and always returned success messages. It now guarantees that operations scheduled on a PG on error exit and don't block forever.

TODO: Figure out if we need to set these envs when we are also overriding the timeout

• TORCH_NCCL_RETHROW_CUDA_ERRORS=0
• TORCH_NCCL_DUMP_ON_TIMEOUT=0
• TORCH_NCCL_PROPAGATE_ERROR=0

Test plan:

pytest torchft/process_group_test.py

I haven't tested this E2E with torchtitan yet but that's next on the list.

[fegin]

Want to confirm my understanding. Now we have ProcessGroupBabyNCCL and ProcessGroupNCCL. The former one uses multi-processing while the later use abort + timeout.

[fegin]

LGTM

[fegin]

Not quite sure why do we need this function if we just return the work?

[fegin]

Nvm, saw it is overridden later.