[xccl] ProcessGroup collective hooks never fire on XPU, so NanCheckHook does not detect NaNs #6144

🐛 Describe the bug

On XPU, ProcessGroup collective hooks never run, so NanCheckHook (the TORCH_DIST_NAN_CHECK NaN check on collective inputs) never fires for XCCL. Found while porting test/distributed/test_c10d_nan_check.py to XPU, where the two NaN-detecting tests are skipped with a link to this issue. The details below were drafted with an AI assistant (Claude Code) and reproduced on our 4x Max 1100 machine.

Root cause. Upstream fires collective hooks from its c10d dispatcher kernels in torch/csrc/distributed/c10d/Ops.cpp, which call ProcessGroup::firePreHook / firePostHook (21 call sites). XPU registers its own kernels, TORCH_LIBRARY_IMPL(c10d, XPU, m) in src/xccl/Register.cpp (e.g. m.impl("allreduce_", allreduce_XPU)). Those kernels never call firePreHook / firePostHook, both at the PyTorch pin (bd00905) and on torch-xpu-ops main. Any hook registered with pg->registerPreHook(...) is therefore silently skipped for XPU tensors. NanCheckHook::attach is one of them (torch/csrc/distributed/c10d/hooks/NanCheckHook.cpp).

The device check itself works. Calling the op directly aborts as expected:

import torch
t = torch.ones(16, device="xpu"); t[3] = float("nan")
torch.ops.c10d.check_for_nan(t); torch.xpu.synchronize()
# -> NanCheck_XPU.cpp:54: ... Assertion `0` failed.  (core dumped)

Through the hook, the NaN is not detected (torchrun --nproc_per_node 2 repro.py):

import os, torch, torch.distributed as dist
from torch._C._distributed_c10d import NanCheckHook
r = int(os.environ["RANK"]); torch.xpu.set_device(r)
dist.init_process_group("xccl")
NanCheckHook.attach(dist.group.WORLD)
t = torch.ones(16, device="xpu"); t[3] = float("nan")
dist.all_reduce(t); torch.xpu.synchronize()
print(f"rank{r}: all_reduce with NaN input completed, NaN NOT detected")

rank0: all_reduce with NaN input completed, NaN NOT detected
rank1: all_reduce with NaN input completed, NaN NOT detected

On CUDA the same collective traps on the device-side NaN check, and on CPU it raises RuntimeError: NaN found in input tensor.

Affected tests in test/distributed/test_c10d_nan_check.py, both of which expect the check to fire:

• XcclNanCheckHookTest::test_nan_in_input_detected
• XcclNanCheckHookTest::test_env_var_auto_attach

The other three tests in the class pass on XPU, including test_legitimate_data_not_flagged, test_nan_in_recv_buffer_ok and test_off_by_default_and_removable, because they don't depend on the hook firing.

Likely fix: call pg->firePreHook(...) / firePostHook(...) in the XPU c10d kernels in Register.cpp, mirroring Ops.cpp. That would also enable other hooks for XPU, such as the flight-recorder hook and HealthCheckHook.

Closing condition: close once both tests above pass on XPU with their XPU skip removed from test_c10d_nan_check.py.

Versions

torch 2.16.0a0+git0e54bdd (PyTorch main), torch-xpu-ops bd00905, oneCCL 2022.1 (DLE 2026.1.2.25), 4x Intel Data Center GPU Max 1100, backend xccl.