Class ContextParallelConfig

src/diffusers/models/_modeling_parallel.py:42–154

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

@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.
        ulysses_anything (`bool`, *optional*, defaults to `False`):
            Whether to enable "Ulysses Anything" mode, which supports arbitrary sequence lengths and head counts that
            are not evenly divisible by `ulysses_degree`. When enabled, `ulysses_degree` must be greater than 1 and
            `ring_degree` must be 1.
        ring_anything (`bool`, *optional*, defaults to `False`):
            Whether to enable "Ring Anything" mode, which supports arbitrary sequence lengths. When enabled,
            `ring_degree` must be greater than 1 and `ulysses_degree` must be 1.
        mesh (`torch.distributed.device_mesh.DeviceMesh`, *optional*):
            A custom device mesh to use for context parallelism. If provided, this mesh will be used instead of
            creating a new one. This is useful when combining context parallelism with other parallelism strategies
            (e.g., FSDP, tensor parallelism) that share the same device mesh. The mesh must have both "ring" and
            "ulysses" dimensions. Use size 1 for dimensions not being used (e.g., `mesh_shape=(2, 1, 4)` with
            `mesh_dim_names=("ring", "ulysses", "fsdp")` for ring attention only with FSDP).

    """

    ring_degree: int | None = None
    ulysses_degree: int | None = None
    convert_to_fp32: bool = True
    # TODO: support alltoall
    rotate_method: Literal["allgather", "alltoall"] = "allgather"
    mesh: torch.distributed.device_mesh.DeviceMesh | None = None
    # Whether to enable ulysses anything attention to support
    # any sequence lengths and any head numbers.
    ulysses_anything: bool = False
    # Whether to enable ring anything attention to support any sequence lengths.
    ring_anything: bool = False

    _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