توازي الموتر

يقوم توازي الموترات بتقسيم مصفوفات الأوزان لنموذج عبر الأجهزة. يحتفظ كل جهاز بشريحة عمودية ("colwise") أو أفقية ("rowwise") من كل طبقة ويحسب نتيجة جزئية. يجمع AllReduce أو AllGather عند حدود الطبقة النتائج الجزئية لتكوين المخرجات الكاملة. على عكس توازي السياق، يقلل من ذاكرة الأوزان لكل جهاز، مما يساعد عندما لا يتسع النموذج لجهاز واحد. راجع الاستدلال الموزع للمقارنة مع الاستراتيجيات الأخرى.

تحميل نموذج مقسّم

Pass a TensorParallelConfig to the parallel_config argument of the model’s from_pretrained(). tp_degree is the number of devices to shard across and must divide the model’s number of attention heads. The model’s class must define a _tp_plan, which maps module names to a sharding style. To check whether a model supports tensor parallelism, look for this attribute on its class, such as Flux2Transformer2DModel.

لإضافة خطة إلى نموذج آخر، راجع إضافة توازي الموترات إلى نموذج.

عند التحميل بهذه الطريقة، يتم تقسيم نقطة الحفظ (checkpoint) أثناء قراءتها. يقرأ كل رتبة (rank) شريحتها الخاصة من كل وزن ويضعها مباشرةً على وحدة معالجة الرسوميات (GPU) الخاصة به، لذا لا تحتفظ أي رتبة بوزن بحجمه الكامل. ينخفض حجم الذاكرة لكل رتبة مع زيادة tp_degree.

import os
import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, Flux2Transformer2DModel, TensorParallelConfig

def main():
    dist.init_process_group(backend="nccl")
    rank, world_size = dist.get_rank(), dist.get_world_size()
    # 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)

    # Each rank reads only its own shard of every planned weight, straight onto its own GPU.
    transformer = Flux2Transformer2DModel.from_pretrained(
        "black-forest-labs/FLUX.2-dev",
        subfolder="transformer",
        dtype=torch.bfloat16,
        parallel_config=TensorParallelConfig(tp_degree=world_size),
    )

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-dev", transformer=transformer, dtype=torch.bfloat16
    )
    # The transformer is already on its device; move the remaining components individually. Do not call
    # `pipeline.to(device)` — that would move every rank's shards onto the same device.
    pipeline.text_encoder.to(device)
    pipeline.vae.to(device)

    generator = torch.Generator().manual_seed(42)
    image = pipeline(prompt="a cat holding a sign that says hello", generator=generator).images[0]
    if rank == 0:
        image.save("output.png")
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

احفظ السكربت باسم tensor_parallel_flux2.py ثم شغّله باستخدام torchrun.

torchrun --nproc-per-node 4 tensor_parallel_flux2.py

يُعيّن السكربت tp_degree إلى world_size، لذا فإن --nproc-per-node 4 يُقسّم الترانسفورمر على 4 وحدات معالجة رسومية.

يتم تقسيم نموذج التحويل فقط. يحتفظ كل رتبة بنسخة كاملة من مشفّر النص، لذا يحتاج كل معالج رسومي إلى ذاكرة كافية لاستيعاب مشفّر النص بالإضافة إلى حصته من نموذج التحويل.

أداء التحميل

يكون التحميل باستخدام parallel_config أسرع ويستهلك ذاكرة أقل لكل رتبة مقارنة بتحميل النموذج الكامل ثم استدعاء enable_parallelism(). تُظهر الجدول أدناه النتائج الخاصة بـ black-forest-labs/FLUX.2-dev (المحوّل فقط، 32 مليار معامل، bf16) على 4x A10G (23 جيجابايت).

القيود

لا يعمل توازي الموترات مع بعض خيارات التحميل والميزات، ولا يمكن حفظ نموذج مجزأ.

• لا يعمل التحميل باستخدام parallel_config المتوازي على مستوى الموترات مع device_map أو quantization_config أو low_cpu_mem_usage=False أو use_flashpack=True أو أوزان غير safetensors. فكل واحدة من هذه الخيارات تثير خطأً بدلاً من تحميل نقطة الفحص الكاملة بصمت.
• يؤدي دمج توازي الموترات مع التكميم أو التفريغ أو محولات LoRA إلى ظهور خطأ، بغض النظر عن كيفية تقسيم النموذج.
• لتقسيم نموذج موجود مسبقًا في الذاكرة، استدعِ enable_parallelism() باستخدام نفس الإعداد. يقوم هذا بتحميل النموذج الكامل أولًا ثم تقسيمه، لذا يحتاج كل رتبة إلى ذاكرة كافية للنقطة الفحصية الكاملة. التحميل باستخدام parallel_config يتجنب هذا التكلفة.
• يُطلق حفظ نموذج متوازي الموترات خطأً في save_pretrained(). احفظ النموذج قبل تقسيمه.

إضافة توازي الموترات إلى نموذج

