Buckets:

mishig's picture
|
download
raw
20.1 kB

ParaAttention

قد تُشكّل نماذج توليد الصور والفيديو الكبيرة، مثل FLUX.1-dev وHunyuanVideo، تحدياً في الاستدلال للتطبيقات الفورية والنشر بسبب حجمها.

ParaAttention هي مكتبة تنفّذ توازي السياق وتخزين كتلة البداية مؤقتًا، ويمكن دمجها مع تقنيات أخرى (torch.compile، fp8 dynamic quantization) لتسريع الاستدلال.

ستُظهر لك هذه الدليل كيفية تطبيق ParaAttention على FLUX.1-dev و HunyuanVideo على وحدات معالجة الرسوميات NVIDIA L20. لم تُطبَّق أي تحسينات على معيارنا المرجعي الأساسي، باستثناء HunyuanVideo لتجنب أخطاء نفاد الذاكرة.

يُظهر معيارنا الأساسي أن FLUX.1-dev قادر على توليد صورة بدقة 1024x1024 في 28 خطوة خلال 26.36 ثانية، وأن HunyuanVideo قادر على توليد 129 إطارًا بدقة 720p في 30 خطوة خلال 3675.71 ثانية.

للحصول على استنتاج أسرع حتى مع توازٍ سياقي، جرّب استخدام وحدات معالجة رسومية من NVIDIA A100 أو H100 (إن توفرت) مع دعم NVLink، خاصةً عند وجود عدد كبير من وحدات المعالجة الرسومية.

ذاكرة التخزين المؤقت للكتلة الأولى

يُقلّل تخزين مخرجات كتل transformers في النموذج وإعادة استخدامها في خطوات الاستدلال التالية من تكلفة الحساب ويجعل الاستدلال أسرع.

ومع ذلك، يصعب تحديد متى ينبغي إعادة استخدام ذاكرة التخزين المؤقت لضمان جودة الصور أو مقاطع الفيديو المولّدة. يستخدم ParaAttention مباشرةً الفرق المتبقي في مخرج أول كتلة تحويل لتقريب الفرق بين مخرجات النموذج. عندما يكون الفرق صغيرًا بما يكفي، يُعاد استخدام الفرق المتبقي من خطوات الاستدلال السابقة. وبعبارة أخرى، يتم تجاوز خطوة إزالة الضوضاء.

يُحقق هذا تسريعًا بمقدار مرتين في استنتاج FLUX.1-dev وHunyuanVideo مع جودة ممتازة جدًا.

كيف يعمل AdaCache، حيث يُعد First Block Cache أحد متغيرات

لتطبيق ذاكرة التخزين المؤقت للكتلة الأولى على FLUX.1-dev، استدعِ apply_cache_on_pipe كما هو موضح أدناه. القيمة الافتراضية لفرق البقايا في نماذج FLUX هي 0.08.

import time
import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(pipe, residual_diff_threshold=0.08)

# Enable memory savings
# pipe.enable_model_cpu_offload()
# pipe.enable_sequential_cpu_offload()

begin = time.time()
image = pipe(
    "A cat holding a sign that says hello world",
    num_inference_steps=28,
).images[0]
end = time.time()
print(f"Time: {end - begin:.2f}s")

print("Saving image to flux.png")
image.save("flux.png")
التحسينات الأصل FBCache rdt=0.06 FBCache rdt=0.08 FBCache rdt=0.10 FBCache rdt=0.12
معاينة Original FBCache rdt=0.06 FBCache rdt=0.08 FBCache rdt=0.10 FBCache rdt=0.12
الوقت الحقيقي (ثانية) 26.36 21.83 17.01 16.00 13.78

قلّص كاش الكتلة الأولى (First Block Cache) زمن الاستدلال إلى 17.01 ثانية مقارنةً بالخط الأساس، أي أسرع بمقدار 1.55 مرة، مع الحفاظ على فقدان جودة شبه معدوم.

لتطبيق First Block Cache على HunyuanVideo، استخدم apply_cache_on_pipe كما هو موضح أدناه. القيمة الافتراضية لفرق البقايا في نماذج HunyuanVideo هي 0.06.

import time
import torch
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

