143 lines
4.8 KiB
Python
143 lines
4.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from typing import Protocol
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
from vllm.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class SpatialShardExecutor(Protocol):
|
|
group: dist.ProcessGroup
|
|
parallel_mode: str
|
|
parallel_size: int
|
|
|
|
def _sync_final_result(
|
|
self,
|
|
rank0_result: torch.Tensor,
|
|
output_ndim: int,
|
|
output_device: torch.device,
|
|
output_dtype: torch.dtype,
|
|
) -> torch.Tensor: ...
|
|
|
|
|
|
class SpatialShardVAE(Protocol):
|
|
distributed_executor: SpatialShardExecutor
|
|
decoder: torch.nn.Module
|
|
post_quant_conv: torch.nn.Module
|
|
dtype: torch.dtype
|
|
_conv_idx: list[int]
|
|
_feat_map: list[object]
|
|
|
|
def is_distributed_enabled(self) -> bool: ...
|
|
|
|
def clear_cache(self) -> None: ...
|
|
|
|
|
|
def select_auto_spatial_shard_split_dim(
|
|
z_shape: tuple[int, ...],
|
|
world_size: int,
|
|
) -> str | None:
|
|
"""Select the longer spatial axis that can be split across every rank."""
|
|
if len(z_shape) != 5 or world_size <= 1:
|
|
return None
|
|
|
|
_, _, _, height, width = z_shape
|
|
candidates = [
|
|
(extent, split_dim) for extent, split_dim in ((height, "height"), (width, "width")) if extent >= world_size
|
|
]
|
|
if not candidates:
|
|
return None
|
|
# Prefer width for square latents. It was modestly faster in the original
|
|
# four-GPU Wan benchmark and keeps the choice deterministic across models.
|
|
return max(candidates, key=lambda candidate: (candidate[0], candidate[1] == "width"))[1]
|
|
|
|
|
|
def resolve_spatial_shard_split_dim(
|
|
mode: str,
|
|
z: torch.Tensor | None = None,
|
|
world_size: int | None = None,
|
|
) -> str | None:
|
|
if mode == "spatial_shard_width":
|
|
return "width"
|
|
if mode == "spatial_shard_height":
|
|
return "height"
|
|
if mode == "auto" and z is not None and world_size is not None:
|
|
return select_auto_spatial_shard_split_dim(tuple(z.shape), world_size)
|
|
return None
|
|
|
|
|
|
def spatial_shard_decode_enabled(vae: SpatialShardVAE, z: torch.Tensor, *, model_name: str) -> bool:
|
|
"""Apply the shared request- and topology-level spatial decode policy."""
|
|
executor = vae.distributed_executor
|
|
if executor.parallel_mode == "tile":
|
|
return False
|
|
if z.ndim != 5:
|
|
logger.warning(
|
|
"%s VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.", model_name
|
|
)
|
|
return False
|
|
if not vae.is_distributed_enabled():
|
|
return False
|
|
|
|
group = executor.group
|
|
world_size = dist.get_world_size(group=group)
|
|
requested_size = int(executor.parallel_size)
|
|
requested_split_dim = resolve_spatial_shard_split_dim(executor.parallel_mode)
|
|
if requested_size != world_size:
|
|
if executor.parallel_mode == "auto":
|
|
logger.debug(
|
|
"%s VAE auto decode selected tile mode for a partial DiT group (requested=%s dit_group=%s)",
|
|
model_name,
|
|
requested_size,
|
|
world_size,
|
|
)
|
|
elif requested_split_dim is not None:
|
|
logger.warning(
|
|
"%s VAE spatial sharded decode currently requires vae_patch_parallel_size "
|
|
"to match the DIT group size; falling back to tiled decode. "
|
|
"requested=%s dit_group=%s split_dim=%s",
|
|
model_name,
|
|
requested_size,
|
|
world_size,
|
|
requested_split_dim,
|
|
)
|
|
return False
|
|
return resolve_spatial_shard_split_dim(executor.parallel_mode, z, world_size) is not None
|
|
|
|
|
|
def prepare_spatial_shard_decode(
|
|
vae: SpatialShardVAE,
|
|
*,
|
|
install: Callable[[SpatialShardVAE, dist.ProcessGroup, str], None],
|
|
) -> None:
|
|
"""Install one VAE's dynamic wrappers while allocations are still tagged."""
|
|
executor = vae.distributed_executor
|
|
mode = executor.parallel_mode
|
|
if mode not in {"auto", "spatial_shard_height", "spatial_shard_width"}:
|
|
return
|
|
requested_size = int(executor.parallel_size)
|
|
world_size = dist.get_world_size(group=executor.group)
|
|
if requested_size <= 1 or requested_size != world_size:
|
|
return
|
|
split_dim = "width" if mode == "spatial_shard_width" else "height"
|
|
install(vae, executor.group, split_dim)
|
|
|
|
|
|
def prepare_pipeline_spatial_shard_decode(pipeline: torch.nn.Module) -> None:
|
|
"""Prepare every VAE that advertises spatial-shard decode capability."""
|
|
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
|
|
|
|
if not dist.is_initialized():
|
|
return
|
|
for vae in ModuleDiscovery.discover(pipeline).vaes:
|
|
prepare = getattr(vae, "_prepare_spatial_shard_decode", None)
|
|
if callable(prepare):
|
|
prepare()
|