التوازي في الموترات للاستدلال

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

يدعم النموذج توازي الموترات إذا عرّفت إعداده الخاصية base_model_tp_plan. يمكنك التحقق من نموذج محمّل باستخدام الخاصية supports_tp_plan.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
print(model.supports_tp_plan)

يتناول هذا الدليل تفعيل توازي الموترات (Tensor Parallelism) في مكتبة Transformers والاستراتيجيات المتاحة لتقسيم الموارد.

تقسيم النموذج

قم بتكوين عدد أجهزة التوازي المجهري باستخدام tp_size في DistributedConfig.

• اضبط DistributedConfig(tp_size=N) لاستخدام الخطة المحددة مسبقًا للنموذج.
• حدّد tp_plan يدويًا وقم بتمريره إلى DistributedConfig مع tp_size.

يمكنك أيضاً تعيين tp_plan="auto" لطلب الخطة المحددة مسبقاً بشكل صريح. عند إغفال tp_size وتعيين tp_plan، يُشتق tp_size من WORLD_SIZE مقسوماً على أحجام التوازي الأخرى. تمرير tp_plan مباشرة إلى from_pretrained() أصبح غير موصى به وسيُزال في الإصدار v5.18.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig

# model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct" # better to visualize all the possible strategies
distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype=torch.bfloat16,
    distributed_config=distributed_config,
)
print(model.tp_plan)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
prompt = "Can I help"
inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)

# distributed run
outputs = model(inputs)

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

from transformers import AutoModelForCausalLM, DistributedConfig

tp_plan = {
    "model.layers.*.self_attn.q_proj": "colwise",
    "model.layers.*.self_attn.k_proj": "colwise",
    "model.layers.*.self_attn.v_proj": "colwise",
    "model.layers.*.self_attn.o_proj": "rowwise",
    ...
}

distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype="auto",
    distributed_config=distributed_config,
)
print(model.tp_plan)

استراتيجيات التقسيم

ALL_PARALLEL_STYLES، وهو ParallelInterface مُعرَّف في distributed/parallel_utils.py، يُقرِن كل اسم استراتيجية يمكن استخدامه في خطة إلى مثيل مُهيَّأ من الاستراتيجية. لا تتفاعل معه مباشرةً لتقسيم نموذج، لكنه القائمة المرجعية للأسماء المتاحة. يسجِّل distributed/tensor_parallel.py استراتيجيات توازي الموترات فيه:

