diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 32aaf93bf..542c29fa1 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -353,6 +353,15 @@ block's shard is copied to a device buffer and reconstructed with `all_gather_into_tensor` on a communication stream while the current block is executing. +![AllGather DLO sharded host weights and lockstep reconstruction](../../figures/dlo/dlo_allgather_sharded_host.svg) + +The figure uses DP2 with TP2 to expose both parallel dimensions. DLO creates +one AllGather group across the DP replicas for each TP coordinate; TP0 workers +reconstruct TP0 blocks together and TP1 workers reconstruct TP1 blocks +together. Ordinary TP collectives remain inside each DP replica. Unlike the +no-AllGather runtime-cache path, these persistent pinned host shards are +process-owned, and every participating rank must follow the same block order. + ```text Compute: [Block N] [Block N+1] [Block N+2] H2D: [shard N+1] [shard N+2] diff --git a/docs/design/figures/dlo/dlo_allgather_sharded_host.svg b/docs/design/figures/dlo/dlo_allgather_sharded_host.svg new file mode 100644 index 000000000..2478786ed --- /dev/null +++ b/docs/design/figures/dlo/dlo_allgather_sharded_host.svg @@ -0,0 +1,165 @@ + + AllGather DLO sharded host weights and lockstep device reconstruction + An external orchestrator submits one synchronized wave to a DP2 by TP2 diffusion engine. Each DP replica contains TP0 and TP1 workers with private pinned host shards and two rotating GPU weight slots. DLO AllGather groups connect matching TP coordinates across DP replicas, reconstructing full rank-local blocks on each device in lockstep. Ordinary TP groups remain within each DP replica. + + + + + + + + + + + + + + + + AllGather DLO: sharded host weights, lockstep reconstruction + DP2 × TP2 shown: DLO groups span DP replicas per TP coordinate; ordinary TP stays within each replica + + + + External router / orchestrator + submits one synchronized DP wave · outside vLLM-Omni DLO + + + ONE NODE · ONE SYNCHRONIZED ENGINE, PROCESS-OWNED HOST SHARDS + + + request wave + + + + One synchronized diffusion engine · DP2 × TP2 + two DP requests follow the same block order; DLO reuses existing DP groups at each TP coordinate + + + + DP replica 0 · request 0 + ordinary TP group + + + TP0 · GPU 0 + DLO backend · local shard H2D + + SLOT 0 · CURRENT + full TP-local N · compute + + SLOT 1 · PREFETCH + full TP-local N+1 after AG + + + TP1 · GPU 1 + DLO backend · local shard H2D + + SLOT 0 · CURRENT + full TP-local N · compute + + SLOT 1 · PREFETCH + full TP-local N+1 after AG + + + + DP replica 1 · request 1 + ordinary TP group + + + TP0 · GPU 2 + DLO backend · local shard H2D + + SLOT 0 · CURRENT + full TP-local N · compute + + SLOT 1 · PREFETCH + full TP-local N+1 after AG + + + TP1 · GPU 3 + DLO backend · local shard H2D + + SLOT 0 · CURRENT + full TP-local N · compute + + SLOT 1 · PREFETCH + full TP-local N+1 after AG + + + DLO AG · TP0 GROUP + + DLO AG · TP1 GROUP + + + + + + ⇄ + DLO AllGather requires lockstep block order + every participant enters each block collective in the same sequence + + + + local shard H2D ×2 + + local shard H2D ×2 + + + + GROUP TOPOLOGY + DP2 × TP2 + DLO groups: same TP coordinate + across DP replicas + TP stays the model-parallel group. + + + PROCESS-OWNED PINNED HOST SHARDS · NOT RUNTIME-CACHE MMAP + + DP replica 0 · private + TP0 shard 0/2 · TP1 shard 0/2 + + DP replica 1 · private + TP0 shard 1/2 · TP1 shard 1/2 + Across each DLO group, shards total approximately one rank-local runtime block copy + + + HOT PATH · PER WORKER + Prefetch next block + local pinned shard → GPU + DLO AllGather → full block + rotate slots; overlap compute +