fix(fsdp): run torch_musa's FSDP2 replacements only on MUSA meshes

[yeahdongcn]

What / Why

torch_musa replaces parts of FSDP2 at import time with versions that create MUSA streams whatever the device mesh is. fully_shard on a CPU (gloo) mesh therefore fails in the root pre-forward under torch_musa:

• with MUSA devices visible, a MUSA stream waits on a CPU stream: AttributeError: 'NoneType' object has no attribute 'wait';
• with MUSA_VISIBLE_DEVICES=, the communication-context setup calls torch.musa.current_stream(): RuntimeError: No MUSA GPUs are available.

PyTorch runs the same CPU-mesh FSDP2 without error. vLLM-Omni's MAGI-2 test_compiled_regions_match_eager_under_dlo_and_hsdp (HSDP on a 4-rank gloo CPU mesh) fails this way on MUSA.

Change

New src/torchada/_fsdp2.py, installed by _patch_fsdp2_device_dispatch in _patch.py:

• torch.distributed.fsdp.fully_shard sends a non-MUSA mesh straight to PyTorch's fully_shard; MUSA meshes and mesh=None keep torch_musa's wrapper.
• FSDPCommContext.lazy_init, FSDPParamGroup.wait_for_unshard and post_backward, and FSDPState._root_post_backward_final_callback dispatch per call: a musa device runs torch_musa's function, any other device runs PyTorch's. The torch_musa module globals that hold these functions are the same dispatchers, so the methods still pickle by reference.
• PyTorch's methods that torch_musa replaces at import are rebuilt from the installed PyTorch source after checking that it matches the loaded code. If any of them cannot be recovered, nothing is installed and torch_musa's FSDP2 stays as it is, with one warning.
• A warning names any other method of these classes that torch_musa replaces (it still runs for every mesh), and a fully_shard attribute that is neither torch_musa's wrapper nor PyTorch's and so is not routed. A reference to torch_musa's wrapper bound before import torchada is not routed.

MUSA meshes keep running torch_musa's functions, with its streams, overlap level and collectives. Nothing is imported that torchada would not already import.

Verification

Tests: tests/test_fsdp2_device_dispatch.py, with a fake torch_musa FSDP2 layer (tests/_fsdp2_fakes.py) and a multi-process worker (tests/_fsdp2_mesh_worker.py):

• on any host, the dispatch, recovery, idempotency, fail-closed and pickling behaviour, and 2- and 4-rank gloo runs whose outputs and FSDP2 call traces match plain PyTorch;
• on a MUSA host, the same CPU-mesh runs against the real torch_musa (devices visible and hidden), and MUSA-mesh runs (musa and cuda meshes, 1-D and 1x2, default, overlap level 4 and both custom collective types) that must match torch_musa without this change. A MUSA-mesh comparison is skipped, naming the configuration and torch_musa's error or timeout, where torch_musa alone cannot complete it.

Results on MTT S5000 (torch 2.11.0.post2, torch_musa 2.11.0.post2+musa5.2.0, 2 cards):

• The new test file: 53 passed, 38 skipped. The skips are the 32 fake-layer cases, which skip once the real torch_musa is loaded, and 6 MUSA-mesh configurations that torch_musa alone cannot run: FSDP2 HSDP on a 1x2 MUSA mesh fails with RuntimeError: MCCL error ... unhandled musa error in the gradient all-reduce (with this change it fails the same way), and TORCH_MUSA_FSDP2_OVERLAP_LEVEL=4 hangs. On origin/main 7 of the MUSA-host tests fail; the ones that run FSDP2 on a CPU mesh fail with the two errors above.
• CPU-mesh reproducer, 1x2 and 1x4 meshes, devices visible and hidden: outputs and call traces identical to plain PyTorch, no MUSA initialisation; origin/main and torch_musa alone fail with the two errors.
• MUSA-mesh comparison against origin/main, 4 settings x 3 meshes: the 6 configurations torch_musa completes give identical outputs, FSDP2 call traces and overlap level; the other 6 fail identically with and without this change.
• vLLM-Omni MAGI-2 test_compiled_regions_match_eager_under_dlo_and_hsdp on vllm-project/vllm-omni main and on #8595, with this change and fix(device): register the device factory as the FX device builtin #124: passes with devices visible and hidden.
• pytest tests/ (other files): 2 failed, 591 passed, 9 skipped with and without this change, the same two TestInductorTemplateHeuristics failures.
• On CPU (torch 2.11.0+cpu), on this branch: the new test file 64 passed, 27 skipped (MUSA only); the full suite has the same failures as main.
• The MUSA runs above used this change on 24e89b2, before fix(musa): keep torch.isfinite asynchronous on MUSA floating tensors #120, fix(cpp_extension): accept torch 2.6+ include_paths/library_paths call signatures #121, fix(device): register the device factory as the FX device builtin #124 and test: pin the torch_musa version in the Inductor template heuristic tests #125 were merged. It rebases onto main without conflicts and its diff is unchanged.

Not covered / notes

• Only torch 2.11 was run.
• ShardedGradScaler on a CPU mesh is still torch_musa's version.
• The version is not bumped here.

AI assistance: Claude Code drafted the change, the tests and this description and ran the MUSA validation listed above.