evanchuenganddavid6666666
744c65b730
Signed-off-by: evanchueng <evanchueng@outlook.com> Signed-off-by: david6666666 <530634352@qq.com> Co-authored-by: david6666666 <530634352@qq.com>
113 lines
3.6 KiB
Python
113 lines
3.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import torch
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
from .base import OffloadBackend, OffloadConfig, OffloadStrategy
|
|
from .block_discovery import get_blocks_attr_names, get_blocks_from_dit, set_blocks_attr_names
|
|
from .distributed_layerwise_backend import (
|
|
DistributedLayerwiseOffloadBackend,
|
|
DistributedLayerwiseOffloadHook,
|
|
apply_distributed_block_hook,
|
|
remove_distributed_block_hook,
|
|
)
|
|
from .layerwise_backend import LayerWiseOffloadBackend
|
|
from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading
|
|
from .sequential_backend import (
|
|
ModelLevelOffloadBackend,
|
|
apply_sequential_offload,
|
|
remove_sequential_offload,
|
|
)
|
|
from .tensor_utils import (
|
|
dtype_size,
|
|
is_dtensor,
|
|
is_materialized_tensor,
|
|
make_offload_placeholder,
|
|
set_tensor_storage,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
__all__ = [
|
|
"OffloadBackend",
|
|
"OffloadConfig",
|
|
"OffloadPlan",
|
|
"OffloadStrategy",
|
|
"LayerWiseOffloadBackend",
|
|
"DistributedLayerwiseOffloadBackend",
|
|
"DistributedLayerwiseOffloadHook",
|
|
"ModelLevelOffloadBackend",
|
|
"apply_sequential_offload",
|
|
"remove_sequential_offload",
|
|
"apply_distributed_block_hook",
|
|
"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",
|
|
"dtype_size",
|
|
"is_dtensor",
|
|
"is_materialized_tensor",
|
|
"make_offload_placeholder",
|
|
"set_tensor_storage",
|
|
]
|
|
|
|
|
|
def get_offload_backend(
|
|
od_config: OmniDiffusionConfig,
|
|
device: torch.device | 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)
|
|
|
|
Returns:
|
|
OffloadBackend instance or None if offloading disabled
|
|
|
|
Example:
|
|
>>> backend = get_offload_backend(od_config, device=torch.device("cuda:0"))
|
|
>>> if backend:
|
|
... backend.enable(pipeline)
|
|
"""
|
|
# Extract and validate configuration
|
|
config = OffloadConfig.from_od_config(od_config)
|
|
|
|
# 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:
|
|
logger.warning(
|
|
"Current device: %s does not support CPU offloading. Skipping offloading.",
|
|
current_omni_platform.get_device_name(),
|
|
)
|
|
return None
|
|
|
|
# Detect device if not provided
|
|
if device is None:
|
|
try:
|
|
device = current_omni_platform.get_torch_device()
|
|
except (NotImplementedError, AttributeError) as exc:
|
|
logger.error("Failed to detect device: %s. Skipping offloading.", exc)
|
|
return None
|
|
|
|
# Create appropriate backend
|
|
if config.strategy == OffloadStrategy.MODEL_LEVEL:
|
|
return ModelLevelOffloadBackend(config, device)
|
|
elif config.strategy == OffloadStrategy.LAYER_WISE:
|
|
return LayerWiseOffloadBackend(config, device)
|
|
elif config.strategy == OffloadStrategy.DISTRIBUTED_LAYER_WISE:
|
|
return DistributedLayerwiseOffloadBackend(config, device)
|
|
else:
|
|
logger.error("Unknown offload strategy: %s", config.strategy)
|
|
return None
|