هذا القسم مخصص للمساهمين الذين يرغبون في إضافة دعم للتوازي الموتر (Tensor Parallelism) إلى نموذج في مكتبة Diffusers.

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

يمكن أن تحتوي كل مفتاح على * واحد كحد أقصى، ويجب أن يشير الجزء قبل * إلى nn.ModuleList، بحيث تغطي كل مدخلة كل كتلة. أما المفتاح الذي لا يحتوي على * فيُطبَّق على النموذج نفسه. والمسارات نسبية بالنسبة إلى النموذج.

عمودي وأفقي

يجب دائمًا إقران "colwise" مع "rowwise" بهذا الترتيب. إذ تترك إسقاط "colwise" ناتجه مقسّمًا، بينما يستهلك إسقاط "rowwise" التالي ذلك الجزء مباشرة، ثم يعيد AllReduce واحد عند حدّ الكتلة بناء النتيجة. أما تقسيم الزوج بأي طريقة أخرى فيُجبر على تجميع في المنتصف ويتطلب تواصلًا أكبر بكثير.

في الانتباه، تعطي الإسقاطات العمودية لـ Q و K و V كل رتبة مجموعة فرعية من الرؤوس، لذا يجب أن يقسم tp_degree عدد الرؤوس. تتبع النسخ المكررة لتيار المُشفّر (add_q_proj، to_add_out، ff_context) نفس النمط الذي تتبعه نظيراتها في تيار الصورة.

الاستيفاءات المدمجة

عندما يجمع Linear عدة متجهات منطقية على طول البُعد المُجزَّأ، فإن التقسيمات البسيطة "colwise"/"rowwise" تعبر عبر عملية الدمج مباشرةً مما يؤدي إلى عدم محاذاة القطع. استخدم PackedColwiseParallel/PackedRowwiseParallel بدلاً من ذلك، حيث تُجزَّأ كل كتلة مجمّعة بشكل مستقل.

# in src/diffusers/models/transformers/your_model.py
from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel

مرّر blocks، وهو قائمة من الأعداد الصحيحة النسبية التي يقبل مجموعها القسمة على البُعد المُجمَّع. على سبيل المثال، استخدم [1, 1] لبوابة SwiGLU وتحويل تصاعدي بنصفين متساويين، أو استخدم [1, 1, 1, 3, 3] لتحويل Q وK وV وبوابة وتحويل تصاعدي مُدمجة مع mlp_ratio=3.

"transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]),

عندما تكون أحجام الكتل معروفة فقط من ملف الإعداد، احذف blocks وخزّن الأحجام المطلقة على Linear أثناء __init__ بدلاً من ذلك، كـ _tp_packed_col_blocks أو _tp_packed_row_blocks.

# in the attention module's __init__
self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim]

# in the model's _tp_plan
"single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(),

يستخدم الخطة في transformer_flux2.py كلا الشكلين. يمرر زوج SwiGLU ذو التدفق المزدوج blocks مباشرةً. أما إسقاط QKV و MLP المدمج ذو التدفق الأحادي فيخزّن أحجامه على Linear داخل Flux2ParallelSelfAttention.__init__، وذلك لأنها تعتمد على mlp_ratio وmlp_mult_factor.

ما يجب استبعاده

يتم تكرار الوحدات غير المدرجة في الخطة على كل رتبة. اترك طبقات التطبيع، وتضمينات البقع والنص، والطبقتين الأخيرتين norm_out وproj_out. فهي صغيرة، لذا فإن تقسيمها يوفر القليل من الذاكرة ويضيف اتصالات. تعديل AdaLN (img_mod وtxt_mod) يكون عادةً صغيرًا بما يكفي ليُترك خارج الخطة أيضًا. عندما يكون كبيرًا، مثل adaln_proj.linear في MiniMax-H3، قسّمه باستخدام "rowwise_split_input".

قيود والتحقق

• يجب أن يقسم tp_degree config.num_attention_heads.
• يجب أن يكون كل كتلة مكدّسة بشكل فردي قابلة للقسمة على tp_degree، وليس مجموعها فقط.

from_pretrained() وenable_parallelism() يتحققان من كلا الأمرين قبل تحميل أي وزن أو تقسيمه.

تحقّق رقمياً من خطة جديدة بدلاً من التحقق البصري: ولّد النتائج باستخدام بذرة ثابتة على جهاز واحد، ثم أعد التوليد تحت توازي الموترات، وقارن المخرجات. غالباً ما يؤدي وضع "colwise"/"rowwise" في المكان الخطأ إلى تشغيل النموذج وإنتاج صورة تبدو معقولة لكنها خاطئة.

ابدأ من خطة موجودة مسبقًا لهندسة معمارية مشابهة.

• لا يحتوي transformer_qwenimage.py على إسقاطات مدمجة، لذا كل عنصر هو ببساطة "colwise" أو "rowwise".
• يضيف transformer_flux.py إسقاطًا صفّيًا مكدّسًا واحدًا.
• يستخدم transformer_flux2.py كلا الأسلوبين المعبّأة.