model_id = "tencent/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    dtype=torch.bfloat16,
    revision="refs/pr/18",
)
pipe = HunyuanVideoPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=torch.float16,
    revision="refs/pr/18",
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(pipe, residual_diff_threshold=0.6)

pipe.vae.enable_tiling()

begin = time.time()
output = pipe(
    prompt="A cat walks on the grass, realistic",
    height=720,
    width=1280,
    num_frames=129,
    num_inference_steps=30,
).frames[0]
end = time.time()
print(f"Time: {end - begin:.2f}s")

print("Saving video to hunyuan_video.mp4")
export_to_video(output, "hunyuan_video.mp4", fps=15)

متصفحك لا يدعم وسم الفيديو.

HunyuanVideo بدون FBCache

متصفحك لا يدعم وسم الفيديو.

HunyuanVideo مع FBCache

قلّص كاش الكتلة الأولى زمن الاستدلال إلى 2271.06 ثانية مقارنةً بالخط الأساسي، أي أسرع بمقدار 1.62 مرة، مع الحفاظ على فقدان جودة شبه معدوم.

تكميم fp8

يسرّع fp8 مع التكميم الديناميكي الاستدلال بشكل أكبر ويقلل من استخدام الذاكرة. يجب تكميم كل من التفعيلات والأوزان لاستخدام نوى NVIDIA Tensor Cores ذات 8 بت.

استخدم float8_weight_only وfloat8_dynamic_activation_float8_weight لتكميم نموذج مُشفِّر النص (text encoder) ونموذج التحويل (transformer).

طريقة التكميم الافتراضية هي التكميم لكل مشغّل (per tensor quantization)، ولكن إذا كانت وحدة معالجة الرسوميات (GPU) لديك تدعم التكميم على مستوى الصف (row-wise quantization)، فيمكنك أيضًا تجربتها للحصول على دقة أفضل.

قم بتثبيت torchao باستخدام الأمر أدناه.

pip3 install -U torch torchao

يختار torch.compile مع mode="max-autotune-no-cudagraphs" أو mode="max-autotune" أفضل نواة للأداء. قد يستغرق التجميع وقتًا طويلاً إذا كانت هذه أول مرة يُستدعى فيها النموذج، لكنه يستحق العناء بعد اكتمال تجميع النموذج.

هذا المثال يقوم بتكميم نموذج التحويل فقط، لكن يمكنك أيضًا تكميم مشفّر النص لتقليل استخدام الذاكرة بشكل أكبر.

يمكن أن يُحدث التكميم الديناميكي تغييرًا كبيرًا في توزيع مخرجات النموذج، لذا عليك تغيير residual_diff_threshold إلى قيمة أكبر حتى يسري مفعوله.

import time
import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(
    pipe,
    residual_diff_threshold=0.12,  # Use a larger value to make the cache take effect
)

from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only

quantize_(pipe.text_encoder, float8_weight_only())
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
pipe.transformer = torch.compile(
   pipe.transformer, mode="max-autotune-no-cudagraphs",
)

# Enable memory savings
# pipe.enable_model_cpu_offload()
# pipe.enable_sequential_cpu_offload()

for i in range(2):
    begin = time.time()
    image = pipe(
        "A cat holding a sign that says hello world",
        num_inference_steps=28,
    ).images[0]
    end = time.time()
    if i == 0:
        print(f"Warm up time: {end - begin:.2f}s")
    else:
        print(f"Time: {end - begin:.2f}s")

print("Saving image to flux.png")
image.save("flux.png")

قلّص كلٌّ من fp8 dynamic quantization وtorch.compile زمن الاستدلال إلى 7.56 ثانية مقارنةً بالخط الأساس، أي أسرع بمقدار 3.48 مرة.

import time
import torch
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

model_id = "tencent/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    dtype=torch.bfloat16,
    revision="refs/pr/18",
)
pipe = HunyuanVideoPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=torch.float16,
    revision="refs/pr/18",
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(pipe)

from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only

quantize_(pipe.text_encoder, float8_weight_only())
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
pipe.transformer = torch.compile(
   pipe.transformer, mode="max-autotune-no-cudagraphs",
)

# Enable memory savings
pipe.vae.enable_tiling()
# pipe.enable_model_cpu_offload()
# pipe.enable_sequential_cpu_offload()

