Çıkarım için tensör paralelliği

Tensor paralelliği, bir model katmanını parçalara ayırarak birden fazla donanım hızlandırıcısının aynı anda üzerinde çalışmasını sağlar. Bu sayede, tek bir GPU'nun bellek kapasitesini aşan modelleri çalıştırabilir ve daha yüksek bir işleme hızı elde edebilirsiniz. GPU'lar her katmanda kısmi sonuçları birbirleriyle paylaştığı için, düğüm içi hızlı iletişim gereklidir.

Bir model, config'i base_model_tp_plan tanımlıyorsa tensor paralelliğini destekler. Yüklenmiş bir modelin supports_tp_plan özelliğini kontrol edin.

from transformers import AutoModelForCausalLM

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

Bu kılavuz, Transformers'ta tensor paralelliğini etkinleştirmeyi ve kullanılabilir bölme stratejilerini kapsar.

Bir modelin bölünmesi

DistributedConfig içinde tp_size ile tensor paralel cihaz sayısını yapılandırın.

• Modelin önceden tanımlı planını kullanmak için DistributedConfig(tp_size=N) ayarlayın.
• Manuel bir tp_plan tanımlayın ve tp_size ile birlikte DistributedConfig nesnesine iletin.

tp_plan="auto" değerini ayarlayarak önceden tanımlanmış planı açıkça talep edebilirsiniz. tp_size atlandığında ve bir tp_plan ayarlandığında, tp_size değeri WORLD_SIZE'in diğer paralel boyutlara bölünmesiyle türetilir. tp_plan değerini doğrudan from_pretrained() fonksiyonuna geçirmek artık kullanılmıyor (deprecated) olup v5.18 sürümünde kaldırılacaktır.

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)

Bölme stratejileri

ALL_PARALLEL_STYLES, distributed/parallel_utils.py dosyasında tanımlanan bir ParallelInterface, bir planda kullanabileceğiniz her strateji adını yapılandırılmış bir strateji örneğine eşler. Bir modeli parçalamak için doğrudan onunla etkileşime girmeseniz de, kullanılabilir adların yetkili listesidir. distributed/tensor_parallel.py dosyası, tensor paralel stratejilerini bu listeye kaydeder:

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

Her tensor paralel strateji, o dosyadaki TensorParallelLayer sınıfının bir alt sınıfıdır. ep_dispatch_experts gibi uzman paralel stratejiler distributed/expert_parallel.py dosyasında yer alır, aynı hook'lara sahip olan ExpertParallelLayer sınıfının alt sınıfıdır ve aynı ALL_PARALLEL_STYLES yapısına kaydedilir. En sık başvuracağınız iki strateji olan ColwiseParallel ve RowwiseParallel, input_layouts ve output_layouts yerleşimlerini alır; bu sayede tek bir sınıf birden fazla plan adını kapsar. colwise, çıktısını son boyutta parçalı (sharded) bırakır ve bunu tüketmesi için arkasından gelen bir rowwise katmanına bırakır; colwise_gather_output ise çıktıyı all-gather ile tam bir tensöre geri toplar.

Aşağıdaki tablo her stratejiyi açıklar.

Paketlenmiş stratejiler

Weight packing, birden fazla lineer katmanı tek bir, daha büyük katmanda birleştirir. PackedColwiseParallel ve PackedRowwiseParallel stratejileri, paketlenmiş ağırlıkları doğru şekilde parçalar. Temel ColwiseParallel veya RowwiseParallel stratejileri ise paketlenmiş ağırlıkları yanlış şekilde parçalar.

Aşağıdaki örnek, up_proj ve gate_proj bileşenlerini tek bir gate_up_proj modülünde birleştirir ve gate_up_proj modülünün parçalanması için packed_rowwise stratejisini gerektirir.

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

gate_up_proj modülünün çıktısını hesaplamak için forward geçişinde toplu matris çarpımı kullanın.

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

Sade bir Shard, bu boyutu ardışık bloklara böler; bu nedenle bir rank, her birinin dilimi yerine gate'in sonunu ve up'in başını alır. Paketli stratejiler, bölünmeyi iç içe geçirmek için _StridedShard'ü split_factor ile (varsayılan olarak 2) kullanır; böylece her rank, her iki yarıdan da eşleşen bir dilim alır ve chunk, shard'landırmadan sonra hâlâ hizalı kalır.

Özel bölme stratejileri

distributed/tensor_parallel.py içindeki TensorParallelLayer sınıfından türetme yaparak özel bir bölme stratejisi oluşturun. Stratejinizin ihtiyaç duyduğu kanca (hook) işlevlerini geçersiz kılın; her birinin varsayılan olarak hiçbir şey yapmayan (no-op) bir uygulaması vardır.

Aşağıdaki örnek, ColwiseParallel'in sadeleştirilmiş bir versiyonunu adım adım ele alıyor.

