pytorch / pytorch/pytorch

[Feature Request][DCP] Add "Broadcast Load" arg to dcp.load for HSDP storage optimization

Open
#171,412 5 comments 3 reactions 0 assignees View on GitHub
bot-triaged feature oncall: distributed oncall: distributed checkpointing ptd-bot-triaged triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🚀 The feature, motivation and pitch

In large-scale HSDP (Hybrid Sharded Data Parallel) training, **storage bandwidth saturation** is a critical bottleneck during checkpoint loading.

Current `dcp.load` behavior follows a SPMD pattern where **every rank reads from storage**.
In a HSDP setup with $N$ nodes (Replication Factor = $N$), the storage bandwidth usage scales linearly ($N \times \text{Model Size}$). Since all replicas in a HSDP group require identical data, this redundant I/O traffic is inefficient.

On clusters with shared file systems (e.g., Lustre, Ceph, NFS), this often leads to severe **I/O throttling**, causing job startup times to spike significantly as the cluster scales.

### 💡 Proposed Solution
I propose adding a new argument `experimental_broadcast_replication` (bool) directly to `torch.distributed.checkpoint.load`.

**Proposed Behavior:**
1. **Identify Topology**: Infer the **Replicate Group** from the `DTensor` mesh.
2. **Leader Read**: Only **Rank 0** (Leader) of each Replicate Group executes the storage read. Non-leader ranks skip the read operation for replicated tensors.
3. **Broadcast**: The Leader broadcasts the loaded tensors to other ranks in the same group via high-speed interconnect (NVLink/IB).
4. **CPU Fallback**: For small CPU tensors (e.g., optimizer steps, scalars), maintain the current "all-read" behavior to avoid unnecessary host-device synchronization overhead.

### 🛠 Implementation Plan
I have implemented a prototype by modifying dcp.load and _load_state_dict to accept this new argument, and updated the loading routine to filter the execution plan and perform broadcasts for replicated tensors.

### Alternatives

_No response_

### Additional context

_No response_

cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @aditvenk @weifengpy @LucasLLC @pradeepfn @H-Huang

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.