Yuanheng ZhaoandSamit
8bf849c399
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com> Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com> Signed-off-by: Samit <285365963@qq.com> Co-authored-by: Samit <285365963@qq.com>
95 lines
2.9 KiB
Python
95 lines
2.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
|
|
import torch
|
|
from torch import nn
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class OffloadStrategy(Enum):
|
|
NONE = "none"
|
|
MODEL_LEVEL = "model_level" # Sequential offloading between DiT and encoders
|
|
LAYER_WISE = "layer_wise" # Block-level
|
|
|
|
|
|
@dataclass
|
|
class OffloadConfig:
|
|
strategy: OffloadStrategy
|
|
pin_cpu_memory: bool = True
|
|
|
|
@classmethod
|
|
def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
|
|
"""Extract and validate offload settings from OmniDiffusionConfig.
|
|
|
|
For now, enforces mutual exclusion between model-level and layer-wise offloading.
|
|
Layer-wise takes priority if both are enabled.
|
|
|
|
Args:
|
|
od_config: OmniDiffusionConfig with offload settings
|
|
|
|
Returns:
|
|
OffloadConfig with validated settings
|
|
"""
|
|
enable_cpu_offload = getattr(od_config, "enable_cpu_offload", False)
|
|
enable_layerwise_offload = getattr(od_config, "enable_layerwise_offload", False)
|
|
pin_cpu_memory = getattr(od_config, "pin_cpu_memory", True)
|
|
|
|
# Determine strategy (mutual exclusion, layer-wise takes priority)
|
|
if enable_layerwise_offload:
|
|
strategy = OffloadStrategy.LAYER_WISE
|
|
if enable_cpu_offload:
|
|
logger.info(
|
|
"Both model-level and layer-wise offloading enabled. "
|
|
"Layer-wise takes priority, disabling model-level offloading."
|
|
)
|
|
elif enable_cpu_offload:
|
|
strategy = OffloadStrategy.MODEL_LEVEL
|
|
else:
|
|
strategy = OffloadStrategy.NONE
|
|
|
|
return cls(
|
|
strategy=strategy,
|
|
pin_cpu_memory=pin_cpu_memory,
|
|
)
|
|
|
|
|
|
class OffloadBackend(ABC):
|
|
"""Base class for CPU offload backends"""
|
|
|
|
def __init__(self, config: OffloadConfig, device: torch.device):
|
|
self.config = config
|
|
self.device = device
|
|
self.enabled = False
|
|
|
|
@abstractmethod
|
|
def enable(self, pipeline: nn.Module) -> None:
|
|
"""Enable offloading on the pipeline.
|
|
|
|
Discovers modules, moves them to appropriate devices, and
|
|
registers forward hooks for swapping/prefetching.
|
|
|
|
Args:
|
|
pipeline: Diffusion pipeline model (e.g., Wan22Pipeline)
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def disable(self) -> None:
|
|
"""Disable offloading and cleanup resources.
|
|
|
|
Removes all registered hooks. Does NOT move modules back to
|
|
original devices (caller responsible for that).
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def is_enabled(self) -> bool:
|
|
return self.enabled
|