Song song ngữ cảnh

Context parallelism chia chuỗi đầu vào giữa các GPU, và mỗi GPU xử lý phần riêng của nó. Điều này làm giảm bộ nhớ trên mỗi GPU, hữu ích nhất cho các chuỗi dài như ảnh độ phân giải cao và video. Để so sánh với các chiến lược khác, xem Distributed inference.

Context parallelism được thực thi bên trong attention backend, và chỉ một số backend hỗ trợ tính năng này. Nếu backend của bạn không hỗ trợ, enable_parallelism sẽ gây ra lỗi liệt kê các backend tương thích. Hãy chuyển đổi backend bằng set_attention_backend() và xem các backend khả dụng để có danh sách đầy đủ.

Cấu hình mọi biến thể bên dưới bằng ContextParallelConfig.

Ring Attention

Ring Attention truyền các khối key (K) và value (V) giữa các GPU được sắp xếp thành một vòng. Mỗi GPU tính toán attention với khối K/V mà nó đang giữ, sau đó truyền khối đó sang GPU tiếp theo. Sau một vòng lặp hoàn chỉnh, mọi GPU đều đã xem qua K/V của mọi token, và không có GPU nào giữ toàn bộ chuỗi.

Đặt ring_degree bằng số lượng GPU và truyền cấu hình vào phương thức enable_parallelism() của transformer.

import os
import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, ContextParallelConfig

def setup_distributed():
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    # torchrun sets LOCAL_RANK to this process's GPU index on its node
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)
    return device

def main():
    device = setup_distributed()
    world_size = dist.get_world_size()

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-dev", dtype=torch.bfloat16
    ).to(device)
    pipeline.transformer.set_attention_backend("_native_cudnn")

    cp_config = ContextParallelConfig(ring_degree=world_size)
    pipeline.transformer.enable_parallelism(config=cp_config)

    prompt = """
    cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
    highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
    """

    # Must specify generator so all ranks start with same latents (or pass your own)
    generator = torch.Generator().manual_seed(42)
    image = pipeline(
        prompt,
        guidance_scale=3.5,
        num_inference_steps=50,
        generator=generator,
    ).images[0]

    if dist.get_rank() == 0:
        image.save("output.png")

    if dist.is_initialized():
        dist.destroy_process_group()


if __name__ == "__main__":
    main()

Lưu script dưới tên ring_attention.py và chạy bằng torchrun hoặc trình khởi chạy tương thích PyTorch khác. Đặt --nproc-per-node bằng số lượng GPU.

torchrun --nproc-per-node 2 ring_attention.py

Để thử các biến thể khác, hãy thay thế lời gọi enable_parallelism trong ring_attention.py bằng lời gọi từ phần tương ứng của chúng.

Ulysses Attention

Ulysses Attention sử dụng giao tiếp all-to-all, trong đó mỗi GPU trao đổi dữ liệu với mọi GPU khác. Thao tác all-to-all đầu tiên nhóm lại dữ liệu để mỗi GPU giữ mọi token cho một tập con các attention head. Mỗi GPU tính toán attention cho các head của nó, và thao tác all-to-all thứ hai nhóm lại kết quả theo token cho lớp tiếp theo.

Đặt ulysses_degree bằng số lượng GPU và truyền cấu hình vào enable_parallelism().

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2))

Unified Attention

Unified Sequence Parallelism kết hợp Ring Attention và Ulysses Attention. Nó thực hiện all-to-all để phân phối lại heads và tokens, chạy Ring Attention trên kết quả, sau đó đảo ngược all-to-all để khôi phục bố cục ban đầu. Ulysses song song hóa trên các attention heads và Ring song song hóa trên sequence, do đó khi kết hợp chúng có thể sử dụng nhiều GPU hơn so với chỉ dùng Ulysses.