for i in range(2):
    begin = time.time()
    output = pipe(
        prompt="A cat walks on the grass, realistic",
        height=720,
        width=1280,
        num_frames=129,
        num_inference_steps=1 if i == 0 else 30,
    ).frames[0]
    end = time.time()
    if i == 0:
        print(f"Warm up time: {end - begin:.2f}s")
    else:
        print(f"Time: {end - begin:.2f}s")

print("Saving video to hunyuan_video.mp4")
export_to_video(output, "hunyuan_video.mp4", fps=15)

تتميز بطاقة NVIDIA L20 بذاكرة تبلغ 48 جيجابايت فقط، وقد تواجه أخطاء نفاد الذاكرة (OOM) بعد التجميع إذا لم يتم استدعاء enable_model_cpu_offload، وذلك لأن HunyuanVideo يحتوي على متجهات تنشيط كبيرة جداً عند التشغيل بدقة عالية وعدد كبير من الإطارات. بالنسبة للبطاقات التي تقل ذاكرتها عن 80 جيجابايت، يمكنك محاولة تقليل الدقة وعدد الإطارات لتجنب أخطاء OOM.

عادةً ما تكون نماذج توليد الفيديو الكبيرة مقيدة بحسابات الانتباه أكثر من الطبقات المتصلة بالكامل. ولا تستفيد هذه النماذج بشكل كبير من التكميم وtorch.compile.

توازي السياق

يعمل توازي السياق على موازاة الاستدلال ويتوسع مع استخدام عدة وحدات معالجة رسومية. ويتيح لك التصميم التركيبي لـ ParaAttention دمج توازي السياق مع ذاكرة التخزين المؤقت للكتلة الأولى والكمّية الديناميكية.

راجع مستودع ParaAttention للحصول على تعليمات مفصّلة وأمثلة حول كيفية توسيع نطاق الاستدلال باستخدام وحدات معالجة رسومية متعددة.

إذا كان من الضروري أن تكون عملية الاستدلال مستمرة وقابلة للخدمة، يُقترح استخدام torch.multiprocessing لكتابة معالج استدلال خاص بك. يمكن أن يؤدي ذلك إلى إزالة التكلفة الإضافية المرتبطة بإطلاق العملية وتحميل النموذج وإعادة تجميعه.

يجمع المثال البرمجي أدناه بين First Block Cache، والتكميم الديناميكي fp8، وtorch.compile، وتوازي السياق (Context Parallelism) للحصول على أسرع سرعة استدلال.

import time
import torch
import torch.distributed as dist
from diffusers import FluxPipeline

dist.init_process_group()

torch.cuda.set_device(dist.get_rank())

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.context_parallel import init_context_parallel_mesh
from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
from para_attn.parallel_vae.diffusers_adapters import parallelize_vae

mesh = init_context_parallel_mesh(
    pipe.device.type,
    max_ring_dim_size=2,
)
parallelize_pipe(
    pipe,
    mesh=mesh,
)
parallelize_vae(pipe.vae, mesh=mesh._flatten())

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(
    pipe,
    residual_diff_threshold=0.12,  # Use a larger value to make the cache take effect
)

from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only

quantize_(pipe.text_encoder, float8_weight_only())
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
torch._inductor.config.reorder_for_compute_comm_overlap = True
pipe.transformer = torch.compile(
   pipe.transformer, mode="max-autotune-no-cudagraphs",
)

# Enable memory savings
# pipe.enable_model_cpu_offload(gpu_id=dist.get_rank())
# pipe.enable_sequential_cpu_offload(gpu_id=dist.get_rank())

for i in range(2):
    begin = time.time()
    image = pipe(
        "A cat holding a sign that says hello world",
        num_inference_steps=28,
        output_type="pil" if dist.get_rank() == 0 else "pt",
    ).images[0]
    end = time.time()
    if dist.get_rank() == 0:
        if i == 0:
            print(f"Warm up time: {end - begin:.2f}s")
        else:
            print(f"Time: {end - begin:.2f}s")

if dist.get_rank() == 0:
    print("Saving image to flux.png")
    image.save("flux.png")

dist.destroy_process_group()

احفظه في run_flux.py ثم شغّله باستخدام torchrun.

# Use --nproc_per_node to specify the number of GPUs
torchrun --nproc_per_node=2 run_flux.py