1. TensorParallelLayer sınıfından türetin ve stratejinin çalıştığı yerleşimleri saklayın. Taban sınıf __init__ tanımlamadığı için super() çağrılacak bir şey yoktur.

   Kopyalandı

   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 işlevini uygulayarak tek bir parametreyi DTensor yer tutucusu olarak sarmalayın. Bu işlev meta tensörler üzerinde çalıştığından, distribute_tensor yalnızca metadata oluşturur ve veri taşımaz. Henüz dağıtılacak tam bir tensör olmadığı için src_data_rank=None parametresini iletin.

   Kopyalandı

   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. Giriş ve çıktı dönüşümlerini uygulayın. install_forward, bunları modülün orijinal forward'ının etrafında çağırır, bu nedenle yalnızca tensörleri yerleşimler arasında taşımaları gerekir. Sütun bazlı bölme, replike edilmiş bir giriş bekler ve son boyutta parçalanmış bir çıktı üretir.

   Kopyalandı

   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

   Hazır gelen ColwiseParallel, düz nn.Linear çıkarımı ve yerel tensörler gerektiren kuantize modüller için DTensor gidiş-dönüş yolunu atlayan hızlı yollar ekler. Stratejiniz aynı muameleyi gerektiriyorsa, kopyalamadan önce kaynağı okuyun.

4. Stratejiyi, tp_plan'ın onu adlandırabilmesi için kaydedin. Kayıt işlemi bir strateji örneği alır, bu nedenle __init__'e ilettiğiniz yerleşimler, adla birlikte saklanır. Aşağıdaki () ifadesine dikkat edin.

   Kopyalandı

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

Benchmarks

Tensor parallelism, özellikle büyük batch boyutları veya uzun diziler için çıkarımı önemli ölçüde hızlandırır.

Bu grafik, 512 dizi uzunluğuna sahip Llama için tek bir ileri geçişte beklenen hızlanmayı göstermektedir.

Uygulama tasarımı

Transformers, tensor paralellizmini çerçeveye bağımsız bir şekilde uygular. Basit ve genişletilebilir bir arayüz sağlamak için torch.distributed içindeki DeviceMesh ve DTensor sınıflarına dayanır.

DeviceMesh

DeviceMesh, birlikte iletişim kuran cihazlardan oluşan çok boyutlu bir ızgara oluşturur. Farklı paralelleştirme stratejileri farklı iletişim kalıpları gerektirir. Bu kalıpları yönetmek için birden fazla alt ızgaraya sahip bir DeviceMesh oluşturun.

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"])

Çoğu torch.distributed paralelleştirme stratejisi, mesh'in kendisine veya alt-mesh'ine uygulanır. Mesh, iletişim desenlerini otomatik olarak yönetir.

DTensor

DTensor (Distributed Tensor), olağan tensor işlemlerinin üzerine dağıtık mantığı işler. Tensor paralelliği altındaki model ağırlıkları DTensor olarak saklanır; bu da bir stratejinin iletişimi açık bir kolektif işlem yerine yerleşim değişikliği olarak tanımlamasını sağlar.

placements özniteliği, bir DeviceMesh içindeki cihazlar arasında bir tensörün nasıl yerleştirileceğini PyTorch'a bildirir. Aşağıdaki değerleri kabul eder:

• Shard(dimension), DTensor nesnesini, oluşturulduğu DeviceMesh üzerinde verilen boyut boyunca böler. Sütun bazlı bölme, ağırlığın çıktı-özellik boyutunu ve sapmanın tek boyutunu parçalar.

  Kopyalandı

  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

  Satır bazlı bölme, girdi-özellik boyutunu parçalara ayırır ve önyargıyı kopyalar; çünkü önyargı her sırada değil, indirgeme işleminden sonra bir kez eklenir.

  Kopyalandı

  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 nesnesini DeviceMesh üzerinde çoğaltarak her cihazda tensörün tam bir kopyasını oluşturur.

• Partial(), bir tensörün indirgeme (reduction) işlemi beklediğini belirtir. Satır bazlı bir katmanın ileri (forward) çıktısı Partial() olur ve bu çıktıyı Replicate() durumuna yeniden dağıtmak, all-reduce işlemini tetikler.

Kaynaklar

• Tensor paralelliği hakkında daha fazla ayrıntı, Ultra-Scale Playbook bölümünde yer alır.

• Mixture-of-experts (MoE) bir model kullanıyorsanız expert parallelism kılavuzunu kontrol edin. Bu modeller tensor parallelism ve expert parallelism'i destekler.

• Tensor paralellizminin kısa bir özetini almak ve sütun ile satır paralel yapıların nasıl farklılaştığını öğrenmek için Transformers'ta Tensor Parallelism (TP): 5 Dakikada Anlayın blog yazısını okuyun.

• Eğitim ortamında nasıl kullanılacağını öğrenmek için Tensor parallelism eğitim rehberine bakın.