Sử dụng Unified Attention khi bạn có ít nhất 4 GPU. Đặt cả ulysses_degree và ring_degree lớn hơn 1. Các GPU tạo thành một lưới 2D gồm ulysses_degree * ring_degree thiết bị, và mỗi GPU thuộc về một nhóm Ulysses và một nhóm Ring.

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ring_degree=2))

Bảng chuẩn dưới đây so sánh Ulysses, Ring và Unified Attention trên 4 GPU H100 bằng script này. Mỗi hình ảnh sử dụng 50 bước suy luận.

Ulysses có thông lượng cao nhất, nhưng không thể sử dụng nhiều GPU hơn số đầu chú ý (attention heads) của mô hình. Unified Attention loại bỏ giới hạn đó.

Ulysses Anything Attention

Ulysses Attention yêu cầu độ dài chuỗi và số đầu attention phải chia hết cho ulysses_degree. Ulysses Anything Attention loại bỏ cả hai yêu cầu này. Bạn không thể kết hợp nó với Unified Attention, do đó ring_degree phải bằng 1.

Đặt ulysses_degree lớn hơn 1 và ulysses_anything=True.

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True))

Thêm backend gloo cho tensor CPU vào init_process_group. Điều này tránh được các đồng bộ hóa CUDA bắt buộc do các thao tác chuyển dữ liệu từ host sang thiết bị và từ thiết bị sang host gây ra.

import torch.distributed as dist

dist.init_process_group(backend="cpu:gloo,cuda:nccl")

Bảng đánh giá dưới đây chạy FLUX.1-dev trên 4 GPU L20 bằng script này.

Ở kích thước 1008x1008, khi độ dài chuỗi không chia hết, chỉ có Ulysses Anything Attention được thực thi. Ở kích thước 1024x1024, hiệu năng của nó tương đương với Ulysses Attention tiêu chuẩn.

Ring Anything Attention

Ring Attention yêu cầu độ dài chuỗi phải chia hết cho ring_degree. Ring Anything Attention hỗ trợ mọi độ dài chuỗi. Nó đệm K/V cục bộ của mỗi GPU lên độ dài cục bộ dài nhất, thực hiện all-gather trên K/V đã đệm, và cắt mỗi khối về độ dài thực trước khi tính toán attention.

Bạn không thể kết hợp nó với Unified Attention, do đó ulysses_degree phải bằng 1. Đặt ring_degree lớn hơn 1 và ring_anything=True.

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2, ring_anything=True))

Ring Anything Attention cũng được hưởng lợi từ backend gloo được trình bày trong Ulysses Anything Attention.

Ring Anything Attention chỉ hỗ trợ suy luận, và attn_mask phải là None.

Bảng chuẩn (benchmark) bên dưới chạy FLUX.1-dev trên 4 GPU RTX 4090 (48GB).

Ring Anything Attention chạy ở kích thước 1008x1008 và cho hiệu suất tương đương với Ring Attention tiêu chuẩn. Để tìm hiểu về động lực và các đánh đổi, hãy xem bình luận này.

Bật song song ngữ cảnh khi tải

Thay vì gọi enable_parallelism sau khi tải, hãy truyền ContextParallelConfig vào tham số parallel_config của from_pretrained(). Khởi tạo nhóm tiến trình trước, vì from_pretrained sẽ gây ra lỗi nếu torch.distributed chưa được khởi tạo. Chạy script bằng torchrun, như trong ví dụ Ring Attention.

import os
import torch
from torch import distributed as dist
from diffusers import AutoModel, ContextParallelConfig, DiffusionPipeline

CKPT_ID = "black-forest-labs/FLUX.1-dev"

dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)

cp_config = ContextParallelConfig(ring_degree=dist.get_world_size())
transformer = AutoModel.from_pretrained(
    CKPT_ID,
    subfolder="transformer",
    dtype=torch.bfloat16,
    parallel_config=cp_config
)

pipeline = DiffusionPipeline.from_pretrained(
    CKPT_ID, transformer=transformer, dtype=torch.bfloat16,
).to(device)