264 lines
11 KiB
Python
264 lines
11 KiB
Python
# 🚨🚨🚨 Experimental parallelism support for Diffusers 🚨🚨🚨
|
|
# Experimental changes are subject to change and APIs may break without warning.
|
|
|
|
# Copyright 2025 The HuggingFace Team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union
|
|
|
|
import torch
|
|
|
|
from diffusers.utils import get_logger
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
pass
|
|
|
|
|
|
logger = get_logger(__name__) # pylint: disable=invalid-name
|
|
|
|
|
|
# TODO(aryan): add support for the following:
|
|
# - Unified Attention
|
|
# - More dispatcher attention backends
|
|
# - CFG/Data Parallel
|
|
# - Tensor Parallel
|
|
|
|
|
|
@dataclass
|
|
class ContextParallelConfig:
|
|
"""
|
|
Configuration for context parallelism.
|
|
|
|
Args:
|
|
ring_degree (`int`, *optional*, defaults to `1`):
|
|
Number of devices to use for Ring Attention. Sequence is split across devices. Each device computes
|
|
attention between its local Q and KV chunks passed sequentially around ring. Lower memory (only holds 1/N
|
|
of KV at a time), overlaps compute with communication, but requires N iterations to see all tokens. Best
|
|
for long sequences with limited memory/bandwidth. Number of devices to use for ring attention within a
|
|
context parallel region. Must be a divisor of the total number of devices in the context parallel mesh.
|
|
ulysses_degree (`int`, *optional*, defaults to `1`):
|
|
Number of devices to use for Ulysses Attention. Sequence split is across devices. Each device computes
|
|
local QKV, then all-gathers all KV chunks to compute full attention in one pass. Higher memory (stores all
|
|
KV), requires high-bandwidth all-to-all communication, but lower latency. Best for moderate sequences with
|
|
good interconnect bandwidth.
|
|
convert_to_fp32 (`bool`, *optional*, defaults to `True`):
|
|
Whether to convert output and LSE to float32 for ring attention numerical stability.
|
|
rotate_method (`str`, *optional*, defaults to `"allgather"`):
|
|
Method to use for rotating key/value states across devices in ring attention. Currently, only `"allgather"`
|
|
is supported.
|
|
|
|
"""
|
|
|
|
ring_degree: Optional[int] = None
|
|
ulysses_degree: Optional[int] = None
|
|
convert_to_fp32: bool = True
|
|
# TODO: support alltoall
|
|
rotate_method: Literal["allgather", "alltoall"] = "allgather"
|
|
|
|
_rank: int = None
|
|
_world_size: int = None
|
|
_device: torch.device = None
|
|
_mesh: torch.distributed.device_mesh.DeviceMesh = None
|
|
_flattened_mesh: torch.distributed.device_mesh.DeviceMesh = None
|
|
_ring_mesh: torch.distributed.device_mesh.DeviceMesh = None
|
|
_ulysses_mesh: torch.distributed.device_mesh.DeviceMesh = None
|
|
_ring_local_rank: int = None
|
|
_ulysses_local_rank: int = None
|
|
|
|
def __post_init__(self):
|
|
if self.ring_degree is None:
|
|
self.ring_degree = 1
|
|
if self.ulysses_degree is None:
|
|
self.ulysses_degree = 1
|
|
|
|
if self.ring_degree == 1 and self.ulysses_degree == 1:
|
|
raise ValueError(
|
|
"Either ring_degree or ulysses_degree must be greater than 1 in order to use context parallel inference"
|
|
)
|
|
if self.ring_degree < 1 or self.ulysses_degree < 1:
|
|
raise ValueError("`ring_degree` and `ulysses_degree` must be greater than or equal to 1.")
|
|
if self.ring_degree > 1 and self.ulysses_degree > 1:
|
|
raise ValueError(
|
|
"Unified Ulysses-Ring attention is not yet supported. Please set either `ring_degree` or `ulysses_degree` to 1."
|
|
)
|
|
if self.rotate_method != "allgather":
|
|
raise NotImplementedError(
|
|
f"Only rotate_method='allgather' is supported for now, but got {self.rotate_method}."
|
|
)
|
|
|
|
@property
|
|
def mesh_shape(self) -> Tuple[int, int]:
|
|
return (self.ring_degree, self.ulysses_degree)
|
|
|
|
@property
|
|
def mesh_dim_names(self) -> Tuple[str, str]:
|
|
"""Dimension names for the device mesh."""
|
|
return ("ring", "ulysses")
|
|
|
|
def setup(self, rank: int, world_size: int, device: torch.device, mesh: torch.distributed.device_mesh.DeviceMesh):
|
|
self._rank = rank
|
|
self._world_size = world_size
|
|
self._device = device
|
|
self._mesh = mesh
|
|
|
|
if self.ulysses_degree * self.ring_degree > world_size:
|
|
raise ValueError(
|
|
f"The product of `ring_degree` ({self.ring_degree}) and `ulysses_degree` ({self.ulysses_degree}) must not exceed the world size ({world_size})."
|
|
)
|
|
|
|
self._flattened_mesh = self._mesh._flatten()
|
|
self._ring_mesh = self._mesh["ring"]
|
|
self._ulysses_mesh = self._mesh["ulysses"]
|
|
self._ring_local_rank = self._ring_mesh.get_local_rank()
|
|
self._ulysses_local_rank = self._ulysses_mesh.get_local_rank()
|
|
|
|
|
|
@dataclass
|
|
class ParallelConfig:
|
|
"""
|
|
Configuration for applying different parallelisms.
|
|
|
|
Args:
|
|
context_parallel_config (`ContextParallelConfig`, *optional*):
|
|
Configuration for context parallelism.
|
|
"""
|
|
|
|
context_parallel_config: Optional[ContextParallelConfig] = None
|
|
|
|
_rank: int = None
|
|
_world_size: int = None
|
|
_device: torch.device = None
|
|
_mesh: torch.distributed.device_mesh.DeviceMesh = None
|
|
|
|
def setup(
|
|
self,
|
|
rank: int,
|
|
world_size: int,
|
|
device: torch.device,
|
|
*,
|
|
mesh: Optional[torch.distributed.device_mesh.DeviceMesh] = None,
|
|
):
|
|
self._rank = rank
|
|
self._world_size = world_size
|
|
self._device = device
|
|
self._mesh = mesh
|
|
if self.context_parallel_config is not None:
|
|
self.context_parallel_config.setup(rank, world_size, device, mesh)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ContextParallelInput:
|
|
"""
|
|
Configuration for splitting an input tensor across context parallel region.
|
|
|
|
Args:
|
|
split_dim (`int`):
|
|
The dimension along which to split the tensor.
|
|
expected_dims (`int`, *optional*):
|
|
The expected number of dimensions of the tensor. If provided, a check will be performed to ensure that the
|
|
tensor has the expected number of dimensions before splitting.
|
|
split_output (`bool`, *optional*, defaults to `False`):
|
|
Whether to split the output tensor of the layer along the given `split_dim` instead of the input tensor.
|
|
This is useful for layers whose outputs should be split after it does some preprocessing on the inputs (ex:
|
|
RoPE).
|
|
"""
|
|
|
|
split_dim: int
|
|
expected_dims: Optional[int] = None
|
|
split_output: bool = False
|
|
|
|
def __repr__(self):
|
|
return f"ContextParallelInput(split_dim={self.split_dim}, expected_dims={self.expected_dims}, split_output={self.split_output})"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ContextParallelOutput:
|
|
"""
|
|
Configuration for gathering an output tensor across context parallel region.
|
|
|
|
Args:
|
|
gather_dim (`int`):
|
|
The dimension along which to gather the tensor.
|
|
expected_dims (`int`, *optional*):
|
|
The expected number of dimensions of the tensor. If provided, a check will be performed to ensure that the
|
|
tensor has the expected number of dimensions before gathering.
|
|
"""
|
|
|
|
gather_dim: int
|
|
expected_dims: Optional[int] = None
|
|
|
|
def __repr__(self):
|
|
return f"ContextParallelOutput(gather_dim={self.gather_dim}, expected_dims={self.expected_dims})"
|
|
|
|
|
|
# A dictionary where keys denote the input to be split across context parallel region, and the
|
|
# value denotes the sharding configuration.
|
|
# If the key is a string, it denotes the name of the parameter in the forward function.
|
|
# If the key is an integer, split_output must be set to True, and it denotes the index of the output
|
|
# to be split across context parallel region.
|
|
ContextParallelInputType = Dict[
|
|
Union[str, int], Union[ContextParallelInput, List[ContextParallelInput], Tuple[ContextParallelInput, ...]]
|
|
]
|
|
|
|
# A dictionary where keys denote the output to be gathered across context parallel region, and the
|
|
# value denotes the gathering configuration.
|
|
ContextParallelOutputType = Union[
|
|
ContextParallelOutput, List[ContextParallelOutput], Tuple[ContextParallelOutput, ...]
|
|
]
|
|
|
|
# A dictionary where keys denote the module id, and the value denotes how the inputs/outputs of
|
|
# the module should be split/gathered across context parallel region.
|
|
ContextParallelModelPlan = Dict[str, Union[ContextParallelInputType, ContextParallelOutputType]]
|
|
|
|
|
|
# Example of a ContextParallelModelPlan (QwenImageTransformer2DModel):
|
|
#
|
|
# Each model should define a _cp_plan attribute that contains information on how to shard/gather
|
|
# tensors at different stages of the forward:
|
|
#
|
|
# ```python
|
|
# _cp_plan = {
|
|
# "": {
|
|
# "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
|
|
# "encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
|
|
# "encoder_hidden_states_mask": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False),
|
|
# },
|
|
# "pos_embed": {
|
|
# 0: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True),
|
|
# 1: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True),
|
|
# },
|
|
# "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
|
|
# }
|
|
# ```
|
|
#
|
|
# The dictionary is a set of module names mapped to their respective CP plan. The inputs/outputs of layers will be
|
|
# split/gathered according to this at the respective module level. Here, the following happens:
|
|
# - "":
|
|
# we specify that we want to split the various inputs across the sequence dim in the pre-forward hook (i.e. before
|
|
# the actual forward logic of the QwenImageTransformer2DModel is run, we will splitthe inputs)
|
|
# - "pos_embed":
|
|
# we specify that we want to split the outputs of the RoPE layer. Since there are two outputs (imag & text freqs),
|
|
# we can individually specify how they should be split
|
|
# - "proj_out":
|
|
# before returning to the user, we gather the entire sequence on each rank in the post-forward hook (after the linear
|
|
# layer forward has run).
|
|
#
|
|
# ContextParallelInput:
|
|
# specifies how to split the input tensor in the pre-forward or post-forward hook of the layer it is attached to
|
|
#
|
|
# ContextParallelOutput:
|
|
# specifies how to gather the input tensor in the post-forward hook in the layer it is attached to
|