trunk/6911c3f96e75fee5a8b5e394ebecf1b52ad4a073: [FSDP2] Avoid division-by-one allocation for singleton shards (#196348)
- PyTorch: 705 events in the last 90 days
- PyTorch: 696th Release in the last 90 days
- Previous: earlier the same day · trunk/76872935231bf7ad36cd1923a68c58572bc57c36
What happened
Since we counter OOM in hsdp, with replication degree = 2 and sharding degree = 1. Therefore we are trying to fix the bug @ppl-ai AI-assisted description: We encountered CUDA OOM with DP2: drep=2, FSDP=1, CP=1 , using a (d_rep=2, d_fsdp_cp=1) mesh and FP32 gradient reduction. Although the shard group has size one, foreach_reduce computes input / 1 before copying to the output. In our run, this requested an additiona…
Summary assembled by rule from the sources below