anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Diffusion] Add loader-owned host-weight plans for DLO (TP=1) (#6213)

Browse Source 
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu authored and GitHub committed 2026-08-16 11:02:38 +08:00
1 parent 3875f8ded4
commit 9d2bb23ff6
17 files changed
+2009 -615
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/design/feature/offloader/distributed_layerwise_offload.md
+150 -27
View File
@@ -4,7 +4,7 @@ This document describes distributed layerwise offload (DLO) for diffusion
models. DLO keeps only a small number of DiT blocks on the accelerator and
streams the remaining blocks from host memory. The distributed backend can
either shard those host-side weights across an existing parallel group or keep
the standard loader's rank-local weights and avoid an additional collective.
complete rank-local block sources and avoid an additional collective.
For user-facing commands, see the
[distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
@@ -14,10 +14,20 @@ and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/rec
DLO is implemented for multi-device diffusion execution. The default
AllGather path is the primary path for DP and SP deployments. The
`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
for standard-loader sharding, workstation bring-up, and systems where an
additional DLO collective is undesirable, but it does not reduce host weight
storage across ranks.
`--dlo-no-use-allgather` path streams complete blocks independently and adds no
DLO weight collective.
Host storage is selected separately from the transfer protocol. The loader can
produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on
the same node can share immutable checkpoint pages when direct mmap is
selected, while the ordinary-loader fallback still keeps a private runtime
copy per process.
The Phase A shared-mmap support boundary is TP1. TP greater than one is an
ordinary-loader compatibility path: DLO can consume the resulting TP-local
tensors, but those configurations do not use checkpoint mmap and must not be
used to claim shared-mmap host-memory savings.
The compatibility matrix below describes the current implementation. The
unit-level guards are covered, but not every parallelism combination has a
@@ -41,6 +51,36 @@ TP is deliberately not used as DLO's AllGather group. HSDP has its own
parameter-sharding lifecycle and is not allowed to be sharded a second time by
DLO's AllGather path.
### The loader owns host-weight planning
Before it decides whether ordinary weight materialization can be skipped, the
diffusion loader builds one `HostWeightPlan`. A direct-checkpoint mmap plan is
accepted only when preflight proves all of the following:
- every required DiT parameter and persistent buffer has exactly one source;
- runtime names, checkpoint keys, shapes, and dtypes match;
- the runtime topology is TP1 without HSDP or online quantization; and
- every custom loader operation is represented by a loader-owned checkpoint
adapter.
The exact plan object is handed to DLO. The backend does not rescan checkpoint
files, repeat the capability decision, or reconstruct names from its block
topology. If preflight fails, the loader materializes weights normally and DLO
consumes those runtime tensors.
The plan owns only dedicated DiT component sources. If a pipeline also exposes
ordinary sources for a text encoder or another non-DiT component, the loader
still consumes those sources and includes their loaded names in its strict
coverage check. Only the source prefixes covered by the plan skip ordinary
materialization. A source that mixes DiT and non-DiT weights fails closed to
the complete ordinary-loader path because it cannot be skipped safely as a
unit.
This boundary keeps checkpoint semantics out of DLO and avoids model-pipeline
flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
transforms. Model-specific direct-layout knowledge, when required, lives in a
checkpoint adapter beside the ordinary loader.
### AllGather path
With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -61,6 +101,15 @@ Buffers: [current slot] [prefetch slot] [current slot]
The backend uses two shared device buffers, so accelerator weight residency is
bounded by the largest streamed blocks rather than the complete model.
When direct checkpoint mmap is selected, the checkpoint mappings are only the
source used to prepare each rank's persistent shard. They can be closed after
shard preparation. Across the AllGather group, those private shards total
approximately one runtime model copy.
An effective DLO group size of one performs no collective, even when
`dlo_use_allgather=True`; it follows the rank-local transfer path described
below.
When DP is greater than one, the engine can process one request per DP rank in
the same denoising wave. Because AllGather is a collective, all participating
requests must take the same execution path at every denoising step.
@@ -68,9 +117,20 @@ requests must take the same execution path at every denoising step.
### Rank-local path without DLO AllGather
With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
one. The regular model loader remains responsible for preparing each rank's
weights, including TP-local tensors or HSDP-managed parameters. DLO then
streams those rank-local tensors block by block using H2D copies only.
one and streams complete blocks using H2D copies only. The host backing may be
either a loader-approved checkpoint mapping or ordinary runtime tensors.
For direct mmap, each process retains immutable safetensors views and uses two
bounded pinned host staging slots. Processes on the same node that map the same
files share physical checkpoint pages through the OS page cache. This removes
the persistent private full-model copy per pure-DP process, but each process
still packs and transfers every complete block. Sharing is node-local; each
node has its own page cache.
When direct mmap preflight fails, the regular model loader remains responsible
for preparing each rank's weights, including TP-local tensors or HSDP-managed
parameters. In that fallback, each pure-DP process keeps a private full runtime
copy.
This mode means:
@@ -80,21 +140,18 @@ This mode means:
not shard weights across SP ranks.
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
DLO's additional weight AllGather is disabled.
- Pure DP deployments keep a full host-side model copy per rank, subject to
shared OS page-cache behavior.
- Pure DP deployments share one checkpoint-backed copy per node when direct
mmap is selected; the ordinary-loader fallback keeps one private runtime
copy per rank.
- The scheduler does not require a synchronized DP request wave for DLO.
This path is intentionally not implemented by reusing the AllGather mmap
loader. It relies on the standard loader so model-specific TP/HSDP transforms
remain intact.
## Parallelism compatibility
| Parallelism | DLO + AllGather | DLO without AllGather |
|---|---|---|
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
| **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. |
| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
### Combined dimensions
@@ -122,11 +179,15 @@ AllGather DP multi-concurrency requires:
The no-AllGather path does not impose these DLO-specific synchronized-wave
requirements.
The mmap loader is used only by the supported DLO+AllGather path when the
model has a compatible checkpoint layout. Online quantization is incompatible
with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
quantization. Models that do not meet the mmap requirements use the regular
loader path.
Direct checkpoint mmap can back either transfer path. It is currently limited
to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
ordinary loader. Online quantization remains incompatible with DLO AllGather;
use `--dlo-no-use-allgather` or disable online quantization.
A normalized runtime mmap cache, built through the ordinary loader, is the
proposed general mechanism for sharing transformed TP or quantized layouts.
That cache and its publication/lifecycle protocol are intentionally outside
this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
## Validation coverage
@@ -134,15 +195,75 @@ Current source-level validation includes:
- HSDP + DLO + AllGather rejection;
- HSDP + DLO without AllGather acceptance at configuration level;
- TP rejection in the DLO+AllGather mmap path;
- loader preflight fallback for TP, HSDP, online quantization, unknown custom
loaders, missing keys, and shape/dtype mismatches;
- exact loader-to-backend plan transfer and ordinary-loader fallback;
- rank-local mmap source retention, bounded two-slot staging, and adapter
transforms without parameter-side flags;
- resident-layer requests requiring no-AllGather;
- DP request-wave validation for denoising-step compatibility;
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
tests.
The highest-value missing coverage is end-to-end numerical comparison against
ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
on the target CUDA/NCCL or CANN/HCCL hardware.
### B300 parallel-topology smoke matrix
A four-GPU B300 smoke test covered MiniMax-H3 FL2VA with the same prompt, seed,
CUDNN attention backend, 256x256 output, two denoising steps, and
`dlo_resident_layers=0`. The TP2 rows used DiT DP2xTP2 with the text encoder and
VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap
or shared-mmap host-memory savings.
| Configuration | Result | Warm E2E | Peak device memory | Host PSS |
|---|---:|---:|---:|---:|
| DP4xTP1 AllGather | Passed, 4 concurrent requests | 2.87 s / 4 requests | 13.84 GiB | 211.99 GiB |
| DP4xTP1 no-AllGather | Passed, 1 request | 15.02 s | 13.23 GiB | 187.77 GiB |
| DP2xTP2 AllGather | Passed, 2 concurrent requests | 4.16 s / 2 requests | 12.50 GiB | 211.97 GiB |
| DP2xTP2 no-AllGather | Passed, 1 request | 3.51 s | 11.88 GiB | 314.01 GiB |
Within each topology, the AllGather and no-AllGather video and audio outputs
were byte-identical. All four runs completed without an `ERROR` or traceback
and released their device allocations. For DP4xTP1, no-AllGather direct mmap
reduced total PSS by 24.22 GiB (11.4%) and `Private_Dirty` from 211.33 to
125.32 GiB (40.7%) relative to AllGather. For DP2xTP2, preflight selected the
ordinary loader as designed; no-AllGather PSS was 314.01 GiB, about 48% above
AllGather, because DP replicas did not share checkpoint-backed runtime
weights. This is a functional and memory smoke test, not a production-quality
performance or output-quality benchmark.
### Host-memory measurement
A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the
ordinary-loader fallback with direct mmap. Both runs used
DP=2, TP=1, no DLO AllGather, BF16 weights, two denoising steps, and a
256x256 four-second request. The ordinary-loader workers were sampled after
initialization. The mmap workers were sampled after one completed request, so
the checkpoint working set had been faulted into the page cache; this is the
more conservative point for mmap.
The values below come from `/proc/<worker>/smaps_rollup` and include the whole
worker, not only the DiT. The stable rank-to-rank difference comes from other
pipeline components, so each worker should be compared with the same worker in
the other storage mode.
| Worker | Ordinary RSS | mmap RSS | Ordinary PSS | mmap PSS | PSS reduction |
|---|---:|---:|---:|---:|---:|
| DP worker 0 | 168.27 GiB | 132.76 GiB | 167.84 GiB | 101.43 GiB | 66.40 GiB |
| DP worker 1 | 116.19 GiB | 79.97 GiB | 115.73 GiB | 48.64 GiB | 67.09 GiB |
| **Two-worker total** | — | — | **283.56 GiB** | **150.08 GiB** | **133.48 GiB (47.1%)** |
The direct-mmap workers each reported 62.45 GiB `Shared_Clean` but only
31.20 GiB `Pss_File`, which is the proportional charge expected when the same
resident checkpoint pages are mapped by two workers. `Private_Dirty` also fell
from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
1, a reduction of about 97–98 GiB per worker. RSS understates this benefit
because it counts a shared physical page in every process that maps it; summed
PSS is the appropriate node-memory comparison.
The highest-value missing coverage is broader end-to-end numerical and
lifecycle comparison against ordinary layerwise offload for DP+SP,
HSDP+SP+no-AllGather, and TP greater than one across additional models and
target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not
change the Phase A direct-mmap TP1 support boundary.
## Recommendations
@@ -150,7 +271,9 @@ on the target CUDA/NCCL or CANN/HCCL hardware.
scaling path.
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
not the goal.
- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
the expectation of higher host-memory use and lower validation confidence.
- Use **no-AllGather** when independent replica execution is required. TP1
direct-mmap deployments can share checkpoint pages per node; other layouts
retain the ordinary loader's private host memory behavior and are outside
the Phase A shared-mmap support scope.
- Prefer **HSDP alone** for production HSDP deployments until the combined
HSDP + DLO no-AllGather path has broader end-to-end coverage.
docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md
+48 -13
View File
@@ -3,7 +3,9 @@
Distributed layerwise offloading (DLO) extends block streaming to multi-device
deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
the host weights and reconstructs each layer at runtime. Without AllGather,
each rank streams its standard-loader rank-local weights independently.
each rank streams a complete block independently. Compatible TP1 deployments
can share checkpoint-backed host pages among processes on the same node;
otherwise DLO streams the ordinary loader's rank-local tensors.
See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
for the implementation contract and compatibility matrix.
@@ -59,21 +61,48 @@ omni = Omni(
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
| `--dlo-no-use-allgather` | Keep standard-loader rank-local weights | `false` |
| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
## mmap weight loading
## Host-weight loading
The DLO plus AllGather path:
The diffusion loader chooses host storage before DLO is enabled. It first
attempts to build a complete, validated direct-checkpoint mmap plan. If names,
coverage, shape, dtype, topology, or loader-callback compatibility cannot be
proven, it runs the ordinary model loader instead. DLO consumes that result and
does not make a second checkpoint-compatibility decision.
The shared-mmap optimization in this phase is supported only with TP1. TP
greater than one falls back before model mutation to ordinary TP-aware loading.
DLO may still consume those TP-local tensors, but this is a compatibility path:
it does not share checkpoint-backed runtime weights across DP replicas and
provides no shared-mmap host-memory guarantee.
The mmap plan skips only dedicated DiT weight sources. Other component sources,
such as a text encoder loaded through the shared diffusion loader, continue to
use their ordinary component loader. A checkpoint source that mixes DiT and
non-DiT weights falls back completely rather than leaving an unplanned
component uninitialized.
With direct checkpoint mmap, the loader:
1. saves non-persistent buffers such as RoPE frequencies;
2. moves the normally created transformer to the meta device;
3. loads checkpoint tensors as mmap views backed by the shared OS page cache;
4. calls model-specific `post_load_weights()` conversions; and
5. restores the saved non-persistent buffers.
4. applies any loader-owned bounded layout adapters while packing blocks;
5. restores saved non-persistent buffers; and
6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle
hooks.
This avoids one full checkpoint RSS copy per rank and does not require
model-specific loading code.
For AllGather with a group larger than one, each process copies only its
persistent shard and then releases the source mapping. For no-AllGather, each
process keeps the mapping open and packs complete blocks through two bounded
pinned staging slots. Processes mapping the same files on one node share the
immutable pages; no-AllGather still performs a complete-block H2D copy in each
process.
When the effective DLO group size is one, `dlo_use_allgather=True` does not
perform a collective and uses the same rank-local transfer behavior.
## Declarative topology
@@ -103,16 +132,22 @@ must enter each collective.
## Limitations
- Online FP8 quantization is rejected with the DLO plus AllGather mmap path.
Use `--dlo-no-use-allgather` or disable online quantization.
- Tensor parallel size greater than one is rejected in the mmap path because
it bypasses TP-aware loader callbacks. The no-AllGather path retains the
standard loader and remains experimental with TP.
- Direct checkpoint mmap currently requires TP1. TP greater than one is
outside the Phase A shared-mmap support scope and falls back before model
mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
but it provides no shared-mmap host-memory benefit or guarantee.
- HSDP plus AllGather is rejected to avoid double sharding. HSDP without
AllGather has limited end-to-end validation.
- Online quantization uses the ordinary loader with no-AllGather. It remains
incompatible with DLO AllGather.
- Resident leading layers require `--dlo-no-use-allgather` and a model
`OffloadPlan` that declares eligible `resident_dit_paths`.
- DP concurrency requires an explicit, identical inference-step count.
Sharing transformed TP or quantized runtime layouts through a normalized mmap
cache is a follow-up design in
[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
the direct-checkpoint path.
See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
for an end-to-end example.
examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py → examples/offline_inference/minimax_h3/dlo_lifecycle.py
+60 -25
View File
@@ -1,21 +1,32 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.
"""Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
The DLO modes accept any positive data-parallel size. With AllGather enabled,
the script submits one complete DP wave and verifies recovery after an invalid
wave. It also checks that every mode shuts down all worker processes cleanly.
Examples:
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode dlo --dp-size 4
# Rank-local mmap storage, with no DLO collective or DP wave scheduling.
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
CUDA_VISIBLE_DEVICES=0,1 \
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode dlo-no-allgather --dp-size 2
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
CUDA_VISIBLE_DEVICES=0,1 \
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode request
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode request --tp-size 2
"""
from __future__ import annotations
@@ -40,14 +51,33 @@ DEFAULT_PROMPTS = (
)
def positive_int(value: str) -> int:
parsed = int(value)
if parsed < 1:
raise argparse.ArgumentTypeError("must be a positive integer")
return parsed
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
parser.add_argument(
"--mode",
choices=("dlo-dp2", "request"),
choices=("dlo", "dlo-no-allgather", "request"),
required=True,
help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
help="DLO (with or without AllGather) or ordinary non-DLO request mode",
)
parser.add_argument(
"--dp-size",
type=positive_int,
default=2,
help="DLO data-parallel size (for example 1, 2, 4, or 8; default: 2)",
)
parser.add_argument(
"--tp-size",
type=positive_int,
default=2,
help="Tensor-parallel size used only by request mode (default: 2)",
)
parser.add_argument("--steps", type=int, default=2)
parser.add_argument("--duration", type=float, default=5.0)
@@ -60,10 +90,11 @@ def parse_args() -> argparse.Namespace:
def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
common: dict[str, Any] = {
"model": args.model,
"trust_remote_code": True,
"num_gpus": 2,
"num_gpus": args.dp_size if is_dlo else args.tp_size,
"usp": 1,
"ring": 1,
"vae_parallel_mode": "tile",
@@ -74,22 +105,22 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
"stage_init_timeout": args.init_timeout,
"init_timeout": args.init_timeout,
}
if args.mode == "dlo-dp2":
if is_dlo:
common.update(
tensor_parallel_size=1,
data_parallel_size=2,
data_parallel_size=args.dp_size,
text_encoder_tp_size=1,
vae_patch_parallel_size=1,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
dlo_use_allgather=args.mode == "dlo",
dlo_resident_layers=0,
)
else:
common.update(
tensor_parallel_size=2,
tensor_parallel_size=args.tp_size,
data_parallel_size=1,
text_encoder_tp_size=2,
vae_patch_parallel_size=2,
text_encoder_tp_size=args.tp_size,
vae_patch_parallel_size=args.tp_size,
enable_distributed_layerwise_offload=False,
)
return common
@@ -172,7 +203,8 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
"engine_kwargs": engine_kwargs(args),
}
try:
if args.mode == "dlo-dp2":
uses_dp_wave = args.mode == "dlo" and args.dp_size > 1
if uses_dp_wave:
started = time.perf_counter()
asymmetric = await asyncio.gather(
generate_one(
@@ -182,12 +214,15 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
prompt="",
seed=1001,
),
generate_one(
engine,
args,
request_id="invalid-peer",
prompt=DEFAULT_PROMPTS[0],
seed=1002,
*(
generate_one(
engine,
args,
request_id=f"invalid-peer-{index}",
prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
seed=1001 + index,
)
for index in range(1, args.dp_size)
),
return_exceptions=True,
)
@@ -195,18 +230,18 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
summary["asymmetric_errors"] = [
f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException)
]
if len(summary["asymmetric_errors"]) != 2:
raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")
if len(summary["asymmetric_errors"]) != args.dp_size:
raise RuntimeError("Every request in the asymmetric DP wave must fail before dispatch")
started = time.perf_counter()
request_count = 2 if args.mode == "dlo-dp2" else 1
request_count = args.dp_size if uses_dp_wave else 1
outputs = await asyncio.gather(
*(
generate_one(
engine,
args,
request_id=f"recovery-{index}",
prompt=DEFAULT_PROMPTS[index],
prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
seed=2000 + index,
)
for index in range(request_count)
tests/diffusion/model_loader/test_diffusers_loader.py
+136
View File
@@ -16,6 +16,11 @@ from vllm.config.load import LoadConfig
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.host_weight_plan import (
HostWeightPlan,
HostWeightPlanResult,
TensorBinding,
)
from vllm_omni.diffusion.models.helios import HeliosPipeline
from vllm_omni.diffusion.registry import initialize_model
@@ -258,6 +263,137 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
assert get_current_diffusion_config_or_none() is None
def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
od_config = SimpleNamespace(
dtype=torch.float32,
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
plan = HostWeightPlan(
backing_kind="checkpoint_mmap",
bindings={},
)
loaded_ordinary_weights = False
def load_weights(_model):
nonlocal loaded_ordinary_weights
loaded_ordinary_weights = True
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader.load_weights = load_weights # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(plan),
)
assert loader.load_model(load_device="cpu") is model
assert not loaded_ordinary_weights
assert loader.take_host_weight_plan() is plan
assert loader.take_host_weight_plan() is None
def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
od_config = SimpleNamespace(
dtype=torch.float32,
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
class MixedSourceModel(nn.Module):
def __init__(self):
super().__init__()
self.transformer = nn.Linear(2, 2, bias=False)
self.text_encoder = nn.Linear(2, 2, bias=False)
self.weights_sources = (
DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="transformer."),
DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="text_encoder."),
)
self.loaded_weight_names: list[str] = []
def load_weights(self, weights):
self.loaded_weight_names = [name for name, _ in weights]
return set(self.loaded_weight_names)
model = MixedSourceModel()
plan = HostWeightPlan(
backing_kind="checkpoint_mmap",
bindings={
"transformer.weight": TensorBinding(
checkpoint_key="weight",
file_path="unused",
)
},
planned_source_prefixes=frozenset({"transformer."}),
)
requested_prefixes: list[str] = []
def get_weights(source, model=None):
del model
requested_prefixes.append(source.prefix)
yield source.prefix + "weight", torch.ones(2, 2)
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader._get_weights_iterator = get_weights # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(plan),
)
assert loader.load_model(load_device="cpu") is model
assert requested_prefixes == ["text_encoder."]
assert model.loaded_weight_names == ["text_encoder.weight"]
assert loader.take_host_weight_plan() is plan
def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
od_config = SimpleNamespace(
dtype=torch.float32,
parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
calls: list[str] = []
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader.load_weights = lambda _model: calls.append("load") # type: ignore[method-assign]
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"),
)
assert loader.load_model(load_device="cpu") is model
assert calls == ["load", "process"]
assert loader.take_host_weight_plan() is None
def test_hsdp_processes_quantized_weights_before_sharding(mocker):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
from vllm_omni.diffusion.offloader.module_collector import PipelineModules
tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py
+43
View File
@@ -191,6 +191,49 @@ def test_model_load_weights_transforms_before_calling_vllm_loader():
]
def test_loader_adapter_declares_equivalent_direct_mmap_transform(monkeypatch):
from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
get_direct_mmap_adapter,
)
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear)
monkeypatch.setattr(h3, "Attention", _FakeAttention)
arch = h3.MiniMaxH3DiTArchConfig(
hidden_size=1,
num_attention_heads=2,
attention_head_dim=1,
rope_inv_freq_len=1,
)
attention = h3.MiniMaxH3Attention(
arch,
quant_config=None,
prefix="blocks.0.attn",
)
transformer = object.__new__(h3.MiniMaxH3DiTModel)
nn.Module.__init__(transformer)
transformer.blocks = nn.ModuleList([nn.Module()])
transformer.blocks[0].attn = attention
pipeline = nn.Module()
pipeline.transformer = transformer
adapter = get_direct_mmap_adapter(pipeline)
assert adapter is not None
policy = adapter.policy_for(
"transformer.blocks.0.attn.qkv_proj.weight",
attention.qkv_proj.weight,
)
assert policy is not None
assert policy.allow_custom_loader
assert policy.transform is not None
checkpoint_weight = torch.arange(6, dtype=torch.float32).reshape(6, 1)
assert policy.transform(checkpoint_weight)[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
assert not hasattr(attention.qkv_proj.weight, "mmap_weight_transform")
def test_pipeline_resolves_transformer_component_quant_config():
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_resolve_component_quant_config,
tests/diffusion/offloader/test_distributed_layerwise_backend.py
+515 -80
View File
@@ -18,6 +18,11 @@ from torch import nn
from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
from vllm_omni.diffusion.model_loader.host_weight_plan import (
HostWeightPlan,
TensorBinding,
build_checkpoint_mmap_plan,
)
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
from vllm_omni.diffusion.offloader.block_discovery import (
get_blocks_attr_names,
@@ -30,7 +35,10 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
PinnedResidentLayerGroup,
)
from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan
from vllm_omni.diffusion.offloader.offload_plan import (
OffloadPlan,
get_offload_plan,
)
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -48,6 +56,9 @@ class DummyEvent:
def record(self, _stream) -> None:
return None
def synchronize(self) -> None:
return None
@contextmanager
def dummy_stream(_stream):
@@ -351,8 +362,9 @@ class TestDistributedLayerwiseOffloadHook:
block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
next_block = nn.Module()
next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
next_block.weight.mmap_weight_transform_pending = True
def transform(tensor):
return tensor.flip(0)
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
@@ -361,17 +373,17 @@ class TestDistributedLayerwiseOffloadHook:
dp_size=2,
rank=0,
pin_memory=False,
tensor_transforms={id(next_block.weight): transform},
)
hook.initialize_hook(block)
assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([3.0, 2.0]))
def test_regular_loader_weight_does_not_reapply_mmap_transform(self, dist_group, patched_offload_runtime):
def test_transform_is_not_applied_without_a_loader_plan(self, dist_group, patched_offload_runtime):
block = nn.Module()
block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
next_block = nn.Module()
next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
@@ -385,6 +397,82 @@ class TestDistributedLayerwiseOffloadHook:
assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0]))
def test_rank_local_mmap_retains_source_and_uses_staging(self, dist_group, patched_offload_runtime):
current_block = nn.Linear(2, 2, bias=False)
next_block = nn.Linear(2, 2, bias=False)
next_block.weight.data.copy_(torch.arange(4, dtype=torch.float32).view(2, 2))
source_storage = next_block.weight.untyped_storage().data_ptr()
def transform(tensor):
return tensor.t()
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
device=torch.device("cpu"),
dp_group=None,
dp_size=1,
rank=0,
pin_memory=False,
rank_local_mmap=True,
tensor_transforms={id(next_block.weight): transform},
)
hook.initialize_hook(current_block)
source = hook.cpu_sources[torch.float32][0]["tensor"]
assert source.untyped_storage().data_ptr() == source_storage
assert hook.cpu_shards == {}
assert next_block.weight.numel() == 0
hook.cpu_staging_buffers = [
{torch.float32: torch.empty(4)},
{torch.float32: torch.empty(4)},
]
hook.prefetch_layer(slot=0, non_blocking=False)
assert torch.equal(
next_block.weight,
torch.tensor([[0.0, 2.0], [1.0, 3.0]]),
)
assert next_block.weight.stride() == (1, 2)
assert torch.equal(source, torch.arange(4, dtype=torch.float32).view(2, 2))
def test_rank_local_mmap_staging_is_bounded_by_largest_block(self, patched_offload_runtime):
hooks = []
for size in (4, 9):
current_block = nn.Linear(1, 1, bias=False)
next_block = nn.Module()
next_block.weight = nn.Parameter(torch.arange(size, dtype=torch.float32))
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
device=torch.device("cpu"),
dp_group=None,
dp_size=1,
rank=0,
pin_memory=False,
shared_buffers=[None, None],
rank_local_mmap=True,
)
hook.initialize_hook(current_block)
hooks.append(hook)
resident_block = nn.Module()
resident_block.weight = nn.Parameter(torch.arange(12, dtype=torch.float32))
resident_group = PinnedResidentLayerGroup(
[resident_block],
device=torch.device("cpu"),
copy_stream=DummyStream(),
pin_memory=False,
rank_local_mmap=True,
)
staging = DistributedLayerwiseOffloadBackend._allocate_shared_cpu_staging_buffers(
hooks,
resident_group,
)
assert len(staging) == 2
assert staging[0][torch.float32].numel() == 12
assert staging[1][torch.float32].numel() == 12
class TestPinnedResidentLayerGroup:
def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
@@ -439,23 +527,41 @@ class TestPinnedResidentLayerGroup:
assert block.weight.stride() == expected_stride
group.offload()
def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self):
module = nn.Linear(2, 2, bias=False)
def test_rank_local_mmap_resident_layers_retain_sources(self, patched_offload_runtime):
blocks = [nn.Linear(2, 2, bias=False) for _ in range(3)]
expected = []
source_storage = []
for index, block in enumerate(blocks):
block.weight.data.fill_(index + 1)
expected.append(block.weight.detach().clone())
source_storage.append(block.weight.untyped_storage().data_ptr())
def transform(tensor):
return tensor.flip(0)
group = PinnedResidentLayerGroup(
blocks,
device=torch.device("cpu"),
copy_stream=DummyStream(),
pin_memory=False,
rank_local_mmap=True,
)
module.weight.mmap_weight_transform = transform
for index, state in enumerate(group._states):
source = state["cpu_sources"][torch.float32][0]["tensor"]
assert source.untyped_storage().data_ptr() == source_storage[index]
assert state["cpu_shards"] == {}
assert len(group._cpu_staging_buffers) == 2
backend = object.__new__(DistributedLayerwiseOffloadBackend)
backend._remember_mmap_param_attrs(module)
module.to_empty(device="meta")
group.load()
for block, weight in zip(blocks, expected):
assert torch.equal(block.weight, weight)
replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float())
backend._attach_mmap_param_attrs("weight", replacement, module.weight)
blocks[0].weight.data.zero_()
group.offload()
group.load()
for block, weight in zip(blocks, expected):
assert torch.equal(block.weight, weight)
assert replacement.mmap_weight_transform is transform
assert replacement.mmap_weight_transform_pending is True
group.offload()
assert all(block.weight.numel() == 0 for block in blocks)
def test_shared_allgather_output_is_narrowed_to_current_block(
self, dist_group, patched_offload_runtime, monkeypatch
@@ -598,7 +704,6 @@ class TestMmapWeightLoading:
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
model_path=str(tmp_path),
),
torch.device("cpu"),
)
@@ -606,13 +711,81 @@ class TestMmapWeightLoading:
dits=[pipeline.transformer],
dit_names=["transformer"],
)
plan = HostWeightPlan(
backing_kind="checkpoint_mmap",
bindings={
name: TensorBinding(
checkpoint_key=name,
file_path=str(weight_file),
)
for name in weights
},
)
backend._load_weights_via_mmap(pipeline, modules)
backend._load_weights_via_mmap(pipeline, modules, plan)
assert pipeline.transformer.post_load_calls == 1
assert pipeline.transformer.time_embedder.weight.dtype == torch.float32
assert pipeline.transformer.blocks[0].weight.dtype == torch.bfloat16
def test_allgather_flag_with_group_size_one_uses_rank_local_mmap(
self,
tmp_path,
patched_offload_runtime,
):
class Transformer(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
def __init__(self):
super().__init__()
self.blocks = nn.ModuleList(
[
nn.Linear(2, 2, bias=False),
nn.Linear(2, 2, bias=False),
]
)
class Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = Transformer()
@staticmethod
def _remap_ckpt_key(key):
return key
pipeline = Pipeline()
weights = {name: torch.ones_like(param) for name, param in pipeline.named_parameters()}
checkpoint_file = tmp_path / "model.safetensors"
save_file(weights, str(checkpoint_file))
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path=str(tmp_path),
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.plan is not None
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
dp_size=1,
dlo_use_allgather=True,
),
torch.device("cpu"),
host_weight_plan=result.plan,
)
backend.enable(pipeline)
assert backend._using_rank_local_mmap
assert backend.dp_group is None
assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
backend.disable()
class TestGetBlocksFromDit:
def test_get_blocks_from_dit_single_block_attr(self):
@@ -878,23 +1051,123 @@ class TestOffloadPlan:
class TestMmapValidation:
"""Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights."""
"""Tests for loader preflight and backend plan realization."""
def test_tp_rejected_when_tp_world_size_gt_1(self, tmp_path, patched_offload_runtime, monkeypatch):
"""DLO+AllGather should reject when TP world size > 1."""
import json
from types import SimpleNamespace
from safetensors.torch import save_file
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
# Mock TP world size = 2
monkeypatch.setattr(
"vllm.distributed.parallel_state.get_tensor_model_parallel_world_size",
lambda: 2,
def test_component_source_prefix_builds_model_namespace(self, tmp_path):
source_root = tmp_path / "partition"
checkpoint_dir = source_root / "transformer"
checkpoint_dir.mkdir(parents=True)
checkpoint_file = checkpoint_dir / "model.safetensors"
save_file(
{"blocks.0.weight": torch.arange(4, dtype=torch.float32).reshape(1, 4)},
str(checkpoint_file),
)
class Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = nn.Module()
self.transformer.blocks = nn.ModuleList([nn.Linear(4, 1, bias=False)])
self.weights_sources = [
SimpleNamespace(
model_or_path=str(source_root),
subfolder="transformer",
revision=None,
prefix="transformer.",
)
]
pipeline = Pipeline()
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=tuple(pipeline.weights_sources),
model_path=None,
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.fallback_reason is None
assert result.plan is not None
assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding(
checkpoint_key="blocks.0.weight",
file_path=str(checkpoint_file),
)
assert result.plan.planned_source_prefixes == frozenset({"transformer."})
def test_non_dedicated_component_source_falls_back_before_discovery(self):
pipeline = nn.Module()
pipeline.transformer = nn.Linear(2, 2, bias=False)
sources = (
SimpleNamespace(
model_or_path="unused",
subfolder=None,
revision=None,
prefix="",
),
)
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=sources,
model_path=None,
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.plan is None
assert "dedicated component weight source" in result.fallback_reason
def test_preflight_falls_back_before_mutation_for_missing_model_tensors(self, tmp_path):
source_root = tmp_path / "partition"
checkpoint_dir = source_root / "transformer"
checkpoint_dir.mkdir(parents=True)
save_file(
{"blocks.0.weight": torch.ones((2, 2), dtype=torch.float32)},
str(checkpoint_dir / "model.safetensors"),
)
class Transformer(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
def __init__(self):
super().__init__()
self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False), nn.Linear(2, 2, bias=False)])
class Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = Transformer()
self.weights_sources = [
SimpleNamespace(
model_or_path=str(source_root),
subfolder="transformer",
revision=None,
prefix="transformer.",
)
]
pipeline = Pipeline()
original_weight = pipeline.transformer.blocks[0].weight
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=tuple(pipeline.weights_sources),
model_path=None,
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.plan is None
assert "1 required DiT tensors" in result.fallback_reason
assert pipeline.transformer.blocks[0].weight is original_weight
assert not original_weight.is_meta
def test_tp_falls_back_before_checkpoint_discovery(self):
class SimpleModel(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
@@ -912,41 +1185,118 @@ class TestMmapValidation:
return key
pipeline = SimplePipeline()
weights = {name: torch.ones(p.shape, dtype=p.dtype) for name, p in pipeline.named_parameters() if not p.is_meta}
save_file(weights, str(tmp_path / "model.safetensors"))
(tmp_path / "model.safetensors.index.json").write_text(
json.dumps({"weight_map": {k: "model.safetensors" for k in weights}})
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path="unused",
tensor_parallel_size=2,
use_hsdp=False,
online_quantization=False,
)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
model_path=str(tmp_path),
),
torch.device("cpu"),
)
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
assert result.plan is None
assert result.fallback_reason == "TP=2 requires the ordinary loader"
with pytest.raises(ValueError, match="Tensor Parallel"):
backend._load_weights_via_mmap(pipeline, modules)
@pytest.mark.parametrize(
("use_hsdp", "online_quantization", "expected_reason"),
[
(True, False, "HSDP requires the ordinary loader"),
(False, True, "online quantization requires the ordinary loader"),
],
)
def test_runtime_layouts_that_need_loader_callbacks_fall_back(
self,
use_hsdp,
online_quantization,
expected_reason,
):
pipeline = nn.Module()
pipeline.transformer = nn.Linear(2, 2, bias=False)
def test_tp_allowed_when_tp_world_size_eq_1(self, tmp_path, patched_offload_runtime, monkeypatch):
"""DLO+AllGather should work when TP world size = 1, even if params
have custom weight_loader attributes (e.g. QKVParallelLinear at TP=1)."""
import json
from types import SimpleNamespace
from safetensors.torch import save_file
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
# Mock TP world size = 1 (default)
monkeypatch.setattr(
"vllm.distributed.parallel_state.get_tensor_model_parallel_world_size",
lambda: 1,
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path="unused",
tensor_parallel_size=1,
use_hsdp=use_hsdp,
online_quantization=online_quantization,
)
assert result.plan is None
assert result.fallback_reason == expected_reason
@pytest.mark.parametrize(
("checkpoint_tensor", "expected_reason"),
[
(torch.ones((3, 2), dtype=torch.float32), "shape mismatch"),
(torch.ones((2, 2), dtype=torch.bfloat16), "dtype mismatch"),
],
)
def test_source_metadata_mismatch_falls_back(self, tmp_path, checkpoint_tensor, expected_reason):
checkpoint_file = tmp_path / "model.safetensors"
save_file({"transformer.weight": checkpoint_tensor}, str(checkpoint_file))
class Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = nn.Linear(2, 2, bias=False)
@staticmethod
def _remap_ckpt_key(key):
return key
pipeline = Pipeline()
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path=str(tmp_path),
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.plan is None
assert expected_reason in result.fallback_reason
def test_non_persistent_buffers_do_not_need_checkpoint_bindings(self, tmp_path):
class Transformer(nn.Module):
def __init__(self):
super().__init__()
self.weight = nn.Parameter(torch.ones(2, 2))
self.register_buffer("derived", torch.ones(2), persistent=False)
checkpoint_file = tmp_path / "model.safetensors"
save_file({"transformer.weight": torch.ones(2, 2)}, str(checkpoint_file))
class Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = Transformer()
@staticmethod
def _remap_ckpt_key(key):
return key
pipeline = Pipeline()
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path=str(tmp_path),
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.plan is not None
assert set(result.plan.bindings) == {"transformer.weight"}
def test_unadapted_custom_weight_loader_falls_back_at_tp1(self, tmp_path):
def custom_weight_loader(param, weight):
param.data.copy_(weight)
@@ -976,27 +1326,67 @@ class TestMmapValidation:
json.dumps({"weight_map": {k: "model.safetensors" for k in weights}})
)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
model_path=str(tmp_path),
),
torch.device("cpu"),
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path=str(tmp_path),
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
# Should NOT raise — TP=1 is allowed even with custom weight_loader
backend._load_weights_via_mmap(pipeline, modules)
assert result.plan is None
assert "custom weight_loader" in result.fallback_reason
def test_cosmos3_adapter_preserves_existing_tp1_direct_mmap_contract(self, tmp_path):
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
Cosmos3OmniDiffusersPipeline,
)
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
Cosmos3VFMTransformer,
)
transformer = object.__new__(Cosmos3VFMTransformer)
nn.Module.__init__(transformer)
transformer.proj_in = nn.Linear(2, 2, bias=False)
transformer.proj_in.weight.weight_loader = lambda *_args: None
pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
nn.Module.__init__(pipeline)
pipeline.transformer = transformer
source_root = tmp_path / "partition"
checkpoint_dir = source_root / "transformer"
checkpoint_dir.mkdir(parents=True)
checkpoint_file = checkpoint_dir / "model.safetensors"
save_file({"proj_in.weight": torch.ones(2, 2)}, str(checkpoint_file))
sources = (
SimpleNamespace(
model_or_path=str(source_root),
subfolder="transformer",
revision=None,
prefix="transformer.",
),
)
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", transformer),),
sources=sources,
model_path=None,
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
assert result.plan is not None
assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
checkpoint_key="proj_in.weight",
file_path=str(checkpoint_file),
)
def test_validate_loaded_weights_called(self, tmp_path, patched_offload_runtime):
"""validate_loaded_weights should be called after mmap loading."""
import json
from types import SimpleNamespace
from safetensors.torch import save_file
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
class ModelWithValidate(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
@@ -1038,19 +1428,64 @@ class TestMmapValidation:
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False, model_path=str(tmp_path)
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
),
torch.device("cpu"),
)
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
result = build_checkpoint_mmap_plan(
pipeline,
dit_modules=(("transformer", pipeline.transformer),),
sources=(),
model_path=str(tmp_path),
tensor_parallel_size=1,
use_hsdp=False,
online_quantization=False,
)
backend._load_weights_via_mmap(pipeline, modules)
assert result.plan is not None
backend._load_weights_via_mmap(pipeline, modules, result.plan)
assert pipeline.transformer.validate_called, "validate_loaded_weights should be called"
class TestConfigValidation:
"""Tests for configuration validation in OffloadConfig / execute_request."""
def test_loader_plan_cannot_be_silently_dropped_on_unsupported_platform(self, monkeypatch):
import vllm_omni.diffusion.offloader as offloader_module
od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
dlo_resident_layers=0,
pin_cpu_memory=False,
model="unused",
parallel_config=SimpleNamespace(
use_hsdp=False,
data_parallel_size=1,
sequence_parallel_size=1,
),
)
plan = HostWeightPlan(
backing_kind="checkpoint_mmap",
bindings={},
)
monkeypatch.setattr(
offloader_module.current_omni_platform,
"supports_cpu_offload",
lambda: False,
)
with pytest.raises(RuntimeError, match="skipped ordinary weight materialization"):
offloader_module.get_offload_backend(
od_config,
device=torch.device("cpu"),
host_weight_plan=plan,
)
def test_hsdp_with_allgather_rejected(self):
"""HSDP + DLO + AllGather should raise ValueError (double sharding)."""
from vllm_omni.diffusion.offloader.base import OffloadConfig
tests/diffusion/test_diffusion_model_runner.py
+8 -1
View File
@@ -740,6 +740,9 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
del kwargs
return SimpleNamespace(transformer=torch.nn.Identity())
def take_host_weight_plan(self):
return None
class _DummyMemoryProfiler:
consumed_memory = 0
@@ -779,7 +782,11 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader)
monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
monkeypatch.setattr(
model_runner_module,
"get_offload_backend",
lambda od_config, device, host_weight_plan: None,
)
monkeypatch.setattr(
model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
)
vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py
+8
View File
@@ -9,6 +9,13 @@ from .modelopt import (
)
def get_direct_mmap_adapter(model: nn.Module):
"""Resolve direct-mmap adapters lazily to avoid model-loader cycles."""
from .direct_mmap import get_direct_mmap_adapter as _get_direct_mmap_adapter
return _get_direct_mmap_adapter(model)
def get_checkpoint_adapter(
model: nn.Module,
source: object,
@@ -29,4 +36,5 @@ __all__ = [
"ModelOptMixedPrecisionCheckpointAdapter",
"ModelOptNvFp4CheckpointAdapter",
"get_checkpoint_adapter",
"get_direct_mmap_adapter",
]
vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py
+108
View File
@@ -0,0 +1,108 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Loader-side adapters for proven direct-checkpoint mmap layouts."""
from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass
from functools import partial
from typing import Protocol
import torch
from torch import nn
@dataclass(frozen=True)
class DirectMmapTensorPolicy:
"""Loader-owned proof for one tensor's direct mmap behavior."""
allow_custom_loader: bool = False
transform: Callable[[torch.Tensor], torch.Tensor] | None = None
class DirectMmapAdapter(Protocol):
def policy_for(
self,
runtime_name: str,
target: torch.Tensor,
) -> DirectMmapTensorPolicy | None: ...
class _MiniMaxH3DirectMmapAdapter:
"""TP1 runtime-layout contract already enforced by MiniMax-H3's loader."""
def __init__(self, pipeline: nn.Module) -> None:
self.pipeline = pipeline
def policy_for(
self,
runtime_name: str,
target: torch.Tensor,
) -> DirectMmapTensorPolicy:
del target
transform = None
suffix = ".qkv_proj.weight"
if runtime_name.endswith(suffix):
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
MiniMaxH3Attention,
_reorder_grouped_qkv_to_qkv,
)
attention_path = runtime_name[: -len(suffix)]
attention = self.pipeline.get_submodule(attention_path)
if not isinstance(attention, MiniMaxH3Attention):
raise ValueError(
f"MiniMax-H3 direct-mmap adapter expected attention at {attention_path!r}, "
f"got {type(attention).__name__}"
)
transform = partial(
_reorder_grouped_qkv_to_qkv,
num_query_groups=attention.total_num_heads,
heads_per_group=1,
head_dim=attention.head_dim,
)
return DirectMmapTensorPolicy(
allow_custom_loader=True,
transform=transform,
)
class _Cosmos3DirectMmapAdapter:
"""Preserve Cosmos3's existing TP1 direct-mmap load contract."""
def policy_for(
self,
runtime_name: str,
target: torch.Tensor,
) -> DirectMmapTensorPolicy:
del runtime_name, target
# Cosmos3 has no packed module mapping or TP1 semantic transform after
# its checkpoint-key remap. Its custom vLLM loaders reduce to exact
# copies at TP1; topology and metadata are checked by plan preflight.
return DirectMmapTensorPolicy(allow_custom_loader=True)
def get_direct_mmap_adapter(pipeline: nn.Module) -> DirectMmapAdapter | None:
"""Return a loader-owned adapter without requiring a pipeline flag."""
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
MiniMaxH3DiTModel,
)
if any(isinstance(module, MiniMaxH3DiTModel) for module in pipeline.modules()):
return _MiniMaxH3DirectMmapAdapter(pipeline)
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
Cosmos3VFMTransformer,
)
if any(isinstance(module, Cosmos3VFMTransformer) for module in pipeline.modules()):
return _Cosmos3DirectMmapAdapter()
return None
__all__ = [
"DirectMmapAdapter",
"DirectMmapTensorPolicy",
"get_direct_mmap_adapter",
]
vllm_omni/diffusion/model_loader/diffusers_loader.py
+68 -20
View File
@@ -6,7 +6,7 @@ import glob
import os
import re
import time
from collections.abc import Generator, Iterable
from collections.abc import Generator, Iterable, Sequence
from pathlib import Path
from typing import cast
@@ -34,6 +34,11 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp
from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
get_checkpoint_adapter,
)
from vllm_omni.diffusion.model_loader.host_weight_plan import (
HostWeightPlan,
build_checkpoint_mmap_plan,
has_online_quantization,
)
from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
from vllm_omni.diffusion.registry import initialize_model
@@ -135,6 +140,13 @@ class DiffusersPipelineLoader:
self.od_config = od_config
self.quant_config = od_config.quantization_config
self.parallel_config = od_config.parallel_config
self.host_weight_plan: HostWeightPlan | None = None
def take_host_weight_plan(self) -> HostWeightPlan | None:
"""Transfer the loader-produced plan to the offload backend."""
plan = self.host_weight_plan
self.host_weight_plan = None
return plan
def _prepare_weights(
self,
@@ -293,8 +305,10 @@ class DiffusersPipelineLoader:
def get_all_weights(
self,
model: nn.Module,
sources: Sequence["ComponentSource"] | None = None,
) -> Generator[tuple[str, torch.Tensor], None, None]:
sources = self._get_weight_sources(model)
if sources is None:
sources = self._get_weight_sources(model)
for source in sources:
yield from self._get_weights_iterator(source, model=model)
@@ -386,6 +400,7 @@ class DiffusersPipelineLoader:
device: torch.device | None = None,
) -> nn.Module:
"""Load a model with the given configurations."""
self.host_weight_plan = None
if load_format is None:
load_format = "default"
# CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8),
@@ -417,22 +432,54 @@ class DiffusersPipelineLoader:
else:
model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False)
# Skip load_weights only for DLO+AllGather when the model
# supports mmap loading and no online quantization is active.
# This condition MUST match the gate in
# DistributedLayerwiseOffloadBackend.enable() so that the
# loader skips ⟺ enable() uses mmap.
from vllm_omni.diffusion.offloader.offload_plan import supports_mmap_loading
_dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
_use_ag = getattr(self.od_config, "dlo_use_allgather", True)
_has_online_quant = self._has_online_quant(model)
_supports_mmap = supports_mmap_loading(model)
_tp_size = int(getattr(self.parallel_config, "tensor_parallel_size", 1))
_use_hsdp = bool(getattr(self.parallel_config, "use_hsdp", False))
_dp_size = int(getattr(self.parallel_config, "data_parallel_size", 1))
_sp_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1))
_dlo_group_size = _dp_size if _dp_size > 1 else _sp_size
_skip_load = _dist_offload and _use_ag and _supports_mmap and not _has_online_quant
plan_result = None
weight_sources = self._get_weight_sources(model)
if _dist_offload:
modules = ModuleDiscovery.discover(model)
plan_result = build_checkpoint_mmap_plan(
model,
dit_modules=tuple(zip(modules.dit_names, modules.dits)),
sources=weight_sources,
model_path=str(getattr(self.od_config, "model", "")) or None,
tensor_parallel_size=_tp_size,
use_hsdp=_use_hsdp,
online_quantization=_has_online_quant,
)
self.host_weight_plan = plan_result.plan
_skip_load = self.host_weight_plan is not None
if _skip_load:
logger.info("DLO+AllGather active: skipping load_weights (will load via mmap in enable())")
logger.info(
"DLO host-weight plan active (%s, %s): skipping ordinary materialization for %s",
"AllGather" if _use_ag and _dlo_group_size > 1 else "rank-local",
self.host_weight_plan.backing_kind,
sorted(self.host_weight_plan.planned_source_prefixes) or "legacy DiT sources",
)
ordinary_sources = tuple(
source
for source in weight_sources
if source.prefix not in self.host_weight_plan.planned_source_prefixes
)
if ordinary_sources:
logger.info(
"Loading %d component weight source(s) outside the DLO host-weight plan",
len(ordinary_sources),
)
self.load_weights(
model,
sources=ordinary_sources,
planned_weights=self.host_weight_plan.bindings,
)
else:
if _dist_offload and _use_ag and _has_online_quant:
raise ValueError(
@@ -442,9 +489,10 @@ class DiffusersPipelineLoader:
"Please use --dlo-no-use-allgather or disable online "
"quantization."
)
if _dist_offload and _use_ag and not _supports_mmap:
if _dist_offload and plan_result is not None:
logger.info(
"DLO+AllGather active but model does not support mmap: loading weights via regular loader."
"DLO direct checkpoint mmap unavailable; using ordinary loader: %s",
plan_result.fallback_reason,
)
logger.debug("Loading weights on %s ...", load_device)
if offload_after_quant:
@@ -498,11 +546,7 @@ class DiffusersPipelineLoader:
"""Whether any layer uses an online-quant method that defers weight
materialization onto the ``meta`` device (upstream vLLM
``uses_meta_device=True``, e.g. online FP8)."""
for module in model.modules():
quant_method = getattr(module, "quant_method", None)
if getattr(quant_method, "uses_meta_device", False):
return True
return False
return has_online_quantization(model)
def _apply_skip_softmax_calibration(self, model: nn.Module) -> None:
from vllm_omni.diffusion.attention.backends.trtllm_calibration import (
@@ -605,12 +649,16 @@ class DiffusersPipelineLoader:
model: nn.Module,
*,
stream_online_quant_to_cpu: bool = False,
sources: Sequence["ComponentSource"] | None = None,
planned_weights: Iterable[str] = (),
) -> None:
weights_to_load = self._get_expected_parameter_names(model)
weights = self.get_all_weights(model)
weights = self.get_all_weights(model) if sources is None else self.get_all_weights(model, sources=sources)
if stream_online_quant_to_cpu:
weights = self._stream_online_quant_weights_to_cpu(model, weights)
loaded_weights = model.load_weights(weights)
if loaded_weights is not None:
loaded_weights = set(loaded_weights).union(planned_weights)
self.counter_after_loading_weights = time.perf_counter()
logger.info_once(
vllm_omni/diffusion/model_loader/host_weight_plan.py
+355
View File
@@ -0,0 +1,355 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Loader-owned host-weight plans shared with diffusion offload backends."""
from __future__ import annotations
import glob
import json
import os
from collections.abc import Callable, Sequence
from dataclasses import dataclass
import torch
from safetensors import safe_open
from torch import nn
from vllm.logger import init_logger
from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
get_direct_mmap_adapter,
)
logger = init_logger(__name__)
TensorTransform = Callable[[torch.Tensor], torch.Tensor]
@dataclass(frozen=True)
class TensorBinding:
"""One runtime tensor backed by one safetensors entry."""
checkpoint_key: str
file_path: str
transform: TensorTransform | None = None
@dataclass(frozen=True)
class HostWeightPlan:
"""Complete, prevalidated host backing consumed by an offload backend."""
backing_kind: str
bindings: dict[str, TensorBinding]
planned_source_prefixes: frozenset[str] = frozenset()
@dataclass(frozen=True)
class HostWeightPlanResult:
"""A complete plan or a fail-closed reason to use the ordinary loader."""
plan: HostWeightPlan | None
fallback_reason: str | None = None
class _PlanIncompatibleError(RuntimeError):
pass
_SAFETENSORS_DTYPES: dict[str, torch.dtype] = {
"BOOL": torch.bool,
"U8": torch.uint8,
"I8": torch.int8,
"I16": torch.int16,
"I32": torch.int32,
"I64": torch.int64,
"F16": torch.float16,
"BF16": torch.bfloat16,
"F32": torch.float32,
"F64": torch.float64,
}
for _safe_name, _torch_name in (
("F8_E4M3", "float8_e4m3fn"),
("F8_E5M2", "float8_e5m2"),
):
if (_dtype := getattr(torch, _torch_name, None)) is not None:
_SAFETENSORS_DTYPES[_safe_name] = _dtype
def has_online_quantization(model: nn.Module) -> bool:
"""Whether any module defers online-quantized weights on ``meta``."""
return any(getattr(getattr(module, "quant_method", None), "uses_meta_device", False) for module in model.modules())
def _resolve_source_root(source: object) -> str:
source_root = os.fspath(getattr(source, "model_or_path"))
subfolder = getattr(source, "subfolder", None)
if not os.path.isdir(source_root):
from vllm.model_executor.model_loader.weight_utils import (
download_weights_from_hf,
)
source_root = download_weights_from_hf(
model_name_or_path=source_root,
cache_dir=None,
allow_patterns=["*.safetensors", "*.safetensors.index.json"],
revision=getattr(source, "revision", None),
subfolder=subfolder,
)
return os.path.join(source_root, subfolder) if subfolder is not None else source_root
def _record_binding(
model_to_ckpt: dict[str, tuple[str, str]],
model_name: str,
checkpoint_key: str,
file_path: str,
) -> None:
existing = model_to_ckpt.get(model_name)
candidate = (checkpoint_key, file_path)
if existing is not None and existing != candidate:
raise _PlanIncompatibleError(
f"multiple checkpoint entries map to runtime tensor {model_name!r}: {existing!r} and {candidate!r}"
)
model_to_ckpt[model_name] = candidate
def _add_safetensors_source(
checkpoint_dir: str,
prefix: str,
remap_fn: Callable[[str], str | None] | None,
model_to_ckpt: dict[str, tuple[str, str]],
) -> int:
indexed_keys = 0
index_files = sorted(glob.glob(os.path.join(checkpoint_dir, "*.safetensors.index.json")))
if index_files:
for index_file in index_files:
with open(index_file, encoding="utf-8") as handle:
index = json.load(handle)
for checkpoint_key, filename in index.get("weight_map", {}).items():
source_name = prefix + checkpoint_key
model_name = remap_fn(source_name) if remap_fn is not None else source_name
if model_name is not None:
_record_binding(
model_to_ckpt,
model_name,
checkpoint_key,
os.path.join(checkpoint_dir, filename),
)
indexed_keys += 1
return indexed_keys
for filename in sorted(glob.glob(os.path.join(checkpoint_dir, "*.safetensors"))):
with safe_open(filename, framework="pt", device="cpu") as handle:
for checkpoint_key in handle.keys():
source_name = prefix + checkpoint_key
model_name = remap_fn(source_name) if remap_fn is not None else source_name
if model_name is not None:
_record_binding(model_to_ckpt, model_name, checkpoint_key, filename)
indexed_keys += 1
return indexed_keys
def _build_source_map(
*,
sources: Sequence[object],
model_path: str | None,
remap_fn: Callable[[str], str | None] | None,
) -> dict[str, tuple[str, str]]:
model_to_ckpt: dict[str, tuple[str, str]] = {}
indexed_keys = 0
for source in sources:
indexed_keys += _add_safetensors_source(
_resolve_source_root(source),
getattr(source, "prefix", ""),
remap_fn,
model_to_ckpt,
)
if not sources:
if remap_fn is None:
raise _PlanIncompatibleError(
"the loader has no component weight sources and the pipeline has no checkpoint-key remapper"
)
if not model_path:
raise _PlanIncompatibleError("the loader has no model path for legacy checkpoint remapping")
resolved_model_path = model_path
if not os.path.isdir(resolved_model_path):
from vllm.model_executor.model_loader.weight_utils import (
download_weights_from_hf,
)
resolved_model_path = download_weights_from_hf(
model_name_or_path=resolved_model_path,
cache_dir=None,
allow_patterns=["*.safetensors", "*.safetensors.index.json"],
)
checkpoint_dirs = {
os.path.dirname(path)
for pattern in ("*.safetensors.index.json", "*.safetensors")
for path in glob.glob(os.path.join(resolved_model_path, "**", pattern), recursive=True)
}
for checkpoint_dir in sorted(checkpoint_dirs):
indexed_keys += _add_safetensors_source(
checkpoint_dir,
"",
remap_fn,
model_to_ckpt,
)
if not model_to_ckpt:
raise _PlanIncompatibleError("no compatible safetensors entries were found in the loader's model sources")
logger.info(
"Indexed %d runtime tensor names from %d checkpoint keys for DLO preflight",
len(model_to_ckpt),
indexed_keys,
)
return model_to_ckpt
def _is_persistent_buffer(module: nn.Module, local_name: str) -> bool:
parts = local_name.split(".")
owner = module.get_submodule(".".join(parts[:-1])) if len(parts) > 1 else module
return parts[-1] not in owner._non_persistent_buffers_set
def _collect_required_dit_tensors(
dit_modules: Sequence[tuple[str, nn.Module]],
) -> dict[str, torch.Tensor]:
required: dict[str, torch.Tensor] = {}
for dit_name, dit_module in dit_modules:
for local_name, tensor in dit_module.named_parameters():
required[f"{dit_name}.{local_name}"] = tensor
for local_name, tensor in dit_module.named_buffers():
if _is_persistent_buffer(dit_module, local_name):
required[f"{dit_name}.{local_name}"] = tensor
return required
def _planned_source_prefixes(
sources: Sequence[object],
dit_modules: Sequence[tuple[str, nn.Module]],
) -> frozenset[str]:
"""Identify component sources that the DiT-only plan owns completely."""
if not sources:
return frozenset()
expected = frozenset(f"{dit_name}." for dit_name, _ in dit_modules)
available = {getattr(source, "prefix", "") for source in sources}
missing = sorted(expected - available)
if missing:
raise _PlanIncompatibleError(
f"direct mmap requires a dedicated component weight source for each DiT; missing prefixes: {missing}"
)
return expected
def _validate_source_metadata(
required: dict[str, torch.Tensor],
bindings: dict[str, TensorBinding],
) -> None:
by_file: dict[str, list[tuple[str, torch.Tensor, TensorBinding]]] = {}
for runtime_name, target in required.items():
by_file.setdefault(bindings[runtime_name].file_path, []).append((runtime_name, target, bindings[runtime_name]))
for file_path, entries in by_file.items():
if not os.path.isfile(file_path):
raise _PlanIncompatibleError(f"checkpoint file does not exist: {file_path}")
with safe_open(file_path, framework="pt", device="cpu") as handle:
available = set(handle.keys())
for runtime_name, target, binding in entries:
if binding.checkpoint_key not in available:
raise _PlanIncompatibleError(
f"checkpoint key {binding.checkpoint_key!r} for {runtime_name!r} is missing from {file_path}"
)
source = handle.get_slice(binding.checkpoint_key)
source_shape = tuple(source.get_shape())
source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype())
if source_shape != tuple(target.shape):
raise _PlanIncompatibleError(
f"shape mismatch for {runtime_name!r}: checkpoint={source_shape}, runtime={tuple(target.shape)}"
)
if source_dtype is None or source_dtype != target.dtype:
raise _PlanIncompatibleError(
f"dtype mismatch for {runtime_name!r}: checkpoint={source.get_dtype()}, runtime={target.dtype}"
)
def build_checkpoint_mmap_plan(
pipeline: nn.Module,
*,
dit_modules: Sequence[tuple[str, nn.Module]],
sources: Sequence[object],
model_path: str | None,
tensor_parallel_size: int,
use_hsdp: bool,
online_quantization: bool,
) -> HostWeightPlanResult:
"""Build a complete direct-checkpoint plan or return a fallback reason."""
if tensor_parallel_size != 1:
return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader")
if use_hsdp:
return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
if online_quantization:
return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
if not callable(remap_fn):
remap_fn = None
adapter = get_direct_mmap_adapter(pipeline)
try:
planned_source_prefixes = _planned_source_prefixes(sources, dit_modules)
planned_sources = tuple(
source for source in sources if getattr(source, "prefix", "") in planned_source_prefixes
)
model_to_ckpt = _build_source_map(
sources=planned_sources,
model_path=model_path,
remap_fn=remap_fn,
)
required = _collect_required_dit_tensors(dit_modules)
if not required:
raise _PlanIncompatibleError("no DiT parameters or persistent buffers were discovered")
missing = sorted(set(required) - set(model_to_ckpt))
if missing:
raise _PlanIncompatibleError(
f"{len(missing)} required DiT tensors have no checkpoint binding (first 5: {missing[:5]})"
)
bindings: dict[str, TensorBinding] = {}
for runtime_name, target in required.items():
policy = adapter.policy_for(runtime_name, target) if adapter is not None else None
has_custom_loader = callable(getattr(target, "weight_loader", None))
if has_custom_loader and not (policy and policy.allow_custom_loader):
raise _PlanIncompatibleError(
f"{runtime_name!r} requires a custom weight_loader with no direct-mmap adapter"
)
checkpoint_key, file_path = model_to_ckpt[runtime_name]
bindings[runtime_name] = TensorBinding(
checkpoint_key=checkpoint_key,
file_path=file_path,
transform=policy.transform if policy is not None else None,
)
_validate_source_metadata(required, bindings)
except (OSError, ValueError, _PlanIncompatibleError) as exc:
return HostWeightPlanResult(None, str(exc))
return HostWeightPlanResult(
HostWeightPlan(
backing_kind="checkpoint_mmap",
bindings=bindings,
planned_source_prefixes=planned_source_prefixes,
)
)
__all__ = [
"HostWeightPlan",
"HostWeightPlanResult",
"TensorBinding",
"build_checkpoint_mmap_plan",
"has_online_quantization",
]
vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py
-7
View File
@@ -540,13 +540,6 @@ class MiniMaxH3Pipeline(
encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")},
on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}),
)
# H3's regular loader performs checkpoint-layout conversions (grouped QKV
# and fused MLP). Keep it on that path until mmap can run the same loader
# callbacks for every affected parameter.
_supports_mmap_loading: ClassVar[bool] = False
# TODO(offload): Re-enable after the generic rank-local mmap path can run
# this model's grouped-QKV reorder and fused-MLP packing before TP sharding.
# Do not bypass the regular loader until that equivalence is tested.
_PROFILER_TARGETS: ClassVar[list[str]] = [
"_prepare_reference_videos",
"encode_prompt",
vllm_omni/diffusion/offloader/__init__.py
+25 -3
View File
@@ -5,6 +5,7 @@ import torch
from vllm.logger import init_logger
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan
from vllm_omni.platforms import current_omni_platform
from .base import OffloadBackend, OffloadConfig, OffloadStrategy
@@ -17,7 +18,7 @@ from .distributed_layerwise_backend import (
)
from .layerwise_backend import LayerWiseOffloadBackend
from .module_residency import PinnedModuleStager
from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading
from .offload_plan import OffloadPlan, get_offload_plan
from .sequential_backend import (
ModelLevelOffloadBackend,
apply_sequential_offload,
@@ -49,7 +50,6 @@ __all__ = [
"remove_distributed_block_hook",
"get_offload_backend",
"get_offload_plan",
"supports_mmap_loading",
"get_blocks_attr_names",
"get_blocks_from_dit",
"set_blocks_attr_names",
@@ -64,12 +64,15 @@ __all__ = [
def get_offload_backend(
od_config: OmniDiffusionConfig,
device: torch.device | None = None,
host_weight_plan: HostWeightPlan | None = None,
) -> OffloadBackend | None:
"""Create appropriate offload backend based on configuration.
Args:
od_config: OmniDiffusionConfig with offload settings
device: Target device (auto-detected if None)
host_weight_plan: Exact loader-produced backing plan, if ordinary
weight materialization was skipped.
Returns:
OffloadBackend instance or None if offloading disabled
@@ -82,12 +85,22 @@ def get_offload_backend(
# Extract and validate configuration
config = OffloadConfig.from_od_config(od_config)
if host_weight_plan is not None and config.strategy != OffloadStrategy.DISTRIBUTED_LAYER_WISE:
raise RuntimeError(
"A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected"
)
# Return None if no offloading requested
if config.strategy == OffloadStrategy.NONE:
return None
# Validate platform (CUDA required for now)
if not current_omni_platform.supports_cpu_offload() or current_omni_platform.get_device_count() < 1:
if host_weight_plan is not None:
raise RuntimeError(
"The loader skipped ordinary weight materialization for DLO, "
"but this platform cannot create the required offload backend"
)
logger.warning(
"Current device: %s does not support CPU offloading. Skipping offloading.",
current_omni_platform.get_device_name(),
@@ -99,6 +112,11 @@ def get_offload_backend(
try:
device = current_omni_platform.get_torch_device()
except (NotImplementedError, AttributeError) as exc:
if host_weight_plan is not None:
raise RuntimeError(
"The loader skipped ordinary weight materialization for DLO, "
"but the target device could not be resolved"
) from exc
logger.error("Failed to detect device: %s. Skipping offloading.", exc)
return None
@@ -108,7 +126,11 @@ def get_offload_backend(
elif config.strategy == OffloadStrategy.LAYER_WISE:
return LayerWiseOffloadBackend(config, device)
elif config.strategy == OffloadStrategy.DISTRIBUTED_LAYER_WISE:
return DistributedLayerwiseOffloadBackend(config, device)
return DistributedLayerwiseOffloadBackend(
config,
device,
host_weight_plan=host_weight_plan,
)
else:
logger.error("Unknown offload strategy: %s", config.strategy)
return None
vllm_omni/diffusion/offloader/base.py
+4 -6
View File
@@ -27,11 +27,10 @@ class OffloadConfig:
pin_cpu_memory: bool = True
use_hsdp: bool = False
dp_size: int = 1 # derived from parallel_config, not user-configurable
# True: add DP sharding + AllGather. False: stream the standard loader's
# rank-local tensors (including TP-local shards) with H2D only.
# True: add DP sharding + AllGather. False: stream complete rank-local
# blocks from the loader-selected host backing with H2D only.
dlo_use_allgather: bool = True
dlo_resident_layers: int = 0 # leading DiT layers kept on device
model_path: str | None = None # checkpoint path for mmap weight loading
@classmethod
def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
@@ -58,7 +57,6 @@ class OffloadConfig:
parallel_config = getattr(od_config, "parallel_config", None)
use_hsdp = getattr(parallel_config, "use_hsdp", False) if parallel_config else False
# Derive dp_size from parallel_config — not user-configurable.
# The offload adapts to whatever DP/SP is already configured.
dp_size = 1
@@ -116,7 +114,8 @@ class OffloadConfig:
dp_size = 1
logger.info(
"Distributed layerwise offload: dlo_use_allgather=False, "
"streaming standard-loader rank-local weights (no DP shard or AllGather)"
"streaming complete rank-local blocks (no DLO shard or AllGather); "
"the backend will select mmap or standard-loader host storage"
)
# HSDP already shards parameters into DTensors. Running distributed
@@ -137,7 +136,6 @@ class OffloadConfig:
dp_size=dp_size,
dlo_use_allgather=dlo_use_allgather,
dlo_resident_layers=dlo_resident_layers,
model_path=getattr(od_config, "model", None),
)
vllm_omni/diffusion/offloader/distributed_layerwise_backend.py
+476 -415
View File
File diff suppressed because it is too large. Load diff
vllm_omni/diffusion/offloader/offload_plan.py
-17
View File
@@ -55,20 +55,3 @@ class OffloadPlan:
def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None:
"""Retrieve the OffloadPlan declared by the pipeline, if any."""
return getattr(pipeline, "_offload_plan", None)
def supports_mmap_loading(pipeline: nn.Module) -> bool:
"""Whether the pipeline supports mmap weight loading.
A pipeline supports mmap if any module defines ``_remap_ckpt_key``,
which maps checkpoint keys to model parameter names. Models without
this method must use the regular ``load_weights()`` path.
This check is shared between ``diffusers_loader.py`` (to decide whether
to skip ``load_weights``) and ``DistributedLayerwiseOffloadBackend.enable()``
(to decide whether to call ``_load_weights_via_mmap``), ensuring both
paths use the same gate condition.
"""
return bool(getattr(pipeline, "_supports_mmap_loading", True)) and any(
callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules()
)
vllm_omni/diffusion/worker/diffusion_model_runner.py
+5 -1
View File
@@ -292,7 +292,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
)
# Apply CPU offloading
self.offload_backend = get_offload_backend(self.od_config, device=self.device)
self.offload_backend = get_offload_backend(
self.od_config,
device=self.device,
host_weight_plan=model_loader.take_host_weight_plan(),
)
if self.offload_backend is not None:
logger.info(f" Enabling offloader backend: {self.offload_backend.__class__.__name__}")
self.offload_backend.enable(self.pipeline)