diff --git a/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py b/examples/offline_inference/minimax_h3/dlo_lifecycle.py similarity index 74% rename from examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py rename to examples/offline_inference/minimax_h3/dlo_lifecycle.py index 2cac7a666..5630777cb 100644 --- a/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py +++ b/examples/offline_inference/minimax_h3/dlo_lifecycle.py @@ -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) diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 9b8d4fc33..70741fe94 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -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): diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index c6aa464b9..d184e71ea 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -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