Anjie HouandClaude Opus 5
b6f4df320d
Signed-off-by: specture724 <specture724@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
66 lines
3.0 KiB
Python
66 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Declarative model capabilities for component layerwise offload.
|
|
|
|
Models declare an ``_offload_plan`` class attribute on the pipeline. Both the
|
|
ordinary and distributed layerwise backends consume the same plan instead of
|
|
embedding model-specific block paths in backend code.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
|
|
from torch import nn
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class OffloadPlan:
|
|
"""Optional declarative metadata for component layerwise offload.
|
|
|
|
Models declare this as a class attribute ``_offload_plan`` on the
|
|
pipeline class. When present, both layerwise backends use it instead of
|
|
model-specific backend branches, making new integrations data-driven.
|
|
|
|
If not declared, the offloader falls back to
|
|
``_layerwise_offload_blocks_attrs`` on each DiT module class.
|
|
|
|
:func:`~vllm_omni.diffusion.offloader.plan_resolver.resolve_offload_plan`
|
|
is the only consumer; backends read the resolved artifact it returns.
|
|
|
|
Attributes:
|
|
block_attrs: Maps DiT path → tuple of block-list attribute names.
|
|
e.g. ``{"transformer": ("gen_layers",),
|
|
"transformer.language_model": ("layers",)}``
|
|
offload_submodules: Maps child name → block-list attribute name,
|
|
for large non-DiT submodules within a DiT that should be
|
|
independently offloaded with their own hooks.
|
|
e.g. ``{"context_encoder": "layers"}``
|
|
resident_dit_paths: DiT paths whose leading blocks may be kept on the
|
|
device when ``dlo_resident_layers`` is nonzero. Keeping this
|
|
model-declared avoids applying a consumer-GPU tuning knob to
|
|
auxiliary or dual DiTs unintentionally.
|
|
encoder_component_types: Maps encoder paths to public selector types
|
|
(currently text_encoder). This declaration is used before the
|
|
compatibility name heuristic.
|
|
encoder_block_attrs: Maps encoder paths to streamable block-list paths.
|
|
encoder_dlo_weight_replication: Encoder paths whose loader-produced
|
|
block tensors are identical across the DiT DLO group. Only these
|
|
encoders may use multi-rank AllGather transfer; this must not be
|
|
declared for encoder-TP shards.
|
|
"""
|
|
|
|
on_demand_component_paths: frozenset[str] = field(default_factory=frozenset)
|
|
|
|
block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict)
|
|
offload_submodules: dict[str, str] = field(default_factory=dict)
|
|
resident_dit_paths: frozenset[str] = field(default_factory=frozenset)
|
|
encoder_component_types: dict[str, str] = field(default_factory=dict)
|
|
encoder_block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict)
|
|
encoder_dlo_weight_replication: frozenset[str] = field(default_factory=frozenset)
|
|
|
|
|
|
def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None:
|
|
"""Retrieve the OffloadPlan declared by the pipeline, if any."""
|
|
return getattr(pipeline, "_offload_plan", None)
|