DCP save deadlocks in SavePlan metadata exchange on Ascend NPU — add process_group to dcp.save #9

Setup

Kimi-K3 (2.8T MoE), FSDP2 + expert parallel, 512 / 256 ranks, Ascend 910C, torch 2.10.0, torch_npu 2.10.0.post2, torch.distributed.checkpoint save.

Symptom

dcp.save never returns. It blocks in distW.all_reduce("write", ...) → gather_object — the SavePlan metadata exchange that runs before any shard is written — then the job dies with:

RuntimeError: operator():../torch_npu/csrc/aten/common/LocalScalarDenseNpu.cpp:23
NPU function error: c10_npu::acl::AclrtSynchronizeStreamWithTimeout(copy_stream), error code is 507046
[Error]: In the specified timeout waiting event, all tasks in the specified stream are not completed.
synchronize stream timeout, timeout = 340000ms

Reproduced at 512 and 256 ranks, with and without prior training; not an OOM, no rank died before the hang; the hang precedes shard writing (zero files observed at 256 ranks).

How tested

A/B runs, identical config; single variable — dcp.save(..., process_group=) vs the default (HCCL) group.

Results

default (HCCL): deadlock at both scales, in the metadata exchange before writing
gloo: save completes — 1143 s @512 ranks (after 10 training steps), 1432 s @256 (fresh, instrumented: +0.0 MB NPU device memory), 2812 entries both isolated bench @512 (no model, same plan size): NPU device memory during save 624.6 MB → 0 (≈ world_size × plan size × 2); save time −18%; metadata exchange 5.6× faster

Root cause

Consistent with: torch_npu object collectives stage the pickled SavePlan as device tensors, and the D2H synchronization inside _gather_object stalls on large plans at high world sizes. (Direct attribution between the HCCL all-gather and the D2H sync link is not yet isolated.)

Conclusion

Device-backend object collectives are a reliability and memory hazard for DCP save at scale; a CPU gloo group is stable and cheaper — matching PyTorch's own hybrid-backend support (cpu:gloo,cuda:nccl).

Proposal

In engine/backends/fsdp2/checkpoint/dcp.py, pass process_group= to the two dcp.save(...) calls — in save_model (~L148) and save_optimizer (~L227). Add or reuse a gloo group created at first save in engine/checkpoint/distributed.py (env-gated for rollout). No caller changes. Full logs and repro scripts available on request.