توازي السياق

يقوم Context parallelism بتقسيم تسلسل الإدخال على وحدات معالجة الرسوميات، حيث تعالج كل وحدة معالجة الرسوميات شريحتها الخاصة. يقلل هذا من الذاكرة لكل وحدة معالجة رسوميات، وهو ما يساعد بشكل أكبر مع التسلسلات الطويلة مثل الصور عالية الدقة والفيديو. للمقارنة مع الاستراتيجيات الأخرى، راجع Distributed inference.

يعمل توازي السياق داخل واجهة الانتباه الخلفية، ولا تدعمه بعض الواجهات الخلفية فقط. إذا لم يدعمه واجهتك، يرمي enable_parallelism خطأً يدرج الواجهات الخلفية المتوافقة. بدّل الواجهات الخلفية باستخدام set_attention_backend()، واطّلع على الواجهات الخلفية المتاحة للحصول على القائمة الكاملة.

قم بتكوين كل متغير أدناه باستخدام ContextParallelConfig.

انتباه الحلقة

يمرر Ring Attention كتل المفاتيح (K) والقيم (V) بين وحدات معالجة الرسوميات (GPUs) مرتبة في حلقة. تحسب كل وحدة معالجة الرسوميات الانتباه مقابل كتلة K/V التي تحتفظ بها، ثم تمرر تلك الكتلة إلى وحدة معالجة الرسوميات التالية. بعد دورة كاملة، تكون كل وحدة معالجة رسوميات قد رأت K/V لكل رمز، ولا تحتفظ أي وحدة معالجة رسوميات بالتسلسل الكامل.

اضبط ring_degree على عدد وحدات معالجة الرسومات، ومرّر الإعداد إلى طريقة enable_parallelism() الخاصة بالمحوّل.

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

احفظ السكربت باسم ring_attention.py وشغّله باستخدام torchrun أو أي مُشغّل متوافق مع PyTorch آخر. اضبط --nproc-per-node على عدد وحدات معالجة الرسومات.

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

لمحاولة المتغيرات الأخرى، استبدل استدعاء enable_parallelism في ring_attention.py بتلك الموجودة في قسمها.

انتباه أوليس

Ulysses Attention uses all-to-all communication, where every GPU exchanges data with every other GPU. The first all-to-all regroups the data so each GPU holds every token for a subset of attention heads. Each GPU computes attention for its heads, and a second all-to-all regroups the results by token for the next layer.

اضبط ulysses_degree على عدد وحدات معالجة الرسومات، ومرّر الإعداد إلى enable_parallelism().

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

الانتباه الموحّد

يجمع توازي التسلسل الموحّد بين Ring Attention و Ulysses Attention. يقوم بتنفيذ عملية all-to-all لإعادة توزيع الرؤوس والرموز، ثم ينفذ Ring Attention على النتيجة، ثم يعكس عملية all-to-all لاستعادة الترتيب الأصلي. يقوم Ulysses بالتوازي عبر رؤوس الانتباه، بينما يقوم Ring بالتوازي عبر التسلسل، لذا يمكنهما معًا استخدام عدد أكبر من وحدات معالجة الرسومات مقارنةً باستخدام Ulysses وحده.

استخدم Unified Attention عندما يكون لديك 4 وحدات معالجة رسومية (GPU) على الأقل. اضبط كلًا من ulysses_degree وring_degree على قيمة أكبر من 1. تُشكّل وحدات المعالجة الرسومية شبكة ثنائية الأبعاد مكوّنة من ulysses_degree * ring_degree جهازًا، وينتمي كل جهاز إلى مجموعة Ulysses واحدة ومجموعة Ring واحدة.

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

يقارن المعيار أدناه بين Ulysses وRing وUnified Attention على 4 وحدات معالجة رسومية H100 باستخدام هذا السكربت. يستخدم كل صورة 50 خطوة استنتاج.

يتميز Ulysses بأعلى معدل إنتاجية، لكنه لا يمكنه استخدام عدد أكبر من وحدات معالجة الرسوميات (GPU) من عدد رؤوس الانتباه في النموذج. يزيل Unified Attention هذا الحد.

Ulysses Anything Attention

يتطلب انتباه Ulysses أن تكون طول التسلسل وعدد رؤوس الانتباه قابلين للقسمة على ulysses_degree دون باقي. يزيل Ulysses Anything Attention كلا الشرطين. لا يمكن دمجه مع انتباه موحّد، لذا يجب أن يكون ring_degree مساوياً لـ 1.

اضبط ulysses_degree على قيمة أكبر من 1 وulysses_anything=True.

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

أضِف واجهة gloo لتensors المعالجة على وحدة المعالجة المركزية إلى init_process_group. يتجنّب هذا المزامنة الإجبارية على CUDA التي تسبّبها عمليات النقل من المضيف إلى الجهاز ومن الجهاز إلى المضيف.

import torch.distributed as dist

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

يُنفّذ المعيار أدناه نموذج FLUX.1-dev على 4 وحدات معالجة رسومية من طراز L20 باستخدام هذا السكربت.

عند 1008x1008، حيث لا يقبل طول التسلسل القسمة بالتساوي، يعمل فقط Ulysses Anything Attention. أما عند 1024x1024، فيعمل بنفس أداء Ulysses Attention القياسي.

Ring Anything Attention

يتطلب Ring Attention أن تكون طول التسلسل قابلاً للقسمة على ring_degree دون باقي. يدعم Ring Anything Attention أي طول للتسلسل. يقوم بتعبئة K/V المحلية لكل وحدة معالجة رسومية (GPU) حتى أطول طول محلي، ثم يجمع K/V المُعبّأة عبر جميع الوحدات، ويقطع كل كتلة إلى طولها الحقيقي قبل حساب الانتباه.

لا يمكن دمجها مع Unified Attention، لذا يجب أن يكون ulysses_degree مساوياً لـ 1. اضبط ring_degree على قيمة أكبر من 1 وring_anything=True.

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

يستفيد Ring Anything Attention أيضاً من الخلفية gloo المعروضة في Ulysses Anything Attention.

يدعم Ring Anything Attention الاستدلال فقط، ويجب أن يكون attn_mask بقيمة None.

يُنفّذ الاختبار أدناه نموذج FLUX.1-dev على 4 وحدات معالجة رسومية RTX 4090 (48 جيجابايت).

يعمل Ring Anything Attention بدقة 1008x1008 ويحقق أداءً مماثلاً تقريباً لأداء Ring Attention القياسي. للاطلاع على الدوافع والمقايضات، راجع هذا التعليق.

تفعيل توازٍ السياق أثناء التحميل

بدلاً من استدعاء enable_parallelism بعد التحميل، مرّر ContextParallelConfig إلى معامل parallel_config في from_pretrained(). قم بتهيئة مجموعة العمليات أولاً، لأن from_pretrained يرمي خطأً إذا لم تكن torch.distributed مُهيأة. شغّل السكربت باستخدام torchrun، كما في مثال 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)