ALL_PARALLEL_STYLES.register("embedding_rowwise", EmbeddingParallel(shard_dim=0))
ALL_PARALLEL_STYLES.register("embedding_colwise", EmbeddingParallel(shard_dim=1))
ALL_PARALLEL_STYLES.register("colwise_gather_output", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise_rep", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise", ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)))
ALL_PARALLEL_STYLES.register("rowwise", RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_split_input", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_rep", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("packed_colwise", PackedColwiseParallel())
ALL_PARALLEL_STYLES.register("packed_rowwise", PackedRowwiseParallel())
ALL_PARALLEL_STYLES.register("sequence_parallel", SequenceParallel(use_local_output=True))
ALL_PARALLEL_STYLES.register("replicated_with_grad_allreduce", ReplicatedWithGradAllReduce())
ALL_PARALLEL_STYLES.register("mla_kv_a_proj", MlaKvAProjParallel())
ALL_PARALLEL_STYLES.register("all_reduce", AllReduceParallel())
ALL_PARALLEL_STYLES.register("moe_tp_experts", MoeExpertsParallel())
ALL_PARALLEL_STYLES.register("moe_identity_expert", MoeIdentityParallel())
ALL_PARALLEL_STYLES.register("local_params", LocalParamsParallel())

كل استراتيجية للتوازي على مستوى الموترات هي فئة فرعية من TensorParallelLayer في ذلك الملف. أما استراتيجيات التوازي على مستوى الخبراء، مثل ep_dispatch_experts، فتوجد في distributed/expert_parallel.py، وهي فئات فرعية من ExpertParallelLayer، التي تملك نفس الخطافات، وتُسجَّل في ALL_PARALLEL_STYLES نفسها. الاستراتيجيتان اللتان ستلجأ إليهما في الغالب، ColwiseParallel وRowwiseParallel، تقبلان مواضع input_layouts وoutput_layouts، وهو ما يتيح لفئة واحدة تغطية عدة أسماء مخططات. تترك colwise ناتجها مجزأً على البُعد الأخير ليستهلكه طبقة rowwise تالية، بينما تجمع colwise_gather_output الناتج بالكامل مرة أخرى إلى موتر كامل.

تصف الجدول أدناه كل استراتيجية.

استراتيجيات مكدّسة

تقوم حزمة الأوزان بدمج عدة طبقات خطية في طبقة واحدة أكبر. تقسّم استراتيجيتا PackedColwiseParallel وPackedRowwiseParallel الأوزان المحزومة بشكل صحيح. أما استراتيجيتا ColwiseParallel أو RowwiseParallel الأساسيتان فتقسّمان الأوزان المحزومة بشكل غير صحيح.

يجمع المثال أدناه up_proj وgate_proj في وحدة واحدة اسمها gate_up_proj ويتطلب استخدام استراتيجية packed_rowwise لتقسيم gate_up_proj.

class Llama4TextExperts(nn.Module):
    ...
    self.gate_up_proj = nn.Parameter(torch.zeros(self.num_experts, self.hidden_size, 2 * self.expert_dim))

استخدم ضرب المصفوفات الدفعية في مسار forward لحساب مخرج وحدة gate_up_proj.

def forward(self, hidden_states):
    ...
    gate_up = torch.bmm(hidden_states, self.gate_up_proj) # Compute the output of the gate_up_proj module
    gate, up = gate_up.chunk(2, dim=-1) # Split the output into gate and up

يقسم Shard العادي تلك البُعد إلى كتل متتالية، لذا يستلم الرتبة ذيل gate ورأس up بدلاً من شريحة من كلٍّ منهما. أما استراتيجيات الحزم فتستخدم _StridedShard مع split_factor (الافتراضي 2) لتبديل التقسيم، مما يمنح كل رتبة شريحة متطابقة من كلا النصفين حتى تبقى chunk متطابقة بعد التجزئة.

استراتيجيات التقسيم المخصصة

ورث من TensorParallelLayer في distributed/tensor_parallel.py لإنشاء استراتيجية تقسيم مخصصة. تجاوز فقط الخطافات التي تحتاجها استراتيجيتك، إذ أن لكل خطاف قيمة افتراضية لا تؤدي أي إجراء.

يستعرض المثال أدناه نسخة مختصرة من ColwiseParallel.

1. ورث من TensorParallelLayer وخزّن المواضع التي يعمل بها هذا الاستراتيجي. لا يحدد الصنف الأساسي أي __init__، لذا لا يوجد ما يُستدعى عليه super().

   تم النسخ

   class ColwiseParallel(TensorParallelLayer):
       def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True):
           self.input_layouts = input_layouts or Replicate()
           self.output_layouts = output_layouts if output_layouts is not None else Shard(-1)
           self.use_local_output = use_local_output

2. نفّذ shard_param لتغليف معامل واحد كحامل مؤقت من نوع DTensor. يعمل هذا على متجهات meta، لذا يبني distribute_tensor البيانات الوصفية فقط دون نقل أي بيانات. مرّر src_data_rank=None لأنه لا يوجد متجه كامل بعد يمكن التشتت منه.

   تم النسخ

   def shard_param(self, module, param, mesh):
       meta = module._parameters.get(param)
       if meta is None:
           return
       # Output features live on dim 0 for a 2D weight and dim -1 for a 1D bias
       placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2)
       module._parameters[param] = torch.nn.Parameter(
           distribute_tensor(meta, mesh, [placement], src_data_rank=None),
           requires_grad=meta.requires_grad,
       )

3. نفّذ تحويلات الإدخال والإخراج. يستدعي install_forward هذه التحويلات قبل وبعد دالة forward الأصلية للوحدة، لذا لا يتعين عليها سوى نقل المصفوفات (tensors) بين التخطيطات. يفترض التقسيم على مستوى الأعمدة أن يكون الإدخال مكرراً (replicated) وينتج مخرجات مقسمة على البُعد الأخير.

   تم النسخ

   def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
       x = args[0]
       if not isinstance(x, DTensor):
           x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False)
       if x.placements != (Replicate(),):
           x = x.redistribute(placements=[Replicate()])
       return (x,) + args[1:], kwargs
   
   def transform_output_post_forward(self, module, output, mesh):
       if not isinstance(output, DTensor):
           output = DTensor.from_local(output, mesh, [Shard(-1)], run_check=False)
       if output.placements != (self.output_layouts,):
           output = output.redistribute(placements=[self.output_layouts])
       return output.to_local() if self.use_local_output else output

   يضيف ColwiseParallel المُرفق مسارات سريعة فوق هذا، تتخطى رحلة DTensor ذهاباً وإياباً للاستدلال على وحدات nn.Linear العادية، وللموديولات المُكمّاة التي تحتاج إلى متجهات محلية. اقرأ الكود المصدري قبل نسخه إذا كانت استراتيجيتك تحتاج إلى المعالجة نفسها.

