Paralelisme konteks

Context parallelism membagi urutan input di antara GPU, dan setiap GPU memproses bagiannya sendiri. Hal ini menurunkan memori per GPU, yang paling membantu untuk urutan panjang seperti gambar resolusi tinggi dan video. Untuk membandingkannya dengan strategi lain, lihat Distributed inference.

Context parallelism berjalan di dalam attention backend, dan hanya beberapa backend yang mendukungnya. Jika backend Anda tidak mendukungnya, enable_parallelism akan memicu error yang mencantumkan backend yang kompatibel. Ganti backend dengan set_attention_backend(), dan lihat backend yang tersedia untuk daftar lengkap.

Konfigurasikan setiap varian di bawah dengan ContextParallelConfig.

Ring Attention

Ring Attention meneruskan blok key (K) dan value (V) antar-GPU yang disusun dalam bentuk cincin. Setiap GPU menghitung perhatian terhadap blok K/V yang dipegangnya, lalu meneruskan blok tersebut ke GPU berikutnya. Setelah satu putaran penuh, setiap GPU telah melihat K/V dari setiap token, dan tidak ada GPU yang pernah memegang seluruh rangkaian.

Set ring_degree ke jumlah GPU, lalu lewati config ke metode enable_parallelism() 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()

Simpan skrip sebagai ring_attention.py dan jalankan dengan torchrun atau launcher lain yang kompatibel dengan PyTorch. Atur --nproc-per-node ke jumlah GPU.

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

Untuk mencoba varian lainnya, ganti panggilan enable_parallelism di ring_attention.py dengan yang ada di bagian mereka.

Perhatian Ulysses

Ulysses Attention menggunakan komunikasi all-to-all, di mana setiap GPU bertukar data dengan setiap GPU lainnya. all-to-all pertama mengelompokkan ulang data sehingga setiap GPU memegang setiap token untuk subset dari attention heads. Setiap GPU menghitung attention untuk heads-nya, dan all-to-all kedua mengelompokkan ulang hasil berdasarkan token untuk layer berikutnya.

Set ulysses_degree ke jumlah GPU, dan lewati config ke enable_parallelism().

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

Perhatian Terpadu

Unified Sequence Parallelism menggabungkan Ring Attention dan Ulysses Attention. Proses ini menjalankan all-to-all untuk mendistribusikan ulang heads dan tokens, kemudian menjalankan Ring Attention pada hasilnya, lalu membalikkan all-to-all untuk memulihkan tata letak asli. Ulysses melakukan paralelisasi di antara attention heads, sedangkan Ring melakukan paralelisasi di sepanjang sequence, sehingga secara bersama-sama keduanya dapat menggunakan lebih banyak GPU dibandingkan Ulysses saja.

Gunakan Unified Attention ketika Anda memiliki setidaknya 4 GPU. Atur ulysses_degree dan ring_degree keduanya ke nilai lebih dari 1. GPU membentuk grid 2D dari ulysses_degree * ring_degree perangkat, dan setiap GPU termasuk dalam satu grup Ulysses dan satu grup Ring.

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

Benchmark di bawah ini membandingkan Ulysses, Ring, dan Unified Attention pada 4 GPU H100 dengan skrip ini. Setiap gambar menggunakan 50 langkah inferensi.

Ulysses memiliki throughput tertinggi, tetapi tidak dapat menggunakan lebih banyak GPU daripada jumlah kepala perhatian (attention heads) pada model. Unified Attention menghilangkan batasan tersebut.

Ulysses Anything Attention

Ulysses Attention membutuhkan panjang sekuens dan jumlah kepala perhatian yang dapat dibagi rata oleh ulysses_degree. Ulysses Anything Attention menghilangkan kedua persyaratan tersebut. Anda tidak dapat menggabungkannya dengan Unified Attention, sehingga ring_degree harus bernilai 1.

Set ulysses_degree ke lebih dari 1 dan ulysses_anything=True.

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

Tambahkan backend gloo untuk tensor CPU ke init_process_group. Ini menghindari sinkronisasi CUDA yang dipaksakan akibat transfer dari host ke perangkat dan dari perangkat ke host.

import torch.distributed as dist

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

Benchmark di bawah ini menjalankan FLUX.1-dev pada 4 GPU L20 dengan skrip ini.

Pada 1008x1008, di mana panjang sekuens tidak terbagi rata, hanya Ulysses Anything Attention yang berjalan. Pada 1024x1024, kinerjanya sama dengan Ulysses Attention standar.

Ring Anything Attention

Ring Attention membutuhkan panjang sekuens yang dapat dibagi rata oleh ring_degree. Ring Anything Attention mendukung panjang sekuens apa pun. Ia melakukan padding pada K/V lokal di setiap GPU hingga panjang lokal terpanjang, melakukan all-gather pada K/V yang sudah di-padding, dan memotong kembali setiap blok ke panjang aslinya sebelum menghitung attention.

Anda tidak dapat menggabungkannya dengan Unified Attention, sehingga ulysses_degree harus 1. Atur ring_degree ke lebih dari 1 dan ring_anything=True.

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

Ring Anything Attention juga mendapat manfaat dari backend gloo yang ditampilkan di Ulysses Anything Attention.

Ring Anything Attention hanya mendukung inferensi, dan attn_mask harus None.

Benchmarks di bawah ini menjalankan FLUX.1-dev pada 4 GPU RTX 4090 (48GB).

Ring Anything Attention berjalan pada 1008x1008 dan kinerjanya kurang lebih sama dengan Ring Attention standar. Untuk motivasi dan trade-off-nya, lihat komentar ini.

Aktifkan paralelisme konteks saat memuat

Alih-alih memanggil enable_parallelism setelah memuat, teruskan ContextParallelConfig ke argumen parallel_config dari from_pretrained(). Inisialisasi grup proses terlebih dahulu, karena from_pretrained akan melempar error jika torch.distributed belum diinisialisasi. Jalankan skrip dengan torchrun, seperti pada contoh 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)