perf(fsdp): SHARD_GRAD_OP for inference (-13.7%, -1855ms, bit-identical) - #52
Open
prashant182 wants to merge 3 commits into
Open
perf(fsdp): SHARD_GRAD_OP for inference (-13.7%, -1855ms, bit-identical)#52prashant182 wants to merge 3 commits into
prashant182 wants to merge 3 commits into
Conversation
perf(dit): eliminate ~5600 .item() syncs from DiT forward (-13% generate())
E8 of autoresearch — see /workspace/lingbot-world-artifacts/EXPERIMENTS.md. Switches DiT FSDP from FULL_SHARD to SHARD_GRAD_OP. Functionally equivalent to FSDP2's reshard_after_forward=False, which FSDP1 doesn't expose directly. Effect: parameters are unsharded after the first all-gather and stay resident across all 35 DiT forwards in a generate(), instead of being re-gathered per-forward. Profiling rationale: AllGather was 91% of NCCL time (975 ms / 2-chunk window) at the post-B3 baseline. The within-forward per-block gathers still fire, but the outer re-gather between forwards is gone — that's where the time went. Memory cost: unsharded 14B-param DiT at bf16 ~ 28GB per rank vs 3.5GB sharded across 8 ranks. On 80GB H100s with VAE+T5+KV cache + activations, we have headroom. Measurement (8×H100, 480×832, 81 frames, seed 42): generate() before 13523 ms (B3 baseline) generate() after 11668 ms Δ -1855 ms (-13.7%) MD5 ed2f82628308a3f8acd9b7935bb84401 (bit-identical, Tier A)
Contributor
Author
|
Nudge for review when you have a moment Single-line config flip in |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Switches DiT FSDP from
FULL_SHARDtoSHARD_GRAD_OP. Parameters stay resident across all 35 DiT forwards in agenerate()instead of being re-gathered per-forward. Bit-identical output (MD5ed2f82628308a3f8acd9b7935bb84401), −13.7% generate().SHARD_GRAD_OPis FSDP1's equivalent of FSDP2'sreshard_after_forward=False(the FSDP2 kwarg isn't accepted by the current PyTorch's FSDP1 constructor).Measurement (8×H100, 480×832, 81 frames, seed 42)
Stacked on top of #51 (B3) for the canonical chain:
generate()prewarm()(PR #48)Why
Profiling the post-B3 baseline showed AllGather was 91% of NCCL time — 975 ms over a 2-chunk window. With
FULL_SHARD, FSDP shards params after each forward and re-gathers on the next call. Across the 35 forwards in onegenerate(), that re-gather fires repeatedly.SHARD_GRAD_OPkeeps params unsharded after the first all-gather — the per-forward re-gather disappears.Memory
Unsharded 14B-param DiT at bf16 ≈ 28 GB per rank vs ~3.5 GB sharded across 8 ranks. On 80 GB H100s with VAE + T5 + KV cache + activations resident, there is headroom; no OOM observed across multiple bench runs.
Tradeoffs / scope
FULL_SHARDdefault would need to come back via a flag — happy to add one if requested.wan/distributed/fsdp.py); ~10 LOC including comment block.Stack
Stacked on PR #51 (B3). Reported diff will shrink to one line + comment once that lands.