[optim] run plan/metadata coordination over a gloo group

[yueming-yuan]

Problem

During torch_dist checkpoint load, torch DCP's plan/metadata coordination (_DistWrapper gather/scatter of LoadPlans, coordinator = global rank 0) runs on the default process group. Megatron initializes that group as NCCL-only, so these object collectives are staged through GPU tensors and go over NCCL.

On a 64-rank GLM-5.2 744B run (16×GB300, NVL72), this makes rank 0 open P2P channels to all 63 peers (~35 channels each, 10MB+2MB shareable buffers per channel):

• rank 0 NCCL allocations: 26.25 GB (4,562 allocs) vs ~3 GB on every other rank
• the buffers live on the world communicator and are never freed, so rank 0 permanently carries ~24 GB less usable GPU memory — including through the colocate rollout window, where it becomes the fleet-wide binding constraint on sglang mem_fraction_static

The payload transported is only pickled plans/shard descriptors (names, offsets, lengths, storage keys — a few hundred MB at rank 0); tensor data never crosses ranks in this path (each rank reads its own shards from disk).

Fix

Route the coordination through a lazily-created gloo group, for both the load and save planning paths. Tensor I/O is untouched; only the transport of plan objects changes (CPU/TCP instead of GPU/NCCL), so checkpoint bytes are identical.

Verification (A/B, identical 16×GB300 jobs, only this patch differs)

metric (rank-0 GPU, physical)	before (job 773)	after (job 778)
NCCL alloc total during ckpt load	26.25 GB (4,562 allocs)	2.62 GB (447 allocs)
actor at first update start	46.0 GiB	22.4 GiB
whole-GPU peak during update	188.7 GiB	164.8 GiB
actor residual during rollout	27.7 GiB	4.1 GiB
update_weights wall time	67.7–75.2 s	63.3 s (no regression)

Load correctness covered by the run itself: checkpoint loads, step-0 metrics normal, and the run has per-tensor SHA256 weight verification enabled (--check-rematerialize-param-from-master-weight).

🤖 Generated with Claude Code