PP >=2 Still Hangs in Multinode Despite Previous PR Fix

[chuhan-ouyang]

Bug description

After reading this PR,, I upgraded my environment to use a PyTorch Version (build from src) that includes this fix: pytorch/pytorch@e63c2b2.

However, when training DeepSeek on TorchTitan with PP = 2, DP Shard = 2, TP = 4, I am still getting a hang. NCCL shows a warning of 'mismatched bytes recv'. The hang occurs if PP >= 2.

The hang occurs after this NCCL warning:

[2025-09-01 21:12:17] nid001097:5838:6020 [0] transport/net_socket.cc:561 NCCL WARN NET/Socket : peer 10.249.23.66<41138> message truncated : receiving 8 bytes instead of 4. If you believe your socket network is in a healthy state, there may be a mismatch in collective sizes or environment settings (e.g. NCCL_PROTO, NCCL_ALGO) between ranks
nid001097:5838:6020 [0] NCCL INFO transport/net.cc:1393 -> 5

I tried enforcing the same NCCL_PROTO NCCL_ALGO across ranks but it did not fix the hang.

Attached are the NCCL and training logs.

PP_hang_NCCL.log
PP_hang_output.log

Versions

Environment: Running on 4 nodes, each with 4 A100 GPUs. Python 3.11.7. Torch: 2.9.0a0+git82d2d23.

DeepSeek Config: https://github.com/chuhan-ouyang/torchtitan/blob/deepseek/torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml

Thank you.

[wwwjn]

cc @H-Huang

[H-Huang]

I tried the config and it is running for me. Have you tried running for llama3 and is it working? I haven't see this before so trying to narrow down where this is model specific or environment specific.

Unfortunately I don't see anything useful from the logs. Can you reduce the communication timeout in the config so that it will actually timeout (you would see a ProcessGroupNCCL exception being thrown in the log) and we can get NCCL dumps?

Running with TORCH_NCCL_DUMP_ON_TIMEOUT=true will give the dump files which you can call torchfrtrace -j on. This will let us know what is the last collective issued per rank.