4. سجّل الاستراتيجية حتى يتمكن tp_plan من الإشارة إليها بالاسم. يتطلب التسجيل نسخة من الاستراتيجية، لذا تُحفظ المواضع التي تمررها إلى __init__ مع الاسم. لاحظ () أدناه.

   تم النسخ

   import torch
   
   from transformers import AutoModelForCausalLM, DistributedConfig
   from transformers.distributed.parallel_utils import ALL_PARALLEL_STYLES
   
   ALL_PARALLEL_STYLES.register("colwise_custom", ColwiseParallel())
   tp_plan = {
       "model.layers.*.self_attn.q_proj": "colwise_custom",
       ...
   }
   distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
   model = AutoModelForCausalLM.from_pretrained(
       "meta-llama/Meta-Llama-3-8B-Instruct",
       dtype=torch.bfloat16,
       distributed_config=distributed_config,
   )

المعايير

يُسهم التوازي على مستوى الموترات بشكل كبير في تسريع الاستدلال، لا سيما عند استخدام أحجام دفع كبيرة أو تسلسلات طويلة.

تُظهر هذه المخطط تسريع الأداء المتوقع لتمريرة أمامية واحدة على Llama مع طول تسلسل يبلغ 512.

تنفيذ التصميم

تنفّذ مكتبة Transformers توازي المصفوفات بطريقة مستقلة عن الإطار البرمجي. وتعتمد على DeviceMesh وDTensor من torch.distributed لتوفير واجهة بسيطة وقابلة للتوسّع.

DeviceMesh

DeviceMesh ينشئ شبكة متعددة الأبعاد من الأجهزة تتواصل فيما بينها. تتطلب استراتيجيات التوازي المختلفة أنماط اتصال مختلفة. أنشئ DeviceMesh بشبكات فرعية متعددة للتعامل مع هذه الأنماط.

import torch
from torch.distributed.device_mesh import init_device_mesh

# Create a 1D mesh of 4 accelerators
device_type = torch.accelerator.current_accelerator().type
device_mesh = init_device_mesh(device_type, (4,), mesh_dim_names=["tp"])

تُطبَّق معظم استراتيجيات التوازي في torch.distributed على الشبكة نفسها أو على شبكات فرعية منها. وتتعامل الشبكة تلقائيًا مع أنماط الاتصال.

DTensor

يتعامل DTensor (الموتر الموزّع) مع منطق التوزيع فوق عمليات الموترات المعتادة. تُخزَّن أوزان النموذج تحت توازّي الموترات على شكل DTensors، وهو ما يتيح للإستراتيجية وصف التواصل كتغيير في التخطيط بدلاً من عملية جماعية صريحة.

يُخبر سمة placements مكتبة PyTorch كيفية توزيع متجه (tensor) عبر الأجهزة ضمن DeviceMesh. وتقبل القيم التالية:

• يقوم Shard(dimension) بتقسيم DTensor عبر بُعدٍ محدد على DeviceMesh الذي تم إنشاؤه في إطاره. يؤدي التقسيم العمودي إلى تجزئة بُعد الميزات الإخراجية للوزن والبُعد الوحيد للانحياز.

  تم النسخ

  weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(0)]) # Shard the output features
  bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Shard(-1)]) # Shard the ONLY dimension

  يقوم التقسيم على مستوى الصفوف بتجزئة بُعد الميزات المدخلة بدلاً من ذلك، مع تكرار الانحياز، لأن الانحياز يُضاف مرة واحدة بعد عملية الاختزال وليس على كل رتبة.

  تم النسخ

  weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(-1)]) # Shard the input features
  bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Replicate()]) # Replicate bias across all GPUs

• Replicate() يكرر DTensor عبر DeviceMesh، مما ينشئ نسخة كاملة من الموتر على كل جهاز.

• يُشير Partial() إلى أن المصفوفة (Tensor) بانتظار عملية اختزال (Reduction). يكون ناتج الدالة الأمامية لطبقة صفية (Row-wise) من نوع Partial()، وإعادة توزيعه إلى Replicate() هي ما يُصدر عملية الاختزال الشامل (All-reduce).

الموارد

• توفر قسم Ultra-Scale Playbook الخاص بالتوازي على مستوى المصفوفات مزيدًا من التفاصيل.

• راجع دليل التوازي بين الخبراء إذا كنت تستخدم نموذجاً من نوع مزيج الخبراء (MoE). تدعم هذه النماذج التوازي على مستوى المصفوفات والتوازي بين الخبراء.

• اقرأ منشور المدونة التوازي على مستوى المصفوفات (TP) في Transformers: خمس دقائق لفهمه للحصول على نظرة عامة سريعة على التوازي على مستوى المصفوفات، وتعلّم كيف تختلف إعدادات التوازي العمودي والأفقي.

• اطّلع على دليل التدريب Tensor parallelism لمعرفة كيفية استخدامه في إعداد التدريب.