Blog: Distributed Layerwise Offload — Running 185GB Models on 64GB NPU/GPU - #295
Blog: Distributed Layerwise Offload — Running 185GB Models on 64GB NPU/GPU#295evanchueng wants to merge 7 commits into
Conversation
…U/GPU Signed-off-by: evanchueng <evanchueng@outlook.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7dd02d838f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| | Model | Block Size | Blocks | HBM/card (DP2) | Host RAM/rank | | ||
| |-------|:---------:|:------:|:--------------:|:------------:| | ||
| | Cosmos3-Nano (33 GB) | 368 MB | 72 | ~10 GB | ~38 GB | | ||
| | Cosmos3-Super (124 GB) | 930 MB | 128 | ~15 GB | ~157 GB | |
There was a problem hiding this comment.
In DP2, this column is labeled per-rank, but the values shown are larger than the per-rank shard the post describes (1/dp_size plus shared mmap page cache): e.g. Cosmos3-Super shows ~157GB per rank for a 124GB model. Readers sizing a two-rank deployment would plan for ~314GB instead of the ~157GB total implied by the mmap/shared-cache design. Please either relabel this as total host/cgroup memory or replace it with true per-rank values.
Useful? React with 👍 / 👎.
| | Per-block wall time | ~23ms | 82s DiT / (72×50) | | ||
| | Estimated compute | ~9ms | Wall - prefetch | |
There was a problem hiding this comment.
With the preceding rows, Wall - prefetch is ~23ms - ~24ms, not ~9ms, so the estimated compute, MFU, and prefetch-bound conclusion are mathematically inconsistent. Anyone using this section to reason about H2D/AllGather overlap will get wrong bottleneck guidance unless the formula or measurements are corrected.
Useful? React with 👍 / 👎.
| @@ -0,0 +1,254 @@ | |||
| --- | |||
| layout: post | |||
| title: "Distributed Layerwise Offload: Running 185GB Models on 64GB NPU/GPU" | |||
There was a problem hiding this comment.
| title: "Distributed Layerwise Offload: Running 185GB Models on 64GB NPU/GPU" | |
| title: "Distributed Layerwise Offload: Serving 200B+ DiT models efficiently in vLLM-Omni" |
| --- | ||
| layout: post | ||
| title: "Distributed Layerwise Offload: Running 185GB Models on 64GB NPU/GPU" | ||
| author: "Evan Chueng" |
| |-------|:----:|:------:|:----------:|:--------:|:------------:| | ||
| | Cosmos3-Nano | 33 GB | 72 | 368 MB | ~10 GB | ~38 GB | | ||
| | Cosmos3-Super | 124 GB | 128 | 930 MB | ~15 GB | ~157 GB | | ||
| | Cosmos3-200B (synthetic) | 185 GB | 200 | 930 MB | ~15 GB | ~195 GB | |
There was a problem hiding this comment.
@david6666666 please add B300 experiments as well
| | T2V 832×480, 29帧 (~1s) | 5 | 1.9 MB | 61s | | ||
| | T2V 1280×720, 121帧 (~5s) | 5 | 16.9 MB | 394s | | ||
|
|
||
| The 200B model (185GB) runs on 2 × 64GB NPUs with DLO. **No other deployment method can achieve this** — the model is 2.9× larger than a single card's HBM, and even with 2 cards, HSDP would need 185GB / 128GB = 1.45× HBM utilization (no room for activations). |
There was a problem hiding this comment.
fp8? why a 200B model owns 185GB HBM memroy?
| | T2I 1024×1024 | 1 | 4.2 MB | 23s | | ||
| | T2I 1024×1024 | 5 | 4.2 MB | 58s | | ||
| | T2V 832×480, 29帧 (~1s) | 5 | 1.9 MB | 61s | | ||
| | T2V 1280×720, 121帧 (~5s) | 5 | 16.9 MB | 394s | |
|
|
||
| ### Overlap Analysis (Cosmos3-Nano, measured) | ||
|
|
||
| | Metric | Value | Source | |
There was a problem hiding this comment.
draw the overlap pipeline please, maybe with a SVG
There was a problem hiding this comment.
a gif would be better
| raise ValueError("DP multi-concurrency requires identical extra_args...") | ||
| ``` | ||
|
|
||
| ### RPC Wave ID for Stale Message Prevention |
|
|
||
| ## Performance | ||
|
|
||
| Tested on Ascend 910B3 (64GB HBM per card), DP2 + AllGather: |
There was a problem hiding this comment.
we also need to add the 30s, 4K resolution experiments
There was a problem hiding this comment.
why only DP2? we need a full experiment design in terms of hardware/model/serving workloads
…eview feedback - Title: 'Serving 200B+ DiT Models Efficiently in vLLM-Omni' (per hsliuustc0106) - Add David Zhao as co-author (per hsliuustc0106) - Full rewrite with ASCII diagrams, memory model, validation results - Add Cosmos3-200B synthetic model section with RFC-3 4K results - Clarify 200B = synthetic (100 blocks, 185GB BF16) - All English (per hsliuustc0106) - DP2/DP4 full experiment matrix (per hsliuustc0106) Signed-off-by: evanchueng <evanchueng@outlook.com>
Signed-off-by: evanchueng <evanchueng@outlook.com>
Signed-off-by: evanchueng <evanchueng@outlook.com>
- Add NVIDIA B300 SXM6 validation: T2I 1024x1024, T2V 480p/720p - Correctness verified via byte-identical SHA256 across all strategies - DLO+AG DP4: 1.39x throughput vs HSDP, 30% HBM usage - 720p 10s: within 2.13% of HSDP latency, 47% HBM - Remove synthetic 200B model section - Author: vLLM-Omni Diffusion Team Signed-off-by: evanchueng <evanchueng@outlook.com>
…llm-project#8/vllm-project#9) Replace ASCII art with clean inline SVG showing: - Three-stream timeline (Compute / H2D / AllGather) - Double-buffer slot alternation (Slot 0 / Slot 1) - Event sync arrows (compute waits for AllGather) - Color-coded legend Signed-off-by: evanchueng <evanchueng@outlook.com>
…m-project#8/vllm-project#9) - Generated 35-frame animated GIF (1500x675, 195KB) using matplotlib - Shows progressive reveal of three-stream timeline: Compute (blue) / H2D (orange) / AllGather (green) - Double-buffer slot alternation with event sync arrows - Replaces static inline SVG Signed-off-by: evanchueng <evanchueng@outlook.com>
|
|
||
| ## The Problem: Large Diffusion Models vs. HBM and Host Memory | ||
|
|
||
| Cosmos3-Super (64B parameters, 124 GB in BF16) cannot fit on a single 64 GB HBM device. The natural approach is to shard across multiple devices, but existing solutions each have limitations: |
There was a problem hiding this comment.
we have 2 alternative:
- offloader
- parallelism
please name both
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
I found four actionable documentation issues; inline comments below. I’m leaving the review as COMMENT and not selecting a merge verdict.
| @@ -0,0 +1,457 @@ | |||
| --- | |||
| layout: post | |||
| title: "Distributed Layerwise Offload: Serving 200B+ DiT Models Efficiently in vLLM-Omni" | |||
There was a problem hiding this comment.
[P1] Qualify the 200B+ headline. The measured sections cover Cosmos3-Nano (33 GB) and Cosmos3-Super (124 GB), while the 185/400 GB entries later in the post are explicitly estimates. Calling this “Serving 200B+” makes an extrapolation sound like a demonstrated deployment. Please change the title/summary to say “scaling toward” or “estimated,” or add reproducible 200B evidence.
| - **DP multi-concurrency**: Each DP rank processes a different request in parallel, achieving 3.3× throughput vs. single-request HSDP, with near-linear scaling. | ||
| - **Platform-agnostic**: Works on both NVIDIA GPU (CUDA/NCCL) and Ascend NPU (CANN/HCCL) via vLLM-Omni's platform abstraction layer. | ||
|
|
||
| Tested on Ascend 910B3 with Cosmos3-Nano (33 GB) and Cosmos3-Super (124 GB): all configurations produce correct video output, with cgroup-visible host memory scaling as O(model_size + dp_size × constant) instead of O(dp_size × model_size). |
There was a problem hiding this comment.
[P1] Scope the memory formula to the measured mode. “All configurations” is too broad: the implementation only uses mmap/sharded loading in DLO+AllGather; --dlo-no-use-allgather loads full weights independently, and GPU pin-memory/cgroup accounting differs from Ascend. Please scope this sentence to the measured DLO+AllGather/Ascend configuration and report other modes separately.
| │ cgroup peak: 47 GB │ | ||
| Total RSS: 4 × 33 = 132 GB │ (was 178 GB, -73%) │ | ||
| + page cache: 33 GB └────────────────────────┘ | ||
| = 178 GB ❌ |
There was a problem hiding this comment.
[P2] Fix the baseline arithmetic. 4 × 33 GB + 33 GB = 165 GB, not 178 GB. If 178 GB includes framework/transient overhead, include that missing ~13 GB term explicitly; otherwise this diagram conflicts with the validation table.
|
|
||
| On Ascend NPU, `pin_memory()` allocates DMA-capable memory via `/dev/davinci_manager` (the NPU device driver). This memory resides in CPU kernel space and is not tracked by cgroup — a key finding that explains why cgroup peak is much lower than expected. | ||
|
|
||
| **What you gain.** HBM holds only 2 layers of weights (~2 GB for Nano, ~3 GB for Super), regardless of model size. From Nano (33 GB) to Super (124 GB), idle HBM grows only 22% (11.5 → 14.6 GB) — the model is 3.8× larger but HBM barely changes. |
There was a problem hiding this comment.
[P2] Correct the HBM growth percentage. 11.5 → 14.6 GB is approximately a 27% increase relative to the Nano baseline, not 22%. Please correct this here and at line 373, or state the denominator/measurement definition.
david6666666
left a comment
david6666666
left a comment
There was a problem hiding this comment.
Reviewed the post against the merged DLO implementation (vllm-omni#5397) and RFC #5396. The core technical story checks out - meta+mmap loading, the fixed double-buffer, AllGather overlap, and DP multi-concurrency all match the code, and the memory-model section is the strongest part of the post. Inline comments below cover a few accuracy/clarity points (code attribution, metric definitions, mixed comparison bases) plus minor nits. Nothing blocking from a correctness standpoint.
| **Why it works.** We create the transformer on `torch.device("meta")`, which allocates zero memory — only tensor shapes and dtypes are recorded. We then load weights as mmap views using `safe_open().get_tensor()`, which returns a view into the OS page cache rather than a copy. | ||
|
|
||
| ```python | ||
| # pipeline_cosmos3.py — create transformer on meta device |
There was a problem hiding this comment.
[P2] Small attribution nit: the meta-device conversion doesn't live in pipeline_cosmos3.py. The merged PR (#5397) creates the transformer normally and the offloader converts it via to_empty(device="meta") in _load_weights_via_mmap (distributed_layerwise_backend.py), with download_weights_from_hf() in diffusers_loader.py. Worth updating this snippet (and the 'Cosmos3 meta device pipeline' reference below) so readers don't go hunting for a meta path in the pipeline file.
|
|
||
| ### Performance | ||
|
|
||
| | Strategy | Per-step (ms) | Throughput (fps) | CPU/rank | HBM/card | vs. HSDP | |
There was a problem hiding this comment.
[P2] Two small things: (1) the table doesn't state which model/resolution it's measuring - the HBM/card 20.3 GB matches the Cosmos3-Nano HSDP value from RFC #5396, so a label would help; (2) 'Throughput (fps)' is ambiguous: it reads like the output video's framerate (which is fixed at 29 fps), but it's actually generated video frames per wall-clock second (29 frames/output x concurrency / wave time). The B300 tables below use outputs/s - consider matching that or defining the metric inline.
| | 480p, ~5s (121f) | 102.58s | 15.88 GiB | 41.36s (125f) | 53.73–62.65 GiB | | ||
| | 480p, ~10s (241f) | 226.70s | 17.33 GiB | 82.47s (245f) | 53.74 GiB | | ||
| | 720p, 5s (121f) | 288.29s | 24.95 GiB | 87.47s | 52.19 GiB | | ||
| | 720p, 10s (241f) | 214.53s (DLO+AG+USP4) | 24.99 GiB | 210.05s | 53.73 GiB | |
There was a problem hiding this comment.
[P2] This row quietly switches strategy: the rows above compare DLO+AG DP4 (4 outputs) vs HSDP (1 output), but this one is DLO+AG+USP4 (1 output) sitting under the 'DLO+AG (4 outputs)' header. That's why the 2.13% is apples-to-apples while the other rows aren't. Marking concurrency per row (or splitting the two comparisons) would keep readers from mixing them up.
|
|
||
| | Strategy | Concurrency | Wave latency | Throughput | Process-tree PSS | Peak HBM/card | | ||
| |----------|:-----------:|:------------:|:----------:|:----------------:|:-------------:| | ||
| | DLO+AG DP4 | 4 | 43.69s (median) | 0.0915 outputs/s | 198–202 GiB | 12.62 GiB | |
There was a problem hiding this comment.
[P2] Worth a bridge sentence: ~200 GiB PSS (about 1.6-2x model size) looks like it contradicts the 'minimal host memory' narrative, but it's consistent - on CUDA the pinned shards are counted in process memory, while on Ascend they're hidden from cgroup via /dev/davinci_manager (which is why the Ascend cgroup numbers are ~47 GB). A one-line note that GPU RAM ~ page cache + pinned shards + framework would prevent readers from comparing the two platforms' numbers directly.
|
|
||
| A validation step rejects concurrent requests with different `num_inference_steps` — since AllGather is a collective, mismatched step counts would cause one rank to exit early while others hang. | ||
|
|
||
| **What you gain.** 4 concurrent requests achieve 3.22 fps throughput — 3.3× the HSDP single-request baseline. Scaling is near-linear (4.0×) because AllGather overhead is fixed (~150 ms/step) and amortized across 4 concurrent computations. |
There was a problem hiding this comment.
[Nit] Tiny wording nit: 3.22/0.967 = 3.33x, so 'near-linear (4.0x)' reads oddly - maybe '3.3x (~83% of the ideal 4x)' or just drop the (4.0x).
|
|
||
| - **Meta-device initialization + mmap weight loading**: Weights are loaded as mmap views pointing to shared OS page cache, eliminating O(dp_size × model_size) RSS during model creation. Cold-start peak RSS drops by 73% (178 GB → 47 GB for Cosmos3-Nano DP4). | ||
| - **Weight sharding + AllGather**: Each rank stores only 1/dp_size of the model. Full layer weights are reconstructed at runtime via AllGather, overlapped with computation on dedicated streams. | ||
| - **Fixed double-buffer scheme**: Exactly 2 layers of weights reside on each device at any time, regardless of model size — 22% HBM growth from 17B to 64B model. |
There was a problem hiding this comment.
[Nit] Same 22% figure appears here in the TL;DR. Also, section 3 mixes two number pairs: the text says idle HBM 11.5 -> 14.6 GB (~27%), while the diagram uses 23.1 -> 28.1 GB (~22%). Picking one pair and stating the denominator would make the number unambiguous.
| │ │ | ||
| │ Existing approaches: │ | ||
| │ │ | ||
| │ HSDP: 124/4 = 31 GB/card → 56 GB │ |
There was a problem hiding this comment.
[Nit] Minor ASCII-logic nit: '124/4 = 31 GB/card -> 56 GB' - the arrow reads like 31 becomes 56. I assume 56 GB is 31 GB weights + activations/comm buffers on the 64 GB card; a short annotation would make the OOM point land better.
3 participants
Blog: Distributed Layerwise Offload — Running 185GB Models on 64GB NPU/GPU