تُقلَّص سرعة الاستدلال إلى 8.20 ثانية مقارنةً بالخط الأساسي، أي أسرع بـ 3.21 مرة، باستخدام معالجي رسومات NVIDIA L20. أما على 4 معالجات L20، فتبلغ سرعة الاستدلال 3.90 ثانية، أي أسرع بـ 6.75 مرة.

يجمع المثال البرمجي أدناه بين First Block Cache وContext Parallelism للحصول على أسرع سرعة استدلال.

import time
import torch
import torch.distributed as dist
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

dist.init_process_group()

torch.cuda.set_device(dist.get_rank())

model_id = "tencent/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    dtype=torch.bfloat16,
    revision="refs/pr/18",
)
pipe = HunyuanVideoPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=torch.float16,
    revision="refs/pr/18",
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.context_parallel import init_context_parallel_mesh
from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
from para_attn.parallel_vae.diffusers_adapters import parallelize_vae

mesh = init_context_parallel_mesh(
    pipe.device.type,
)
parallelize_pipe(
    pipe,
    mesh=mesh,
)
parallelize_vae(pipe.vae, mesh=mesh._flatten())

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(pipe)

# from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only
#
# torch._inductor.config.reorder_for_compute_comm_overlap = True
#
# quantize_(pipe.text_encoder, float8_weight_only())
# quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
# pipe.transformer = torch.compile(
#    pipe.transformer, mode="max-autotune-no-cudagraphs",
# )

# Enable memory savings
pipe.vae.enable_tiling()
# pipe.enable_model_cpu_offload(gpu_id=dist.get_rank())
# pipe.enable_sequential_cpu_offload(gpu_id=dist.get_rank())

for i in range(2):
    begin = time.time()
    output = pipe(
        prompt="A cat walks on the grass, realistic",
        height=720,
        width=1280,
        num_frames=129,
        num_inference_steps=1 if i == 0 else 30,
        output_type="pil" if dist.get_rank() == 0 else "pt",
    ).frames[0]
    end = time.time()
    if dist.get_rank() == 0:
        if i == 0:
            print(f"Warm up time: {end - begin:.2f}s")
        else:
            print(f"Time: {end - begin:.2f}s")

if dist.get_rank() == 0:
    print("Saving video to hunyuan_video.mp4")
    export_to_video(output, "hunyuan_video.mp4", fps=15)

dist.destroy_process_group()

احفظه في run_hunyuan_video.py ثم شغّله باستخدام torchrun.

# Use --nproc_per_node to specify the number of GPUs
torchrun --nproc_per_node=8 run_hunyuan_video.py

تُقلَّص سرعة الاستدلال إلى 649.23 ثانية مقارنةً بالخط الأساسي، أي أسرع بـ 5.66 مرة، باستخدام 8 وحدات معالجة رسومية NVIDIA L20.

المعايير

نوع وحدة معالجة الرسوميات عدد وحدات معالجة الرسوميات التحسينات الوقت الحقيقي (ثانية) التسريع
NVIDIA L20 1 الخط الأساسي 26.36 1.00x
NVIDIA L20 1 FBCache (rdt=0.08) 17.01 1.55x
NVIDIA L20 1 FP8 DQ 13.40 1.96x
NVIDIA L20 1 FBCache (rdt=0.12) + FP8 DQ 7.56 3.48x
NVIDIA L20 2 FBCache (rdt=0.12) + FP8 DQ + CP 4.92 5.35x
NVIDIA L20 4 FBCache (rdt=0.12) + FP8 DQ + CP 3.90 6.75x
نوع وحدة معالجة الرسوميات عدد وحدات معالجة الرسوميات التحسينات الوقت الحقيقي (ثانية) التسريع
NVIDIA L20 1 الخط الأساسي 3675.71 1.00x
NVIDIA L20 1 FBCache 2271.06 1.62x
NVIDIA L20 2 FBCache + CP 1132.90 3.24x
NVIDIA L20 4 FBCache + CP 718.15 5.12x
NVIDIA L20 8 FBCache + CP 649.23 5.66x

Xet Storage Details

Size:
20.1 kB
·
Xet hash:
8681c6697a7fb74ddd84204b0a2bf18e751d79229f018bf535807f0bad7fa931

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.