diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index af4eee100..d48509099 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -29,6 +29,7 @@ Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsu | **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. | | **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. | | **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. | +| **Auxiliary component policy** | ✅ Encoders/VAEs may remain resident. | ✅ Encoders/VAEs may remain resident. | See [Parallelism compatibility](#parallelism-compatibility) and [Request and loading constraints](#request-and-loading-constraints) for the @@ -419,6 +420,22 @@ failure. - **HSDP + DP or TP:** rejected independently by the diffusion parallel configuration. +### Auxiliary component policy + +`dlo_offload_components` maps discovered encoder/VAE paths to booleans. A +component set to `False` is materialized on the target device once and is not +entered through the model's manual offload lifecycle. A `True` value permits +the existing `OffloadPlan.on_demand_component_paths` and encoder block-hook +behavior; it does not add support where the model declares none. The optional +`default` key handles unmatched auxiliary components, and an empty map +preserves the current policy. + +The policy is resolved after component discovery and before DLO mutates DiT +storage. Unknown names, including DiT paths, are rejected. DLO continues to +own DiT streaming regardless of this map, which avoids an invalid state where +the loader has skipped ordinary DiT materialization but the backend is asked +to keep that DiT resident. + ## Request and loading constraints AllGather DP multi-concurrency requires: @@ -430,6 +447,15 @@ AllGather DP multi-concurrency requires: The no-AllGather path does not impose these DLO-specific synchronized-wave requirements. +MiniMax-H3 currently builds one global encoder TP group and broadcasts the +resulting conditioning from global rank 0. Therefore an encoder TP group that +spans multiple DP replicas is not valid for independent-prompt AllGather +waves. Exact-recipe DP2/TP2 and DP2/SP2 experiments with encoder TP4 completed +the collective schedule, but two distinct prompts with the same seed produced +byte-identical video and audio in every measured wave. Enabling or disabling +encoder offload produced paired-identical results, so this is a pre-existing +model topology limitation rather than component-policy behavior. + 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. Per-tensor online FP8 linears can use DLO AllGather after the @@ -494,6 +520,9 @@ 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. +These B300 DP waves intentionally reused the same prompt and seed, so they do +not establish independent-prompt isolation across DP replicas. + ### Host-memory measurement A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 2c6bbacb4..63a22bef8 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -36,7 +36,8 @@ vllm serve /path/to/model --omni \ vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ --data-parallel-size 4 \ - --dlo-no-use-allgather + --dlo-no-use-allgather \ + --dlo-offload-components '{"text_encoder": false}' # Sequence parallel deployment vllm serve /path/to/model --omni \ @@ -51,6 +52,7 @@ omni = Omni( model="/path/to/model", enable_distributed_layerwise_offload=True, dlo_use_allgather=True, + dlo_offload_components={"text_encoder": False}, ) ``` @@ -63,6 +65,7 @@ omni = Omni( | `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` | | `--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` | +| `--dlo-offload-components JSON` | Permit or disable offload for discovered encoder/VAE components | `{}` | | `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` | | `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset | | `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` | @@ -244,6 +247,25 @@ class MyPipeline(nn.Module): When no plan exists, discovery falls back to `_layerwise_offload_blocks_attrs` and then heuristic attribute lookup. +## Auxiliary component policy + +`dlo_offload_components` follows the same component-map style as diffusion +quantization. Keys are discovered encoder or VAE paths, values are booleans, +and `default` is the optional fallback: + +```python +dlo_offload_components = { + "text_encoder": False, + "default": True, +} +``` + +`True` permits the model's existing DLO behavior for that component. `False` +keeps the whole component resident, including its child blocks and non-block +modules. An empty map preserves current behavior. Unknown keys are rejected; +the map does not control DiT blocks because enabling DLO gives the backend +ownership of DiT streaming. + ## Data-parallel concurrency With `data_parallel_size > 1` and AllGather enabled, the scheduler can process @@ -268,6 +290,13 @@ must enter each collective. are validated. - Resident leading layers require `--dlo-no-use-allgather` and a model `OffloadPlan` that declares eligible `resident_dit_paths`. +- Component policy applies only to discovered auxiliary encoders and VAEs; + it cannot disable DLO for a DiT. +- MiniMax-H3's encoder TP group is global rather than DP-replica-local. An + encoder TP size that spans multiple DP replicas (for example TP4 with + DP2/TP2 or DP2/SP2 on four ranks) is not supported for independent-prompt + AllGather waves: the current pipeline broadcasts rank 0's conditioning to + both replicas. This limitation is independent of the component policy. - DP concurrency requires an explicit, identical inference-step count. Sharing quantized or otherwise unvalidated transformed layouts through a diff --git a/examples/offline_inference/minimax_h3/dlo_lifecycle.py b/examples/offline_inference/minimax_h3/dlo_lifecycle.py index 5630777cb..b369784d6 100644 --- a/examples/offline_inference/minimax_h3/dlo_lifecycle.py +++ b/examples/offline_inference/minimax_h3/dlo_lifecycle.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Validate MiniMax-H3 request-mode and DLO lifecycle behavior. The DLO modes accept any positive data-parallel size. With AllGather enabled, @@ -85,12 +85,20 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--height", type=int, default=768) parser.add_argument("--batch-wait-ms", type=float, default=500.0) parser.add_argument("--init-timeout", type=float, default=1800.0) + parser.add_argument( + "--dlo-offload-components", + type=json.loads, + default=None, + help="Component policy JSON, e.g. '{\"text_encoder\": false}'", + ) parser.add_argument("--output", type=Path) return parser.parse_args() def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]: is_dlo = args.mode in {"dlo", "dlo-no-allgather"} + if args.dlo_offload_components and not is_dlo: + raise ValueError("--dlo-offload-components requires a DLO mode") common: dict[str, Any] = { "model": args.model, "trust_remote_code": True, @@ -114,6 +122,7 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]: enable_distributed_layerwise_offload=True, dlo_use_allgather=args.mode == "dlo", dlo_resident_layers=0, + dlo_offload_components=({} if args.dlo_offload_components is None else args.dlo_offload_components), ) else: common.update( diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py index b0ebf1038..20c274e47 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py @@ -99,3 +99,22 @@ def test_manual_component_offload_failure_forces_retained_cache_release(mocker): assert component.load_to_device.call_count == 2 assert component.offload_to_cpu.call_count == 2 pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True) + + +def test_resident_component_skips_manual_dlo_lifecycle(mocker): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = mocker.Mock() + pipeline.od_config.enable_layerwise_offload = False + pipeline.od_config.enable_distributed_layerwise_offload = True + pipeline._model_cpu_offload_modules = [] + component = mocker.Mock() + component._omni_dlo_offload_enabled = False + + with pipeline._component_on_device(component): + pass + + component.load_to_device.assert_not_called() + component.offload_to_cpu.assert_not_called() diff --git a/tests/diffusion/offloader/test_component_policy.py b/tests/diffusion/offloader/test_component_policy.py new file mode 100644 index 000000000..a419e00de --- /dev/null +++ b/tests/diffusion/offloader/test_component_policy.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.distributed_layerwise_backend import ( + DistributedLayerwiseOffloadBackend, +) +from vllm_omni.diffusion.offloader.module_collector import PipelineModules +from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +def _od_config(**overrides): + values = { + "enable_cpu_offload": False, + "enable_layerwise_offload": False, + "enable_distributed_layerwise_offload": True, + "dlo_use_allgather": False, + "dlo_resident_layers": 0, + "dlo_offload_components": {}, + "pin_cpu_memory": False, + "parallel_config": SimpleNamespace( + use_hsdp=False, + data_parallel_size=1, + sequence_parallel_size=1, + ), + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _backend(policy: dict[str, bool]) -> DistributedLayerwiseOffloadBackend: + backend = object.__new__(DistributedLayerwiseOffloadBackend) + backend.config = OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dlo_use_allgather=False, + dlo_offload_components=policy, + ) + backend.device = torch.device("cpu") + return backend + + +def _modules() -> PipelineModules: + return PipelineModules( + dits=[], + dit_names=[], + encoders=[nn.Module()], + encoder_names=["text_encoder"], + vaes=[nn.Module(), nn.Module()], + vae_names=["video_vae", "audio_vae"], + ) + + +def test_component_policy_resolves_exact_then_default(): + backend = _backend({"text_encoder": False, "default": True}) + + assert not backend._component_offload_enabled("text_encoder") + assert backend._component_offload_enabled("video_vae") + + +@pytest.mark.parametrize( + ("value", "error", "match"), + [ + ([], TypeError, "must be a mapping"), + (["text_encoder"], TypeError, "must be a mapping"), + (False, TypeError, "must be a mapping"), + ({1: False}, TypeError, "keys must be non-empty strings"), + ({"text_encoder": 0}, TypeError, "must be a boolean"), + ], +) +def test_component_policy_rejects_invalid_shapes(value, error, match): + with pytest.raises(error, match=match): + OffloadConfig.from_od_config(_od_config(dlo_offload_components=value)) + + +def test_component_policy_accepts_none_as_default(): + config = OffloadConfig.from_od_config(_od_config(dlo_offload_components=None)) + + assert config.dlo_offload_components == {} + + +def test_component_policy_requires_dlo(): + with pytest.raises(ValueError, match="requires distributed layerwise offload"): + OffloadConfig.from_od_config( + _od_config( + enable_distributed_layerwise_offload=False, + dlo_offload_components={"text_encoder": False}, + ) + ) + + +def test_component_policy_rejects_unknown_or_dit_names(): + backend = _backend({"transformer": False}) + + with pytest.raises(ValueError, match="unknown auxiliary components.*transformer"): + backend._validate_component_offload_policy(_modules()) + + +@pytest.mark.parametrize( + ("policy", "encoder_stage", "video_stage", "audio_stage", "try_encoder"), + [ + ({}, True, True, True, True), + ({"text_encoder": False}, False, True, True, False), + ({"default": False, "video_vae": True}, False, True, False, False), + ], +) +def test_prepare_auxiliary_components_applies_policy( + mocker, + policy, + encoder_stage, + video_stage, + audio_stage, + try_encoder, +): + backend = _backend(policy) + backend._try_layerwise_offload_encoder = mocker.Mock() + backend._register_on_demand_hook = mocker.Mock() + modules = _modules() + plan = OffloadPlan( + encoder_block_attrs={"text_encoder": ("layers",)}, + on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}), + ) + + backend._prepare_auxiliary_components(modules, plan) + + assert backend._try_layerwise_offload_encoder.called is try_encoder + assert modules.encoders[0]._omni_dlo_offload_enabled is encoder_stage + assert modules.vaes[0]._omni_dlo_offload_enabled is video_stage + assert modules.vaes[1]._omni_dlo_offload_enabled is audio_stage + assert backend._register_on_demand_hook.call_args_list == [ + mocker.call(modules.encoders[0], "text_encoder", stage_on_demand=encoder_stage), + mocker.call(modules.vaes[0], "video_vae", stage_on_demand=video_stage), + mocker.call(modules.vaes[1], "audio_vae", stage_on_demand=audio_stage), + ] diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 452eec286..ed6ae4fa3 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -332,7 +332,7 @@ def test_serve_cli_forwards_distilled_lora_to_diffusion_stage(): ] -def test_serve_cli_forwards_distributed_offload_residency(): +def test_serve_cli_forwards_distributed_offload_options(): """Ensure the two-GPU DLO placement controls reach the diffusion stage.""" parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") @@ -347,6 +347,8 @@ def test_serve_cli_forwards_distributed_offload_residency(): "--dlo-no-use-allgather", "--dlo-resident-layers", "20", + "--dlo-offload-components", + '{"text_encoder": false, "default": true}', ] ) @@ -357,9 +359,14 @@ def test_serve_cli_forwards_distributed_offload_residency(): assert args.enable_distributed_layerwise_offload is True assert args.dlo_use_allgather is False assert args.dlo_resident_layers == 20 + assert args.dlo_offload_components == {"text_encoder": False, "default": True} assert engine_args["enable_distributed_layerwise_offload"] is True assert engine_args["dlo_use_allgather"] is False assert engine_args["dlo_resident_layers"] == 20 + assert engine_args["dlo_offload_components"] == { + "text_encoder": False, + "default": True, + } def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo(): diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 81afc4b8c..838a5bf96 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -697,6 +697,7 @@ class _DiffusionConfigProjection: enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True dlo_resident_layers: int = Field(default=0, ge=0) + dlo_offload_components: dict[str, bool] = Field(default_factory=dict) host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled" host_weight_runtime_root: str | None = None dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 4e9327530..1d117aa60 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -389,6 +389,7 @@ class StageDeployConfig: enable_distributed_layerwise_offload: bool | None = None dlo_use_allgather: bool | None = None dlo_resident_layers: int | None = None + dlo_offload_components: dict[str, bool] | None = None host_weight_runtime_mode: str | None = None host_weight_runtime_root: str | None = None dlo_host_registration_limit_gib: float | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index dece344b2..f55d9df1c 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -795,6 +795,9 @@ class OmniDiffusionConfig: dlo_use_allgather: bool = True # Leading main-DiT blocks kept resident by distributed layerwise offload. dlo_resident_layers: int = 0 + # Auxiliary encoder/VAE policy. False keeps the whole component resident; + # DiT offload remains owned by DLO. + dlo_offload_components: dict[str, bool] = field(default_factory=dict) # Final-layout Host Weight Runtime policy. The loader only activates this # for eligible no-AllGather DLO; all other configurations preserve their # existing loader/storage path. diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index ff3990b32..5937812b5 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -1436,7 +1436,7 @@ class MiniMaxH3Pipeline( with sequential_offload_component(component): yield return - staged = self._uses_manual_component_offload() + staged = self._uses_manual_component_offload() and getattr(component, "_omni_dlo_offload_enabled", True) try: if staged: component.load_to_device() diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index c4ad93482..3e0c56537 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -2,7 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from abc import ABC, abstractmethod -from dataclasses import dataclass +from collections.abc import Mapping +from dataclasses import dataclass, field from enum import Enum from typing import Protocol, runtime_checkable @@ -52,6 +53,9 @@ class OffloadConfig: # 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 + # Per auxiliary component: True permits model-declared DLO offload; + # False keeps the whole component resident. "default" is the fallback. + dlo_offload_components: dict[str, bool] = field(default_factory=dict) # Optional per-worker ceiling for registering an HWR mmap. Zero means no # additional ceiling; pin_cpu_memory controls whether registration is tried. dlo_host_registration_limit_gib: float = 0.0 @@ -125,6 +129,25 @@ class OffloadConfig: # requirements (concurrent requests, dummy run skip). dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True) dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0)) + raw_component_policy = getattr(od_config, "dlo_offload_components", {}) + if raw_component_policy is None: + raw_component_policy = {} + if not isinstance(raw_component_policy, Mapping): + raise TypeError( + "dlo_offload_components must be a mapping of component names to booleans, " + f"got {type(raw_component_policy).__name__}" + ) + dlo_offload_components: dict[str, bool] = {} + for component, enabled in raw_component_policy.items(): + if not isinstance(component, str) or not component: + raise TypeError("dlo_offload_components keys must be non-empty strings") + if not isinstance(enabled, bool): + raise TypeError( + f"dlo_offload_components[{component!r}] must be a boolean, got {type(enabled).__name__}" + ) + dlo_offload_components[component] = enabled + if dlo_offload_components and not enable_distributed_layerwise_offload: + raise ValueError("dlo_offload_components requires distributed layerwise offload") dlo_host_registration_limit_gib = validate_dlo_host_registration_options( limit_gib=getattr(od_config, "dlo_host_registration_limit_gib", 0.0), enable_dlo=enable_distributed_layerwise_offload, @@ -166,6 +189,7 @@ class OffloadConfig: dp_size=dp_size, dlo_use_allgather=dlo_use_allgather, dlo_resident_layers=dlo_resident_layers, + dlo_offload_components=dlo_offload_components, dlo_host_registration_limit_gib=dlo_host_registration_limit_gib, ) diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index 7063c2759..45c6c250a 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -43,7 +43,7 @@ from .host_registration import ( HostRegistrationError, register_host_mappings, ) -from .module_collector import ModuleDiscovery +from .module_collector import ModuleDiscovery, PipelineModules from .offload_plan import ( OffloadPlan, get_offload_plan, @@ -1400,6 +1400,39 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): module.to(self.device) logger.info("Moved %s (%s) to GPU (resident)", label, module.__class__.__name__) + def _component_offload_enabled(self, name: str) -> bool: + policy = self.config.dlo_offload_components + return policy.get(name, policy.get("default", True)) + + def _validate_component_offload_policy(self, modules: PipelineModules) -> None: + known_components = {"default", *modules.encoder_names, *modules.vae_names} + unknown = sorted(self.config.dlo_offload_components.keys() - known_components) + if unknown: + raise ValueError( + f"dlo_offload_components contains unknown auxiliary components: {unknown}. DLO always owns DiT offload." + ) + + def _prepare_auxiliary_components(self, modules: PipelineModules, plan: OffloadPlan | None) -> None: + for encoder, name in zip(modules.encoders, modules.encoder_names, strict=True): + offload = self._component_offload_enabled(name) + encoder._omni_dlo_offload_enabled = offload + if offload: + self._try_layerwise_offload_encoder(encoder, name, plan) + self._register_on_demand_hook( + encoder, + name, + stage_on_demand=(offload and plan is not None and name in plan.on_demand_component_paths), + ) + + for vae, name in zip(modules.vaes, modules.vae_names, strict=True): + offload = self._component_offload_enabled(name) + vae._omni_dlo_offload_enabled = offload + self._register_on_demand_hook( + vae, + name, + stage_on_demand=(offload and plan is not None and name in plan.on_demand_component_paths), + ) + def _try_layerwise_offload_encoder(self, module: nn.Module, name: str, plan: OffloadPlan | None) -> bool: """Stream plan-declared encoder blocks on each rank without AllGather.""" if plan is None or name not in plan.encoder_block_attrs: @@ -1668,6 +1701,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # Retrieve optional declarative OffloadPlan from the pipeline. # When present, replaces heuristic block discovery. plan = get_offload_plan(pipeline) + self._validate_component_offload_policy(modules) if self.config.dlo_resident_layers and (plan is None or not plan.resident_dit_paths): logger.warning( @@ -1728,21 +1762,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): ) logger.info("DLO is using host tensors materialized by the ordinary loader") - # Keep VAE/encoders on CPU; move to GPU on-demand via hooks. - # This saves ~4.3 GB HBM per card (VAE 1.3 + encoder 1.1 + sound 1.9) - # during the DiT forward pass. They are only needed briefly for - # text-encoding (before DiT) and VAE-decode (after DiT). - for enc, enc_name in zip(modules.encoders, modules.encoder_names): - self._try_layerwise_offload_encoder(enc, enc_name, plan) - self._register_on_demand_hook( - enc, "encoder", stage_on_demand=plan is not None and enc_name in plan.on_demand_component_paths - ) - for vae, vae_name in zip(modules.vaes, modules.vae_names): - self._register_on_demand_hook( - vae, - "vae", - stage_on_demand=(plan is not None and vae_name in plan.on_demand_component_paths), - ) + # Apply the user policy to discovered auxiliary components. The default + # keeps the model-declared staged lifecycle; False leaves a component + # resident for repeated requests. + self._prepare_auxiliary_components(modules, plan) # Move resident modules to GPU (small modules needed every forward) for name, module in zip(modules.resident_names, modules.resident_modules): diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 8fdda71cb..147302558 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -569,6 +569,7 @@ class OrchestratorArgs: enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True dlo_resident_layers: int = 0 + dlo_offload_components: dict[str, bool] = field(default_factory=dict) host_weight_runtime_mode: str = "disabled" host_weight_runtime_root: str | None = None dlo_host_registration_limit_gib: float = 0.0 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 4b44e7067..980377bd1 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1059,6 +1059,9 @@ class AsyncOmniEngine: "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), "dlo_use_allgather": kwargs.get("dlo_use_allgather", True), "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0), + "dlo_offload_components": ( + {} if kwargs.get("dlo_offload_components") is None else kwargs["dlo_offload_components"] + ), "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"), "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"), "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0), diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 87f4f8834..ebf1341cd 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -770,6 +770,18 @@ class OmniServeCommand(CLISubcommand): "under the existing pinned-memory policy and fall back to bounded staging when unavailable." ), ) + omni_config_group.add_argument( + "--dlo-offload-components", + type=json.loads, + default=None, + metavar="JSON", + help=( + "Per auxiliary encoder/VAE DLO policy as a JSON object, for example " + "'{\"text_encoder\": false}'. False keeps the whole component " + "resident; true permits the model's existing offload behavior. " + "DiT offload is not controlled by this map." + ), + ) # Video model parameters (e.g., Wan2.2) - engine-level omni_config_group.add_argument( "--boundary-ratio",