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

[Diffusion] Rebase DLO plan and generalize lifecycle example

Browse Source 
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed 2026-08-15 23:40:24 +08:00
1 parent 88614175c0
commit 11687f73bd
3 files changed
+89 -37
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py → examples/offline_inference/minimax_h3/dlo_lifecycle.py
+56 -28
View File
@@ -1,28 +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 \
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
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-no-allgather
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
@@ -47,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", "dlo-dp2-no-allgather", "request"),
choices=("dlo", "dlo-no-allgather", "request"),
required=True,
help="DLO DP2 (with or without AllGather) 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)
@@ -67,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",
@@ -81,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 in {"dlo-dp2", "dlo-dp2-no-allgather"}:
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=args.mode == "dlo-dp2",
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
@@ -179,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(
@@ -189,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,
)
@@ -202,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/offloader/test_distributed_layerwise_backend.py
+3 -2
View File
@@ -404,7 +404,7 @@ class TestDistributedLayerwiseOffloadHook:
source_storage = next_block.weight.untyped_storage().data_ptr()
def transform(tensor):
return tensor.flip(0)
return tensor.t()
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
@@ -431,8 +431,9 @@ class TestDistributedLayerwiseOffloadHook:
assert torch.equal(
next_block.weight,
torch.tensor([[2.0, 3.0], [0.0, 1.0]]),
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):
vllm_omni/diffusion/offloader/distributed_layerwise_backend.py
+30 -7
View File
@@ -257,6 +257,19 @@ class DistributedLayerwiseOffloadHook(ModelHook):
dtype = source.dtype
offset = offsets.get(dtype, 0)
transform = (tensor_transforms or {}).get(id(target))
runtime_source = transform(source) if callable(transform) else source
if runtime_source.dtype != dtype or runtime_source.shape != source.shape:
raise ValueError(
"mmap weight transform changed tensor metadata for "
f"{name!r}: expected dtype={dtype}, shape={tuple(source.shape)}, "
f"got dtype={runtime_source.dtype}, shape={tuple(runtime_source.shape)}"
)
stride = runtime_source.stride()
storage_numel = (
0
if runtime_source.numel() == 0
else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(runtime_source.shape, stride))
)
cpu_sources.setdefault(dtype, []).append(
{
@@ -269,11 +282,12 @@ class DistributedLayerwiseOffloadHook(ModelHook):
{
"name": name,
"offset": offset,
"numel": source.numel(),
"shape": source.shape,
"numel": storage_numel,
"shape": runtime_source.shape,
"stride": stride,
}
)
offsets[dtype] = offset + source.numel()
offsets[dtype] = offset + storage_numel
# The detached source above keeps the mmap storage alive while the
# module parameter/buffer is rebound to the rotating device slot.
@@ -446,15 +460,24 @@ class DistributedLayerwiseOffloadHook(ModelHook):
transform = source_info["transform"]
if callable(transform):
source = transform(source)
if source.dtype != dtype or source.numel() != meta["numel"]:
if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
raise ValueError(
"mmap weight transform changed tensor layout for "
f"{source_info['name']!r}: expected dtype={dtype}, "
f"numel={meta['numel']}, got dtype={source.dtype}, "
f"numel={source.numel()}"
f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
f"stride={source.stride()}"
)
start = meta["offset"]
destination[start : start + meta["numel"]].copy_(source.reshape(-1))
physical_storage = destination[start : start + meta["numel"]]
if source.is_contiguous():
physical_storage.copy_(source.flatten())
else:
torch.as_strided(
physical_storage,
size=source.shape,
stride=source.stride(),
).copy_(source)
staged[dtype] = destination
return staged