From f18924c1cad8d4aa0def12054d5717ef3c859f91 Mon Sep 17 00:00:00 2001 From: Alicia <115451386+congw729@users.noreply.github.com> Date: Wed, 22 Jul 2026 09:45:17 +0800 Subject: [PATCH] [Refactor][Misc] Remove unused code in diffusion hooks/model_loader (#5270) Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com> --- vllm_omni/diffusion/hooks/__init__.py | 6 -- vllm_omni/diffusion/hooks/base.py | 33 ++------ .../diffusion/hooks/sequence_parallel.py | 84 ------------------- .../model_loader/diffusers_loader.py | 11 --- .../diffusion/model_loader/hub_prefetch.py | 33 +------- .../models/qwen_image/pipeline_qwen_image.py | 5 +- .../qwen_image/pipeline_qwen_image_edit.py | 5 +- .../pipeline_qwen_image_edit_plus.py | 5 +- .../qwen_image/pipeline_qwen_image_layered.py | 5 +- 9 files changed, 17 insertions(+), 170 deletions(-) diff --git a/vllm_omni/diffusion/hooks/__init__.py b/vllm_omni/diffusion/hooks/__init__.py index 68f918383..ea1a28a6c 100644 --- a/vllm_omni/diffusion/hooks/__init__.py +++ b/vllm_omni/diffusion/hooks/__init__.py @@ -3,7 +3,6 @@ """Hook mechanism for model forward interception.""" from vllm_omni.diffusion.hooks.base import ( - BaseState, HookRegistry, ModelHook, StateManager, @@ -12,14 +11,11 @@ from vllm_omni.diffusion.hooks.sequence_parallel import ( SequenceParallelGatherHook, SequenceParallelSplitHook, apply_sequence_parallel, - disable_sequence_parallel_for_model, - enable_sequence_parallel_for_model, remove_sequence_parallel, ) __all__ = [ # Base hooks - "BaseState", "StateManager", "ModelHook", "HookRegistry", @@ -28,6 +24,4 @@ __all__ = [ "SequenceParallelGatherHook", "apply_sequence_parallel", "remove_sequence_parallel", - "enable_sequence_parallel_for_model", - "disable_sequence_parallel_for_model", ] diff --git a/vllm_omni/diffusion/hooks/base.py b/vllm_omni/diffusion/hooks/base.py index 6863a4594..8857737ca 100644 --- a/vllm_omni/diffusion/hooks/base.py +++ b/vllm_omni/diffusion/hooks/base.py @@ -20,21 +20,18 @@ from vllm.logger import init_logger logger = init_logger(__name__) -class BaseState: - """Base class for hook state containers.""" - - def reset(self) -> None: # pragma: no cover - default is no-op - pass - - class StateManager: - """Manage per-context hook state instances.""" + """Manage per-context hook state instances. - def __init__(self, state_cls: Callable[[], BaseState], init_args: tuple = (), init_kwargs: dict | None = None): + State classes (e.g. TeaCacheState, MagCacheState) are plain classes; + they only need to be constructible with the given init args. + """ + + def __init__(self, state_cls: Callable[..., Any], init_args: tuple = (), init_kwargs: dict | None = None): self._state_cls = state_cls self._init_args = init_args self._init_kwargs = init_kwargs or {} - self._states: dict[str, BaseState] = {} + self._states: dict[str, Any] = {} self._context: str = "default" @property @@ -45,7 +42,7 @@ class StateManager: def set_context(self, name: str) -> None: self._context = name or "default" - def get_state(self) -> BaseState: + def get_state(self) -> Any: if self._context not in self._states: self._states[self._context] = self._state_cls(*self._init_args, **self._init_kwargs) return self._states[self._context] @@ -332,17 +329,3 @@ class HookRegistry: hook = self._hooks.get(name) if hook is not None: hook.reset_state(self.module) - - def reset(self) -> None: - """Reset all hooks and clear the registry. - - This removes all hooks from the registry and resets each hook's state. - Also restores module.forward to its original implementation. - """ - for name, hook in list(self._hooks.items()): - hook.reset_state(self.module) - self._hooks.clear() - - if hasattr(self.module, "_omni_original_forward"): - self.module.forward = self.module._omni_original_forward # type: ignore[attr-defined] - delattr(self.module, "_omni_original_forward") diff --git a/vllm_omni/diffusion/hooks/sequence_parallel.py b/vllm_omni/diffusion/hooks/sequence_parallel.py index 0c6cda334..e8f88c8a9 100644 --- a/vllm_omni/diffusion/hooks/sequence_parallel.py +++ b/vllm_omni/diffusion/hooks/sequence_parallel.py @@ -679,87 +679,3 @@ def remove_sequence_parallel( continue registry.remove_hook(hook_name) - - -def enable_sequence_parallel_for_model( - model: nn.Module, - config: SequenceParallelConfig | None = None, -) -> None: - """Enable sequence parallelism for a model using its _sp_plan. - - This is a convenience function that reads the model's _sp_plan attribute - and applies sequence parallelism automatically. - - Note: This corresponds to `enable_context_parallel_for_model` in diffusers, - but uses vLLM-Omni's _sp_plan instead of diffusers' _cp_plan. - - The function performs two main tasks: - 1. Applies _sp_plan hooks to shard inputs and gather outputs - 2. Ensures Attention layers are configured for the correct parallel mode - (handled automatically by vLLM-Omni's forward_context mechanism) - - Args: - model: The model to enable SP for. Must have a _sp_plan attribute. - config: Optional config. If None, uses default based on current - parallel state. - - Raises: - ValueError: If model has no _sp_plan defined. - - Note: - vLLM-Omni supports Ulysses + Ring hybrid parallelism: - - ulysses_degree > 1: Uses All-to-All communication over Q/K/V heads - - ring_degree > 1: Uses Ring attention with K/V passing - - Both > 1: Hybrid mode (Ulysses handles head redistribution, - Ring handles K/V circulation) - """ - from vllm_omni.diffusion.distributed.parallel_state import ( - get_ring_parallel_world_size, - get_ulysses_parallel_world_size, - ) - from vllm_omni.diffusion.distributed.sp_plan import get_sp_plan_from_model - - plan = get_sp_plan_from_model(model) - if plan is None: - raise ValueError( - f"Model {model.__class__.__name__} has no _sp_plan defined. " - f"Define _sp_plan as a class attribute or pass a plan explicitly." - ) - - if config is None: - # Create config from current parallel state - ulysses_degree = get_ulysses_parallel_world_size() - ring_degree = get_ring_parallel_world_size() - config = SequenceParallelConfig( - ulysses_degree=ulysses_degree, - ring_degree=ring_degree, - ) - if ulysses_degree > 1 and ring_degree > 1: - mode = "hybrid" - elif ulysses_degree > 1: - mode = "ulysses" - else: - mode = "ring" - logger.info( - f"Created SP config from parallel state: " - f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, " - f"mode={mode}" - ) - - apply_sequence_parallel(model, config, plan) - logger.info(f"Enabled sequence parallelism for {model.__class__.__name__}") - - -def disable_sequence_parallel_for_model(model: nn.Module) -> None: - """Disable sequence parallelism for a model. - - Note: This corresponds to `disable_context_parallel_for_model` in diffusers. - - Args: - model: The model to disable SP for. - """ - from vllm_omni.diffusion.distributed.sp_plan import get_sp_plan_from_model - - plan = get_sp_plan_from_model(model) - if plan is not None: - remove_sequence_parallel(model, plan) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 8ee172e6b..bb048451a 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -12,7 +12,6 @@ from typing import cast import torch from torch import nn -from vllm.config import ModelConfig from vllm.config.load import LoadConfig from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase @@ -81,7 +80,6 @@ def _natural_sort_key(filepath: str) -> list: return [int(s) if s.isdigit() else s for s in re.split(r"(\d+)", os.path.basename(filepath))] -MODEL_INDEX = "model_index.json" DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json" TRANSFORMER_WEIGHTS_INDEX = "model.safetensors.index.json" INDEX_FILES = [DIFFUSION_MODEL_WEIGHTS_INDEX, TRANSFORMER_WEIGHTS_INDEX] @@ -326,15 +324,6 @@ class DiffusersPipelineLoader: return all_parameter_names return {name for name in all_parameter_names if name.startswith(source_prefixes)} - def download_model(self, model_config: ModelConfig) -> None: - self._prepare_weights( - model_name_or_path=model_config.model, - subfolder=None, - revision=model_config.revision, - fall_back_to_pt=True, - allow_patterns_overrides=None, - ) - def load_model( self, load_device: str, diff --git a/vllm_omni/diffusion/model_loader/hub_prefetch.py b/vllm_omni/diffusion/model_loader/hub_prefetch.py index a33dca74d..9a9295ab9 100644 --- a/vllm_omni/diffusion/model_loader/hub_prefetch.py +++ b/vllm_omni/diffusion/model_loader/hub_prefetch.py @@ -433,10 +433,7 @@ def from_pretrained_with_prefetch( ``factory`` is a bound ``SomeModel.from_pretrained`` (or any callable with the same ``(model, *, subfolder, local_files_only, **kwargs)`` signature). - This is a stronger sibling of :func:`retry_on_missing_shard`: that helper - only retries the missing-shard ``OSError`` and never re-prefetches, so it - cannot recover the second face of the same race. Two shapes of partial - -cache failure crash the diffusion server outright: + Two shapes of partial-cache failure crash the diffusion server outright: * ``OSError: does not appear to have a file named text_encoder/model-0000X-of-0000Y.safetensors`` - a shard is still under @@ -488,31 +485,3 @@ def from_pretrained_with_prefetch( assert last_exc is not None # loop only exits via return or a caught exc raise last_exc - - -def retry_on_missing_shard(load_fn, *, max_retries: int = 3, base_delay: float = 5.0): - """Call *load_fn* with retry on the transformers v5 shard-resolution race. - - When the prefetch lock cannot be acquired (e.g. flock unsupported on - the filesystem and dotfile lock times out), ``from_pretrained`` may - still hit the ``cached_files`` race. This wrapper retries with - exponential backoff when the OSError message matches the specific - "does not appear to have a file named" pattern. - """ - for attempt in range(max_retries): - try: - return load_fn() - except OSError as exc: - if "does not appear to have a file" not in str(exc): - raise - if attempt == max_retries - 1: - raise - delay = base_delay * (attempt + 1) - logger.warning( - "from_pretrained failed with shard-resolution race (%s); retrying in %.1fs (attempt %d/%d)", - exc, - delay, - attempt + 1, - max_retries, - ) - time.sleep(delay) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 986c2ef86..474ea946c 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -309,9 +309,8 @@ class QwenImagePipeline( model, subfolder="scheduler", local_files_only=local_files_only ) # ``from_pretrained_with_prefetch`` re-prefetches and retries on a - # half-written cache (missing-shard ``OSError`` *and* the default - # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard`` - # could not recover) instead of crashing the worker. + # half-written cache (missing-shard ``OSError`` and the default + # -config size-mismatch ``RuntimeError``) instead of crashing the worker. self.text_encoder = from_pretrained_with_prefetch( Qwen2_5_VLForConditionalGeneration.from_pretrained, model, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index 3d6e937c7..3f0d16002 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py @@ -263,9 +263,8 @@ class QwenImageEditPipeline( model, subfolder="scheduler", local_files_only=local_files_only ) # ``from_pretrained_with_prefetch`` re-prefetches and retries on a - # half-written cache (missing-shard ``OSError`` *and* the default - # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard`` - # could not recover) instead of crashing the worker. + # half-written cache (missing-shard ``OSError`` and the default + # -config size-mismatch ``RuntimeError``) instead of crashing the worker. self.text_encoder = from_pretrained_with_prefetch( Qwen2_5_VLForConditionalGeneration.from_pretrained, model, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index a11637163..8b5c4e3b1 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py @@ -230,9 +230,8 @@ class QwenImageEditPlusPipeline( model, subfolder="scheduler", local_files_only=local_files_only ) # ``from_pretrained_with_prefetch`` re-prefetches and retries on a - # half-written cache (missing-shard ``OSError`` *and* the default - # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard`` - # could not recover) instead of crashing the worker. + # half-written cache (missing-shard ``OSError`` and the default + # -config size-mismatch ``RuntimeError``) instead of crashing the worker. self.text_encoder = from_pretrained_with_prefetch( Qwen2_5_VLForConditionalGeneration.from_pretrained, model, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py index e2fe67976..830c4d6c1 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py @@ -237,9 +237,8 @@ class QwenImageLayeredPipeline( model, subfolder="scheduler", local_files_only=local_files_only ) # ``from_pretrained_with_prefetch`` re-prefetches and retries on a - # half-written cache (missing-shard ``OSError`` *and* the default - # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard`` - # could not recover) instead of crashing the worker. + # half-written cache (missing-shard ``OSError`` and the default + # -config size-mismatch ``RuntimeError``) instead of crashing the worker. self.text_encoder = from_pretrained_with_prefetch( Qwen2_5_VLForConditionalGeneration.from_pretrained, model,