[Bugfix] Fix AllReduceFusionPass crash with PP+TP configurations

[haosdent]

Purpose

Fix NCCL "Duplicate GPU detected" error when using -tp 2 -pp 2 -cc.fuse_allreduce_rms=True.
FIX #35426

PR #34109 introduced a regression by using TorchDistBackend for multi-group TP configurations, which has upstream issues with bcast(), IPC socket collisions, and device assignment. Tracked at flashinfer-ai/flashinfer#2647.

This PR gracefully disables AllReduceFusionPass when PP > 1 or DP > 1 until the upstream flashinfer fixes land.

Test Plan

Test Result

[gemini-code-assist]

Code Review

This pull request addresses a crash in AllReduceFusionPass when using both pipeline parallelism (PP) and tensor parallelism (TP). The fix involves two main parts: first, adding guards to gracefully disable the pass for PP+TP and DP+TP configurations, which is a sensible workaround for an upstream limitation in flashinfer. Second, it re-introduces a patched _TPCommBackend to fix communication issues in flashinfer's distributed backend when multiple TP groups are active. A new regression test is also added to verify that the pass is correctly disabled in a PP+TP setup, ensuring the fix is effective and preventing future regressions. The changes are well-structured and address the described issue effectively. I have one suggestion to improve maintainability by using a standard torch.distributed API instead of a custom one.

[haosdent]

@ProExpertProg @hjjq Can take a look when you are available, many thanks!

[ilmarkov]

Thanks for the PR!
When flashinfer fixes the issue with SymmMemBackend, they may also fix the issue with broadcast_object_list and IPC socket opIds so we wouldn't have to create custom _TPCommBackend. Do we actually need this _TPCommBackend at the moment given that we disable the fusion for DP>1 and PP>1?

[haosdent]

When flashinfer fixes the issue with SymmMemBackend, they may also fix the issue with broadcast_object_list and IPC socket opIds so we wouldn't have to create custom _TPCommBackend.

Thanks @ilmarkov I see, if flashinfer would fix together, then I think we don't need it and could just disable first. Let me update the PR later

BTW, do you have the issue links, and I could put them in the comments as references.

[ilmarkov]

@haosdent I have filed an issue in Flashinfer: flashinfer-ai/flashinfer#2647

[haosdent]

Thanks @ilmarkov , I have updated the PR, can take a look when you are available