Bağlam paralelliği

Bağlam paralelliği, girdi dizisini GPU'lar arasında böler ve her GPU kendi dilimini işler. Bu, GPU başına bellek kullanımını azaltır ve yüksek çözünürlüklü görüntüler ile video gibi uzun dizilerde en çok fayda sağlar. Diğer stratejilerle karşılaştırmak için Dağıtık çıkarım bölümüne bakın.

Bağlam paralelliği, dikkat (attention) arka uçlarının içinde çalışır ve yalnızca bazı arka uçlar bunu destekler. Kendi arka ucunuz desteklemiyorsa, enable_parallelism uyumlu arka uçları listeleyen bir hata fırlatır. Arka uçları set_attention_backend() ile değiştirin ve tam liste için mevcut arka uçlara bakın.

Aşağıdaki her varyantı ContextParallelConfig ile yapılandırın.

Halka Dikkat

Ring Attention, GPU'lerin bir halka şeklinde dizildiği sistemlerde anahtar (K) ve değer (V) bloklarını GPU'lar arasında iletir. Her GPU, üzerinde tuttuğu K/V bloğuyla dikkat (attention) hesabını yapar ve ardından bu bloğu bir sonraki GPU'ya iletir. Tam bir tur tamamlandığında, her GPU tüm token'ların K/V değerlerini görmüş olur ve hiçbir GPU tam diziyi üzerinde tutmaz.

ring_degree değerini GPU sayısına ayarlayın ve yapılandırmayı dönüştürücünün enable_parallelism() yöntemine iletin.

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()

Betik dosyasını ring_attention.py olarak kaydedin ve torchrun veya başka bir PyTorch uyumlu başlatıcı ile çalıştırın. --nproc-per-node değerini GPU sayısına ayarlayın.

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

Diğer varyantları denemek için ring_attention.py içindeki enable_parallelism çağrısını, ilgili bölümden olanıyla değiştirin.

Ulysses Dikkat

Ulysses Attention, her GPU'nün diğer tüm GPU'larla veri alışverişi yaptığı all-to-all iletişimini kullanır. İlk all-to-all, verileri yeniden gruplayarak her GPU'nun dikkat başlıklarının bir alt kümesi için tüm token'ları tutmasını sağlar. Her GPU kendi başlıkları için dikkat hesaplamasını yapar ve ikinci all-to-all, sonuçları bir sonraki katman için token bazında yeniden gruplar.

ulysses_degree değerini GPU sayısına ayarlayın ve yapılandırmayı enable_parallelism() fonksiyonuna iletin.

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

Birleşik Dikkat

Unified Sequence Parallelism, Ring Attention ve Ulysses Attention'ı birleştirir. Başları ve jetonları yeniden dağıtmak için bir all-to-all çalıştırır, sonucu üzerinde Ring Attention çalıştırır ve ardından orijinal düzeni geri getirmek için all-to-all'ı tersine çevirir. Ulysses, dikkat başları arasında paralelleştirme yaparken Ring, dizi üzerinde paralelleştirme yapar; bu nedenle birlikte, yalnızca Ulysses'e kıyasla daha fazla GPU kullanabilirler.

En az 4 GPU'nuz varsa Unified Attention kullanın. Hem ulysses_degree hem de ring_degree değerlerini 1'den büyük bir sayıya ayarlayın. GPU'lar, ulysses_degree * ring_degree cihazlık 2B bir ızgara oluşturur ve her GPU bir Ulysses grubuna ve bir Ring grubuna aittir.

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

Aşağıdaki benchmark, bu betik ile 4 H100 GPU üzerinde Ulysses, Ring ve Unified Attention'ı karşılaştırır. Her görüntü 50 çıkarım adımı kullanır.

Ulysses en yüksek verimi sağlar ancak modeldeki dikkat başı sayısından fazla GPU kullanamaz. Unified Attention bu sınırı kaldırır.

Ulysses Anything Dikkat

Ulysses Attention, dizi uzunluğunun ve dikkat başlarının sayısının ulysses_degree ile tam bölünebilmesini gerektirir. Ulysses Anything Attention her iki gereksinimi de ortadan kaldırır. Unified Attention ile birleştirilemediğinden ring_degree değeri 1 olmalıdır.

ulysses_degree değerini 1'den büyük bir sayıya ve ulysses_anything=True olarak ayarlayın.

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

init_process_group için CPU tensörlerine yönelik bir gloo arka ucu ekleyin. Bu, hosttan cihaza ve cihazdan hosta aktarımların neden olduğu zorunlu CUDA senkronizasyonlarını önler.

import torch.distributed as dist

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

Aşağıdaki benchmark, FLUX.1-dev modelini 4 L20 GPU üzerinde bu betikle çalıştırır.

1008x1008 boyutunda, dizi uzunluğu tam olarak bölünmediği için yalnızca Ulysses Anything Attention çalışır. 1024x1024 boyutunda ise standart Ulysses Attention ile aynı performansı gösterir.

Ring Anything Attention

Ring Attention, dizi uzunluğunun ring_degree değerine tam bölünmesini gerektirir. Ring Anything Attention ise herhangi bir dizi uzunluğunu destekler. Her GPU’nun yerel K/V’sini en uzun yerel uzunluğa kadar doldurur, doldurulmuş K/V’leri tüm GPU’larda toplar ve dikkat hesaplamadan önce her bloğu gerçek uzunluğuna geri keser.

Unified Attention ile birleştirilemez, bu nedenle ulysses_degree değeri 1 olmalıdır. ring_degree değerini 1'den büyük bir değere ve ring_anything=True olarak ayarlayın.

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

Ring Anything Attention, Ulysses Anything Attention bölümünde gösterilen gloo arka ucundan da faydalanır.

Ring Anything Attention yalnızca çıkarımı destekler ve attn_mask None olmalıdır.

Aşağıdaki benchmark, 4 RTX 4090 (48GB) GPU üzerinde FLUX.1-dev çalıştırır.

Ring Anything Attention 1008x1008 çözünürlüğünde çalışır ve standart Ring Attention ile yaklaşık aynı performansı gösterir. Motivasyon ve ödünler için bu yorumu inceleyin.

Yükleme sırasında bağlam paralelliğini etkinleştir

enable_parallelism çağrısını yükleme işleminden sonra yapmak yerine, ContextParallelConfig nesnesini from_pretrained() fonksiyonunun parallel_config argümanına iletin. from_pretrained, torch.distributed başlatılmamışsa hata fırlattığı için süreç grubunu (process group) önce başlatın. Ring Attention örneğinde olduğu gibi betiği torchrun ile başlatın.

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)