768 lines
32 KiB
Python
768 lines
32 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project and the HuggingFace Team.
|
|
# All rights reserved.
|
|
#
|
|
# This module is adapted from HuggingFace diffusers library:
|
|
# diffusers/src/diffusers/hooks/context_parallel.py
|
|
#
|
|
# NOTE: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers.
|
|
# We use the term "Sequence Parallelism" to align with vLLM-Omni's existing terminology.
|
|
#
|
|
# Key adaptations from diffusers:
|
|
# - ModuleForwardMetadata: parameter lookup logic (adapted)
|
|
# - SequenceParallelSplitHook/GatherHook: hook structure (adapted from ContextParallel*)
|
|
# - apply_sequence_parallel: registration logic (adapted from apply_context_parallel)
|
|
#
|
|
# Key differences from diffusers:
|
|
# - Uses vLLM-Omni's SequenceParallelGroupCoordinator instead of DeviceMesh
|
|
# - Uses sp_shard/sp_gather from sp_sharding.py instead of funcol operations
|
|
# - Supports Ulysses + Ring hybrid parallelism
|
|
#
|
|
"""Sequence Parallelism hooks for non-intrusive SP support.
|
|
|
|
This module implements the hook-based mechanism for applying sequence parallelism
|
|
to models without modifying their forward() methods.
|
|
|
|
Usage:
|
|
1. Define _sp_plan on your model class (corresponds to diffusers' _cp_plan)
|
|
2. Call apply_sequence_parallel(model, config, plan) to enable SP
|
|
3. Call remove_sequence_parallel(model, plan) to disable SP
|
|
|
|
The hooks automatically shard inputs before forward and gather outputs after,
|
|
based on the plan specification.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import inspect
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.distributed.sp_plan import (
|
|
AnySequenceParallelInput,
|
|
SequenceParallelConfig,
|
|
SequenceParallelInput,
|
|
SequenceParallelInputType,
|
|
SequenceParallelModelPlan,
|
|
SequenceParallelOutput,
|
|
SequenceParallelOutputType,
|
|
SequenceParallelPartialInput,
|
|
)
|
|
from vllm_omni.diffusion.distributed.sp_sharding import sp_gather, sp_shard
|
|
from vllm_omni.diffusion.hooks.base import HookRegistry, ModelHook
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# Hook name templates for identifying SP hooks
|
|
_SP_INPUT_HOOK_TEMPLATE = "sp_input---{}"
|
|
_SP_OUTPUT_HOOK_TEMPLATE = "sp_output---{}"
|
|
|
|
|
|
@dataclass
|
|
class ModuleForwardMetadata:
|
|
"""Metadata for mapping forward() parameter names to args/kwargs positions.
|
|
|
|
This caches the inspection of a module's forward signature to efficiently
|
|
locate parameters by name in subsequent calls.
|
|
"""
|
|
|
|
cached_parameter_indices: dict[str, int] | None = None
|
|
_cls: type[nn.Module] | None = None
|
|
|
|
def _get_parameter_from_args_kwargs(
|
|
self,
|
|
identifier: str | int,
|
|
args: Sequence[Any] = (),
|
|
kwargs: dict[str, Any] | None = None,
|
|
) -> tuple[Any, bool, int | None]:
|
|
"""Get a parameter value from args or kwargs by name.
|
|
|
|
Args:
|
|
identifier: The parameter name to look up, or a positional index.
|
|
args: Positional arguments passed to forward.
|
|
kwargs: Keyword arguments passed to forward.
|
|
|
|
Returns:
|
|
Tuple of (value, is_kwarg, index).
|
|
- value: The parameter value (or None if not found)
|
|
- is_kwarg: True if found in kwargs
|
|
- index: Position in args if found there
|
|
|
|
Raises:
|
|
ValueError: If parameter not found in signature.
|
|
"""
|
|
kwargs = kwargs or {}
|
|
|
|
if isinstance(identifier, int):
|
|
if identifier < len(args):
|
|
return args[identifier], False, identifier
|
|
return None, False, identifier
|
|
|
|
# First check kwargs
|
|
if identifier in kwargs:
|
|
return kwargs[identifier], True, None
|
|
|
|
# Check cached indices
|
|
if self.cached_parameter_indices is not None:
|
|
index = self.cached_parameter_indices.get(identifier, None)
|
|
if index is None:
|
|
raise ValueError(f"Parameter '{identifier}' not found in cached indices.")
|
|
if index < len(args):
|
|
return args[index], False, index
|
|
return None, False, index
|
|
|
|
# Build cache from forward signature
|
|
if self._cls is None:
|
|
raise ValueError("Model class is not set for metadata.")
|
|
|
|
parameters = list(inspect.signature(self._cls.forward).parameters.keys())
|
|
parameters = parameters[1:] # Skip `self`
|
|
self.cached_parameter_indices = {param: i for i, param in enumerate(parameters)}
|
|
|
|
if identifier not in self.cached_parameter_indices:
|
|
raise ValueError(f"Parameter '{identifier}' not found in function signature.")
|
|
|
|
index = self.cached_parameter_indices[identifier]
|
|
|
|
if index >= len(args):
|
|
return None, False, index
|
|
|
|
return args[index], False, index
|
|
|
|
|
|
def _unwrap_module(module: nn.Module) -> nn.Module:
|
|
"""Unwrap a module from any wrappers to get the original class.
|
|
|
|
Args:
|
|
module: Potentially wrapped module.
|
|
|
|
Returns:
|
|
The unwrapped module.
|
|
"""
|
|
# Handle common wrappers
|
|
while hasattr(module, "_modules") and len(module._modules) == 1:
|
|
inner = next(iter(module._modules.values()))
|
|
if inner is not None:
|
|
module = inner
|
|
else:
|
|
break
|
|
return module
|
|
|
|
|
|
class SequenceParallelSplitHook(ModelHook):
|
|
"""Hook for splitting inputs before a module's forward pass.
|
|
|
|
This hook is registered to modules that need their inputs sharded
|
|
across sequence parallel ranks. It intercepts the forward call,
|
|
shards specified inputs according to the plan, and passes the
|
|
sharded inputs to the original forward.
|
|
|
|
For split_output=True inputs, it shards the output instead.
|
|
|
|
Supports both SequenceParallelInput (full split) and SequenceParallelPartialInput
|
|
(partial split for text/image separation).
|
|
|
|
Note: This corresponds to `ContextParallelSplitHook` in diffusers.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
metadata: SequenceParallelInputType,
|
|
config: SequenceParallelConfig,
|
|
) -> None:
|
|
super().__init__()
|
|
self.metadata = metadata
|
|
self.config = config
|
|
self.module_forward_metadata: ModuleForwardMetadata | None = None
|
|
# Cache for text lengths resolved from kwargs
|
|
self._text_len_cache: dict[str, int] = {}
|
|
self._shard_groups = {
|
|
sp_input.shard_group
|
|
for value in metadata.values()
|
|
for sp_input in (value if isinstance(value, (list, tuple)) else (value,))
|
|
if isinstance(sp_input, SequenceParallelInput) and sp_input.shard_group is not None
|
|
}
|
|
|
|
def initialize_hook(self, module: nn.Module) -> nn.Module:
|
|
cls = _unwrap_module(module).__class__
|
|
self.module_forward_metadata = ModuleForwardMetadata(_cls=cls)
|
|
return module
|
|
|
|
def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
|
|
"""Shard inputs before forward."""
|
|
args_list = list(args)
|
|
# Clear text length cache for this forward pass
|
|
self._text_len_cache.clear()
|
|
if self._shard_groups:
|
|
from vllm_omni.diffusion.forward_context import (
|
|
get_forward_context,
|
|
is_forward_context_available,
|
|
)
|
|
|
|
if is_forward_context_available():
|
|
ctx = get_forward_context()
|
|
for shard_group in self._shard_groups:
|
|
ctx.sp_shard_metadata.pop(shard_group, None)
|
|
|
|
for name, spm in self.metadata.items():
|
|
# Skip if this is a split_output entry (handled in post_forward)
|
|
if isinstance(spm, (SequenceParallelInput, SequenceParallelPartialInput)) and spm.split_output:
|
|
continue
|
|
|
|
# Get the parameter value
|
|
if self.module_forward_metadata is None:
|
|
raise RuntimeError("SequenceParallelSplitHook.initialize_hook must run before forward")
|
|
input_val, is_kwarg, index = self.module_forward_metadata._get_parameter_from_args_kwargs(
|
|
name, args_list, kwargs
|
|
)
|
|
|
|
if input_val is None:
|
|
continue
|
|
|
|
# Shard the input
|
|
if isinstance(input_val, torch.Tensor):
|
|
if not isinstance(spm, (SequenceParallelInput, SequenceParallelPartialInput)):
|
|
raise ValueError(
|
|
f"Expected SequenceParallelInput for tensor parameter '{name}', got {type(spm).__name__}"
|
|
)
|
|
input_val = self._prepare_sp_input(input_val, spm, args_list, kwargs)
|
|
elif isinstance(input_val, (list, tuple)):
|
|
# Handle list/tuple of tensors with per-element config
|
|
if not isinstance(spm, (list, tuple)):
|
|
raise ValueError(
|
|
f"Expected list/tuple of SequenceParallelInput for parameter '{name}' "
|
|
f"which is a list/tuple, but got {type(spm).__name__}"
|
|
)
|
|
if len(input_val) != len(spm):
|
|
raise ValueError(f"Expected {len(spm)} elements for parameter '{name}', got {len(input_val)}")
|
|
sharded_input_val = []
|
|
for i, x in enumerate(input_val):
|
|
if torch.is_tensor(x) and not spm[i].split_output:
|
|
x = self._prepare_sp_input(x, spm[i], args_list, kwargs)
|
|
sharded_input_val.append(x)
|
|
input_val = type(input_val)(sharded_input_val)
|
|
else:
|
|
raise ValueError(f"Unsupported input type for sharding: {type(input_val).__name__}")
|
|
|
|
# Update args or kwargs
|
|
if is_kwarg:
|
|
if not isinstance(name, str):
|
|
raise TypeError(f"Keyword parameter name must be str, got {type(name).__name__}")
|
|
kwargs[name] = input_val
|
|
elif index is not None and index < len(args_list):
|
|
args_list[index] = input_val
|
|
else:
|
|
raise ValueError(f"Failed to update parameter '{name}' after sharding.")
|
|
|
|
# Output splitting only needs integer text lengths. Saving args/kwargs
|
|
# here would keep full hidden-state inputs alive after output sharding,
|
|
# including when the module raises before post_forward can run.
|
|
for spm in self.metadata.values():
|
|
if isinstance(spm, SequenceParallelPartialInput) and spm.split_output:
|
|
self._resolve_text_len(spm, args_list, kwargs)
|
|
|
|
return tuple(args_list), kwargs
|
|
|
|
def post_forward(self, module: nn.Module, output: Any) -> Any:
|
|
"""Shard outputs for split_output=True entries."""
|
|
from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
|
|
|
|
is_tensor = isinstance(output, torch.Tensor)
|
|
is_tensor_list = isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output)
|
|
|
|
if not is_tensor and not is_tensor_list:
|
|
# No tensor outputs to shard
|
|
logger.warning_once("No tensor outputs to shard.")
|
|
return output
|
|
|
|
output_list = [output] if is_tensor else list(output)
|
|
actually_sharded = False
|
|
equal_rank_seq_lens = True
|
|
|
|
for index, spm in self.metadata.items():
|
|
if not isinstance(index, int):
|
|
continue
|
|
if not isinstance(spm, (SequenceParallelInput, SequenceParallelPartialInput)) or not spm.split_output:
|
|
continue
|
|
if index >= len(output_list):
|
|
raise ValueError(f"Index {index} out of bounds for output of length {len(output_list)}.")
|
|
|
|
original = output_list[index]
|
|
output_list[index] = self._prepare_sp_input(original, spm)
|
|
if output_list[index] is not original:
|
|
actually_sharded = True
|
|
equal_rank_seq_lens &= isinstance(spm, SequenceParallelInput) and spm.auto_pad
|
|
|
|
# Mark SP as active only if at least one tensor was actually sharded
|
|
if actually_sharded and is_forward_context_available():
|
|
ctx = get_forward_context()
|
|
ctx._sp_shard_depth += 1
|
|
ctx._sp_equal_pad_stack.append(equal_rank_seq_lens)
|
|
|
|
return output_list[0] if is_tensor else type(output)(output_list)
|
|
|
|
def _resolve_text_len(
|
|
self,
|
|
sp_input: SequenceParallelPartialInput,
|
|
args: Sequence[Any],
|
|
kwargs: dict[str, Any],
|
|
) -> int:
|
|
"""Resolve text length from the source specification."""
|
|
source = sp_input.text_len_source
|
|
|
|
if isinstance(source, int):
|
|
return source
|
|
|
|
# String source - look up from kwargs or args
|
|
if source in self._text_len_cache:
|
|
return self._text_len_cache[source]
|
|
|
|
# Try to get from kwargs/args
|
|
try:
|
|
if self.module_forward_metadata is None:
|
|
raise RuntimeError("SequenceParallelSplitHook.initialize_hook must run before forward")
|
|
val, _, _ = self.module_forward_metadata._get_parameter_from_args_kwargs(source, args, kwargs)
|
|
if val is None:
|
|
raise ValueError(f"Parameter '{source}' is None, cannot determine text length.")
|
|
if isinstance(val, torch.Tensor):
|
|
# TODO: Currently assumes batch_size=1, where shape[0] is sequence length.
|
|
# For batch inference support, this should be updated to handle
|
|
# shape (batch_size, seq_len, ...) where text_len varies per sample.
|
|
text_len = val.shape[0]
|
|
elif isinstance(val, int):
|
|
text_len = val
|
|
else:
|
|
raise ValueError(f"Cannot determine text length from '{source}' of type {type(val).__name__}")
|
|
self._text_len_cache[source] = text_len
|
|
return text_len
|
|
except ValueError as e:
|
|
raise ValueError(f"Failed to resolve text_len_source '{source}': {e}") from e
|
|
|
|
def _prepare_sp_input(
|
|
self,
|
|
x: torch.Tensor,
|
|
sp_input: AnySequenceParallelInput,
|
|
args: Sequence[Any] = (),
|
|
kwargs: dict[str, Any] | None = None,
|
|
) -> torch.Tensor:
|
|
"""Shard a tensor according to the input specification."""
|
|
kwargs = kwargs or {}
|
|
|
|
if sp_input.expected_dims is not None and x.dim() != sp_input.expected_dims:
|
|
logger.warning_once(f"Expected tensor with {sp_input.expected_dims} dims, got {x.dim()}. Skipping split.")
|
|
return x
|
|
|
|
split_dim = (
|
|
sp_input.split_dim if isinstance(sp_input, (SequenceParallelInput, SequenceParallelPartialInput)) else None
|
|
)
|
|
if split_dim is not None and x.size(split_dim) == 0:
|
|
logger.warning_once("Skip sharding for zero-sized tensors.")
|
|
return x
|
|
|
|
def _raise_strict_divisibility_error(*, dim: int, seq_len: int, sp_size: int) -> None:
|
|
# Keep message actionable: strict mode must be evenly shardable at the split hook level.
|
|
msg = (
|
|
"Sequence Parallel strict mode requires the sequence length to be evenly shardable at split time, "
|
|
f"but got seq_len={seq_len} (dim={dim}) not divisible by sequence_parallel_size={sp_size}. "
|
|
)
|
|
suggestions: list[str] = []
|
|
if self.config.ulysses_degree > 1:
|
|
suggestions.append("set parallel_config.ulysses_mode='advanced_uaa' (UAA)")
|
|
suggestions.append(
|
|
"enable auto_pad=True in the model _sp_plan (requires attention_mask support and ring_degree=1)"
|
|
)
|
|
suggestions.append("choose seq_len divisible by sequence_parallel_size (or adjust SP degrees)")
|
|
msg += "Try: " + "; ".join(suggestions) + "."
|
|
raise ValueError(msg)
|
|
|
|
def _maybe_validate_strict_divisibility(*, dim: int, seq_len: int) -> None:
|
|
from vllm_omni.diffusion.forward_context import get_ulysses_mode
|
|
|
|
if get_ulysses_mode(default="strict") != "strict":
|
|
return
|
|
sp_size = int(self.config.sequence_parallel_size)
|
|
if sp_size <= 1:
|
|
return
|
|
if seq_len < sp_size or (seq_len % sp_size) != 0:
|
|
_raise_strict_divisibility_error(dim=dim, seq_len=int(seq_len), sp_size=sp_size)
|
|
|
|
if isinstance(sp_input, SequenceParallelInput):
|
|
# Full split with optional auto-padding
|
|
if sp_input.auto_pad:
|
|
shard = self._shard_with_auto_pad(
|
|
x,
|
|
sp_input.split_dim,
|
|
sp_input.shard_group,
|
|
)
|
|
else:
|
|
_maybe_validate_strict_divisibility(dim=sp_input.split_dim, seq_len=x.size(sp_input.split_dim))
|
|
shard = sp_shard(x, sp_input.split_dim, validate=False)
|
|
# contiguous() may be a no-op for a single-batch sequence slice;
|
|
# clone explicitly to release its full (possibly padded) backing.
|
|
if sp_input.clone_shard:
|
|
from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size
|
|
|
|
if get_sequence_parallel_world_size() > 1:
|
|
return shard.clone(memory_format=torch.contiguous_format)
|
|
return shard
|
|
elif isinstance(sp_input, SequenceParallelPartialInput):
|
|
# Partial split: keep text portion, split image portion
|
|
text_len = self._resolve_text_len(sp_input, args, kwargs)
|
|
dim = sp_input.split_dim
|
|
|
|
# Split tensor into text and image portions
|
|
text_part = x.narrow(dim, 0, text_len)
|
|
image_part = x.narrow(dim, text_len, x.size(dim) - text_len)
|
|
|
|
# Only shard the image portion
|
|
_maybe_validate_strict_divisibility(dim=dim, seq_len=image_part.size(dim))
|
|
image_part_sharded = sp_shard(image_part, dim, validate=False)
|
|
|
|
# Concatenate back: [text_full, image_sharded]
|
|
return torch.cat([text_part, image_part_sharded], dim=dim)
|
|
else:
|
|
raise ValueError(f"Unsupported input config type: {type(sp_input).__name__}")
|
|
|
|
def _shard_with_auto_pad(
|
|
self,
|
|
x: torch.Tensor,
|
|
dim: int,
|
|
shard_group: str | None,
|
|
) -> torch.Tensor:
|
|
"""Pad and shard a tensor while recording its global sequence length.
|
|
|
|
Models with several independent sequences read the recorded length back
|
|
via `shard_group` to build masks. Legacy singleton padding fields are
|
|
also populated for single-sequence models.
|
|
"""
|
|
from vllm_omni.diffusion.attention.selector import get_attn_backend_for_capability
|
|
from vllm_omni.diffusion.distributed.parallel_state import (
|
|
get_ring_parallel_world_size,
|
|
get_sequence_parallel_rank,
|
|
get_sequence_parallel_world_size,
|
|
)
|
|
from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
|
|
|
|
world_size = get_sequence_parallel_world_size()
|
|
if world_size == 1:
|
|
return x
|
|
|
|
seq_len = x.size(dim)
|
|
rank = get_sequence_parallel_rank()
|
|
remainder = seq_len % world_size
|
|
|
|
if shard_group is not None and is_forward_context_available():
|
|
ctx = get_forward_context()
|
|
existing = ctx.sp_shard_metadata.setdefault(shard_group, seq_len)
|
|
if existing != seq_len:
|
|
raise ValueError(
|
|
f"All tensors in shard_group={shard_group!r} must have the "
|
|
f"same global sequence length, but got {existing} and {seq_len}."
|
|
)
|
|
|
|
if remainder == 0:
|
|
# Keyed groups record their length even when no padding is needed.
|
|
return sp_shard(x, dim, validate=False)
|
|
|
|
# Check backend compatibility
|
|
attention_config = None
|
|
if is_forward_context_available():
|
|
od_config = get_forward_context().omni_diffusion_config
|
|
if od_config is not None:
|
|
attention_config = od_config.diffusion_attention_config
|
|
|
|
attn_backend = get_attn_backend_for_capability(
|
|
role="self",
|
|
attention_config=attention_config,
|
|
)
|
|
attention_spec = None
|
|
if attention_config is not None:
|
|
attention_spec, _ = attention_config.resolve_with_source(role="self")
|
|
if not attn_backend.supports_attention_mask(attention_spec):
|
|
raise ValueError(
|
|
f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). "
|
|
f"Cannot use {attn_backend.get_name()} which does not support attention_mask. "
|
|
f"Please switch to SDPA or Ascend attention backend."
|
|
)
|
|
|
|
# Ring attention does not support attention_mask
|
|
if get_ring_parallel_world_size() > 1:
|
|
raise ValueError(
|
|
f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). "
|
|
f"Cannot use Ring attention which does not support attention_mask. "
|
|
f"Please switch to Ulysses SP only."
|
|
)
|
|
|
|
# Calculate padding
|
|
pad_size = world_size - remainder
|
|
padded_seq_len = seq_len + pad_size
|
|
|
|
# Pad the tensor
|
|
pad_shape = list(x.shape)
|
|
pad_shape[dim] = pad_size
|
|
padding = torch.zeros(pad_shape, dtype=x.dtype, device=x.device)
|
|
x_padded = torch.cat([x, padding], dim=dim)
|
|
|
|
# Store padding info in forward context (only once, for primary tensor)
|
|
# Attention layers will create masks dynamically using this info.
|
|
# Keyed groups record their length in sp_shard_metadata above and must
|
|
# not clobber the legacy singleton fields, whose semantics assume a
|
|
# single sequence per boundary.
|
|
if shard_group is None and is_forward_context_available():
|
|
ctx = get_forward_context()
|
|
# Only set if not already set (first auto_pad tensor wins)
|
|
if ctx.sp_original_seq_len is None:
|
|
ctx.sp_padding_size = pad_size
|
|
ctx.sp_original_seq_len = seq_len
|
|
logger.debug(
|
|
f"Auto-padded sequence from {seq_len} to {padded_seq_len} "
|
|
f"(pad_size={pad_size}, world_size={world_size}, dim={dim})"
|
|
)
|
|
|
|
# Shard the padded tensor
|
|
return x_padded.chunk(world_size, dim=dim)[rank]
|
|
|
|
|
|
class SequenceParallelGatherHook(ModelHook):
|
|
"""Hook for gathering outputs after a module's forward pass.
|
|
|
|
This hook is registered to modules that need their outputs gathered
|
|
from all sequence parallel ranks. It intercepts the output and gathers
|
|
it according to the plan specification.
|
|
|
|
Note: This corresponds to `ContextParallelGatherHook` in diffusers.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
metadata: SequenceParallelOutputType,
|
|
config: SequenceParallelConfig,
|
|
) -> None:
|
|
super().__init__()
|
|
if isinstance(metadata, SequenceParallelOutput):
|
|
metadata = [metadata]
|
|
elif isinstance(metadata, tuple):
|
|
metadata = list(metadata)
|
|
self.metadata = metadata
|
|
self.config = config
|
|
|
|
def initialize_hook(self, module: nn.Module) -> nn.Module:
|
|
return module
|
|
|
|
def post_forward(self, module: nn.Module, output: Any) -> Any:
|
|
"""Gather outputs after forward and remove padding if applied."""
|
|
from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
|
|
|
|
is_tensor = isinstance(output, torch.Tensor)
|
|
|
|
if is_tensor:
|
|
output = [output]
|
|
elif not (isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output)):
|
|
raise ValueError(f"Expected tensor or list/tuple of tensors, got {type(output).__name__}")
|
|
|
|
output = list(output)
|
|
|
|
if len(output) != len(self.metadata):
|
|
raise ValueError(f"Expected {len(self.metadata)} outputs, got {len(output)}.")
|
|
|
|
actually_gathered = False
|
|
|
|
for i, spm in enumerate(self.metadata):
|
|
if spm is None:
|
|
continue
|
|
|
|
x = output[i]
|
|
if spm.expected_dims is not None and x.dim() != spm.expected_dims:
|
|
logger.warning_once(
|
|
f"Expected output tensor with {spm.expected_dims} dims, got {x.dim()}. Skipping gather."
|
|
)
|
|
continue
|
|
|
|
# Gather from all ranks
|
|
gathered = sp_gather(x, spm.gather_dim, validate=False)
|
|
|
|
# Remove padding if it was applied
|
|
original_seq_len = None
|
|
if is_forward_context_available():
|
|
ctx = get_forward_context()
|
|
if spm.shard_group is None:
|
|
original_seq_len = ctx.sp_original_seq_len
|
|
else:
|
|
original_seq_len = ctx.sp_shard_metadata.get(spm.shard_group)
|
|
if original_seq_len is None:
|
|
raise RuntimeError(f"No SP shard metadata was registered for shard_group={spm.shard_group!r}.")
|
|
if original_seq_len is not None and gathered.size(spm.gather_dim) > original_seq_len:
|
|
gathered = gathered.narrow(spm.gather_dim, 0, original_seq_len)
|
|
logger.debug(f"Removed padding: gathered shape {gathered.shape} (original_seq_len={original_seq_len})")
|
|
|
|
output[i] = gathered
|
|
actually_gathered = True
|
|
|
|
# Mark SP as inactive only if at least one tensor was actually gathered
|
|
if actually_gathered and is_forward_context_available():
|
|
ctx = get_forward_context()
|
|
ctx._sp_shard_depth = max(0, ctx._sp_shard_depth - 1)
|
|
if ctx._sp_equal_pad_stack:
|
|
ctx._sp_equal_pad_stack.pop()
|
|
for spm in self.metadata:
|
|
if spm is not None and spm.shard_group is not None:
|
|
ctx.sp_shard_metadata.pop(spm.shard_group, None)
|
|
|
|
return output[0] if is_tensor else type(output)(output)
|
|
|
|
|
|
def _get_submodule_by_name(model: nn.Module, name: str) -> nn.Module | list[nn.Module]:
|
|
"""Get a submodule by dotted name, supporting wildcards.
|
|
|
|
Args:
|
|
model: The root module.
|
|
name: Dotted path to submodule. Use "*" to match all children
|
|
of a ModuleList.
|
|
|
|
Returns:
|
|
The submodule or list of submodules if wildcard used.
|
|
|
|
Raises:
|
|
ValueError: If the path is invalid or module not found.
|
|
"""
|
|
if name.count("*") > 1:
|
|
raise ValueError("Wildcard '*' can only be used once in the name")
|
|
return _find_submodule_by_name(model, name)
|
|
|
|
|
|
def _find_submodule_by_name(model: nn.Module, name: str) -> nn.Module | list[nn.Module]:
|
|
"""Recursive helper for _get_submodule_by_name."""
|
|
if name == "":
|
|
return model
|
|
|
|
first_atom, remaining_name = name.split(".", 1) if "." in name else (name, "")
|
|
|
|
if first_atom == "*":
|
|
if not isinstance(model, nn.ModuleList):
|
|
raise ValueError("Wildcard '*' can only be used with ModuleList")
|
|
submodules = []
|
|
for submodule in model:
|
|
subsubmodules = _find_submodule_by_name(submodule, remaining_name)
|
|
if not isinstance(subsubmodules, list):
|
|
subsubmodules = [subsubmodules]
|
|
submodules.extend(subsubmodules)
|
|
return submodules
|
|
else:
|
|
if hasattr(model, first_atom):
|
|
submodule = getattr(model, first_atom)
|
|
return _find_submodule_by_name(submodule, remaining_name)
|
|
else:
|
|
raise ValueError(f"'{first_atom}' is not a submodule of '{model.__class__.__name__}'")
|
|
|
|
|
|
def apply_sequence_parallel(
|
|
module: nn.Module,
|
|
config: SequenceParallelConfig,
|
|
plan: SequenceParallelModelPlan,
|
|
) -> None:
|
|
"""Apply sequence parallel hooks to a model according to the plan.
|
|
|
|
This function registers hooks on the specified submodules to automatically
|
|
shard inputs and gather outputs for sequence parallelism.
|
|
|
|
Note: This corresponds to `apply_context_parallel` in diffusers.
|
|
|
|
The complete SP flow is:
|
|
1. Input sharding (SequenceParallelSplitHook): Split sequence across SP ranks
|
|
2. Attention parallelism (handled by vLLM-Omni's Attention layer):
|
|
- Ulysses: All-to-All over Q/K/V heads
|
|
- Ring: K/V circulation in ring topology
|
|
- Hybrid: Both (Ulysses handles head redistribution, Ring handles K/V)
|
|
3. Output gathering (SequenceParallelGatherHook): Gather sequence from SP ranks
|
|
|
|
Args:
|
|
module: The model to apply SP to.
|
|
config: The sequence parallel configuration.
|
|
plan: Dictionary mapping module names to input/output specifications.
|
|
|
|
Example:
|
|
config = SequenceParallelConfig(ulysses_degree=2)
|
|
plan = {
|
|
"": {"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3)},
|
|
"proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
|
|
}
|
|
apply_sequence_parallel(model, config, plan)
|
|
|
|
Note:
|
|
vLLM-Omni's Attention layer automatically handles the internal
|
|
parallelism (Ulysses All-to-All or Ring attention) based on the
|
|
forward_context configuration. This function only handles input/output
|
|
sharding for the model as a whole.
|
|
"""
|
|
logger.debug(
|
|
f"Applying sequence parallel with config: ulysses={config.ulysses_degree}, "
|
|
f"ring={config.ring_degree}, plan keys: {list(plan.keys())}"
|
|
)
|
|
|
|
for module_id, sp_model_plan in plan.items():
|
|
submodule = _get_submodule_by_name(module, module_id)
|
|
if not isinstance(submodule, list):
|
|
submodule = [submodule]
|
|
|
|
logger.debug(f"Applying SP hooks to '{module_id}' ({len(submodule)} module(s))")
|
|
|
|
for m in submodule:
|
|
hook: SequenceParallelSplitHook | SequenceParallelGatherHook
|
|
if isinstance(sp_model_plan, dict):
|
|
# Input specification
|
|
hook = SequenceParallelSplitHook(sp_model_plan, config)
|
|
hook_name = _SP_INPUT_HOOK_TEMPLATE.format(module_id)
|
|
elif isinstance(sp_model_plan, (SequenceParallelOutput, list, tuple)):
|
|
# Output specification
|
|
if isinstance(sp_model_plan, SequenceParallelOutput):
|
|
sp_model_plan = [sp_model_plan]
|
|
elif isinstance(sp_model_plan, tuple):
|
|
sp_model_plan = list(sp_model_plan)
|
|
if not all(isinstance(x, SequenceParallelOutput) or x is None for x in sp_model_plan):
|
|
raise ValueError(f"Expected SequenceParallelOutput elements, got {sp_model_plan}")
|
|
hook = SequenceParallelGatherHook(sp_model_plan, config)
|
|
hook_name = _SP_OUTPUT_HOOK_TEMPLATE.format(module_id)
|
|
else:
|
|
raise ValueError(f"Unsupported plan type: {type(sp_model_plan).__name__}")
|
|
|
|
registry = HookRegistry.get_or_create(m)
|
|
registry.register_hook(hook_name, hook)
|
|
|
|
|
|
def remove_sequence_parallel(
|
|
module: nn.Module,
|
|
plan: SequenceParallelModelPlan,
|
|
) -> None:
|
|
"""Remove sequence parallel hooks from a model.
|
|
|
|
Note: This corresponds to `remove_context_parallel` in diffusers.
|
|
|
|
Args:
|
|
module: The model to remove SP from.
|
|
plan: The same plan used when applying SP.
|
|
"""
|
|
for module_id, sp_model_plan in plan.items():
|
|
submodule = _get_submodule_by_name(module, module_id)
|
|
if not isinstance(submodule, list):
|
|
submodule = [submodule]
|
|
|
|
for m in submodule:
|
|
registry = getattr(m, "_hook_registry", None)
|
|
if registry is None:
|
|
continue
|
|
|
|
if isinstance(sp_model_plan, dict):
|
|
hook_name = _SP_INPUT_HOOK_TEMPLATE.format(module_id)
|
|
elif isinstance(sp_model_plan, (SequenceParallelOutput, list, tuple)):
|
|
hook_name = _SP_OUTPUT_HOOK_TEMPLATE.format(module_id)
|
|
else:
|
|
continue
|
|
|
|
registry.remove_hook(hook_name)
|