🤗 accelerate'ta dizisel paralellik

Bu kılavuz, 🤗accelerate içinde dizisel paralelliğin (sequence parallelism) kullanımının temellerini ele alacaktır.

Çok ilgili olan Context Parallelism'e de bakın.

Neden dizilim paralelliği?

Büyük dil modellerinin ve son zamanlarda akıl yürütme modellerinin ortaya çıkmasıyla birlikte, dizi uzunluğu hızla artmaktadır. Buna, dikkat mekanizmasının karesel bellek karmaşıklığının eklenmesi, uzun dizilere sahip modelleri eğitmek için daha verimli yöntemlere duyulan ihtiyacı doğurmuştur. 128k dizi uzunluğunda, standart dikkat uygulaması varsayıldığında, dikkat matrisinin bellek gereksinimi bf16 hassasiyetinde 128k * 128k * 2 bytes * num_heads = ~32 GB * num_heads'dir. Elbette, bu dikkat ağırlıklarını somutlaştırmayan flash attention veya SDPA kullanıldığında bu gereksinim dramatik biçimde azalır, ancak bellek gereksinimlerindeki artış hâlâ göz ardı edilemez düzeydedir.

Ulysses Sequence parallelism, dikkat (attention) hesabına giren girdilerin sıra (sequence) boyutu boyunca parçalanmasına ve dikkatin normal şekilde hesaplanmasına olanak tanır; ancak her GPU'da yalnızca dikkat başlıklarının (attention heads) bir dilimi kullanılır. Bu sayede, birkaç ek araçla birlikte, 15M+ sıra uzunluğuna ölçeklenebilen uzun sıralı modeller eğitilebilir. Ulysses SP'nin TiledMLP, Liger-Kernel, aktivite kontrol noktası (activation checkpoint) CPU'ya aktarımı ve birkaç başka hile ile nasıl güçlendirileceğini görmek için lütfen şu makaleye başvurun: Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences.

Ulysses SP, FSDP CP'den nasıl farklıdır

Context Parallelism belgesinde, başlık boyutunda dilimleme yerine Ring Attention kullanan ve yine dizi boyutunda dilimleme yapan Context Parallelism adlı başka bir teknolojinin dağıtımı hakkında bilgi edinebilirsiniz.

Aşağıdaki makaleler, bu iki teknoloji arasındaki farkları çok ayrıntılı bir şekilde açıklar:

• https://insujang.github.io/2024-01-11/tensor-parallelism-and-sequence-parallelism-detailed-analysis/
• https://huggingface.co/blog/exploding-gradients/ulysses-ring-attention

Bir makaleden uyarlanan kısa bir özet:

• Ulysses SP nispeten düşük bir iletişim yüküne sahiptir, ancak Attention Heads sayısıyla sınırlıdır ve bu nedenle ağ topolojisi için belirli gereksinimleri vardır (tek bir replika için attention heads sayısı, katılan GPU sayısına bölünebilmelidir). All-to-all iletişimi gecikmeye duyarlıdır ve Deepspeed gerektirir.
• FSDP CP Ring-Attention'ın P2P halka iletişimi, yukarıda belirtilen bölünebilirlik gereksinimlerine sahip değildir, ancak daha yüksek bir iletişim hacmine sahiptir.

Son olarak, makalede USP: A Unified Sequence Parallelism Approach for Long Context Generative AI açıklandığı gibi SP + CP'nin birleştirilmesi, daha uzun bir dizi uzunluğunu desteklemek için mümkün olmalıdır; ancak bu henüz 🤗accelerate içine entegre edilmemiştir.

Desteklenen dizi paralelliği (sequence parallelism) arka uçları

Şu anda tek dizi paralelliği arka ucu deepspeed'tir; bu, Arctic Long Sequence Training teknolojisi'nin bir parçası olan modernize edilmiş Ulysses SP'den gelir. Kendi kodunuza doğrudan entegre etmek isterseniz bir öğretici de mevcuttur.

Sıralı paralellik nasıl kullanılır?

from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

+# Example: 4 GPUs with sp_size=4, dp_shard_size=1
+# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
+     sp_backend="deepspeed",
+     sp_size=4,
+     dp_shard_size=1,  # Explicit: no data parallelism
+     sp_handler=DeepSpeedSequenceParallelConfig(
+         sp_seq_length_is_variable: true,
+         sp_attn_implementation="sdpa",
+     ),
+ )

accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

🤗accelerate içindeki diğer tüm özelliklerde olduğu gibi, sıralı paralelliği de accelerate launch komutuna ilgili bayrakları geçirerek etkinleştirebilirsiniz. Bu durumda durum farklı değildir:

accelerate launch --parallelism-config-sp-size 8  ...

Ayrıca sp_size ve diğer yapılandırmaları accelerate config komutunda da ayarlayabilirsiniz; bunlar accelerate yapılandırma dosyanızda saklanır, böylece her betiğinizi başlattığınızda bunları geçmenize gerek kalmaz.

sequence parallelism, data parallelism ile birleştirilir. Ek GPU gerektirmez. Dolayısıyla 8 GPU'nuz varsa şunu yapabilirsiniz: --parallelism-config-dp-shard-size 8 --parallelism-config-sp-size 8. Veya bunları programatik olarak ayarlamak için ParallelismConfig sınıfını kullanabilirsiniz.

Önemli: dp_replicate_size × dp_shard_size × sp_size = num_processes eşitliğinin sağlandığından emin olmalısınız. Örneğin, 8 GPU ve sp_size=8 kullanıldığında dp_shard_size=1 gereklidir (çünkü 1 × 1 × 8 = 8). 4 GPU ve sp_size=2 kullanıldığında, 2B paralellik için dp_shard_size=2 kullanılabilir (çünkü 1 × 2 × 2 = 4).

ALST/Ulysses SP arka uç yapılandırması

ALST/UlyssesSP, bu makalede açıklandığı gibi, dikkat baş paralelliği kullanarak dizi paralelliğini uygular. Basitlik için, kullanıcı açısından aynı olan dizi paralelliğinin kavramını ve kurulumunu yeniden kullanırız: tek bir partiyi işlemek için birden çok GPU kullanılır.

ALST'nin neyi mümkün kıldığına dair bir fikir vermek için: bu sayede tek bir H100 GPU üzerinde 500K token ile bf16'da eğitim yapabildik, tek bir düğümde 3.7M, Llama-8B için ise yalnızca dört düğüm kullanarak 15M tokena ulaştık. HF Accelerate'in bu özelliği, ALST'nin üç bileşeninden yalnızca 1'ini etkinleştirir, bu nedenle elde edilebilir dizi uzunluğu daha küçük olacaktır. ALST'nin tam gücünden yararlanmak için TiledMLP, CPU'ya aktivasyon checkpoint offload ve birkaç diğer özelliğin etkinleştirilmesini istersiniz. Detaylar için lütfen bu öğreticiye başvurun.

deepspeed arka ucunu yapılandırmak için:

# Example: 4 GPUs with sp_size=4, dp_shard_size=1
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,
    dp_shard_size=1,  # Explicit: no data parallelism
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length=256,
        sp_seq_length_is_variable=True,
        sp_attn_implementation="sdpa",
    ),
)
accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

• sp_backend: burada deepspeed olarak ayarlayın
• sp_size, sıralı paralellik derecesidir - yukarıdaki örnekte 4'tür, bu nedenle tek bir batch'i işlemek için 4 GPU kullanılacaktır (aynı GPU'lar üzerinde DP=4 uygulanırken)
• sp_seq_length ve sp_seq_length_is_variable dizi uzunluklarıyla başa çıkmak için kullanılır. sp_seq_length_is_variable=True ise backend, batch'ler arasında değişebilecek bir dizi uzunluğuyla çalışır; bu durumda sp_seq_length değeri, dizide paralel derekeye bölünebilir bir değere ayarlanabilir ya da hiç ayarlanmayabilir. Bu durumda her forward çağrısında dizideki değişkenler girdiden türetilir. False ise seq_length, batch'in dizideki uzunluk boyutuyla eşleşmelidir ve bu boyut her zaman aynı olacak şekilde dolgu (padding) ile tamamlanmalıdır. Varsayılan değer True'dur.
• sp_attn_implementation, sdpa, flash_attention_2 veya flash_attention_3 değerlerinden biridir. Bu sequence parallel uygulaması attention_mask yerine position_ids kullandığından, eager position_ids ile çalışma desteğine sahip olana kadar burada çalışamaz. Ayrıca, sdpa'nin tek bir örnekte birleştirilen birden fazla örneği doğru şekilde işlemediğini, tüm örneği tek bir bütün olarak ele alacağını lütfen unutmayın. Örnekler birleştirilmediyse sdpa doğru şekilde çalışır. Bu nedenle, her zaman çalışan Flash Attention ideal seçim olmalıdır.

Bu değerleri DeepSpeedSequenceParallelConfig nesnesinde ayarlamak yerine, aynı sonucu ortam değişkenlerini kullanarak da elde edebilirsiniz - bunlar yukarıdaki listenin sonuna karşılık gelir.

• PARALLELISM_CONFIG_SP_BACKEND
• PARALLELISM_CONFIG_SP_SEQ_LENGTH
• PARALLELISM_CONFIG_SP_SEQ_LENGTH_IS_VARIABLE
• PARALLELISM_CONFIG_SP_ATTN_IMPLEMENTATION

Kodda geçirilmezse, sp_size, --parallelism_config_sp_size komut satırı argümanı ile ayarlanabilir. Diğer argümanlar için de durum aynıdır. Ayrıca accelerate config dosyası tarzında yapılandırma da yapabilirsiniz, örneğin 2 GPU için:

distributed_type: DEEPSPEED
deepspeed_config:
  deepspeed_config_file: path/to/ds_config.json
machine_rank: 0
num_machines: 1
num_processes: 2
parallelism_config:
  parallelism_config_dp_replicate_size: 1
  parallelism_config_dp_shard_size: 1  # Must satisfy: 1 × 1 × 2 = 2 num_processes
  parallelism_config_sp_size: 2
  parallelism_config_sp_backend: deepspeed
  parallelism_config_sp_seq_length_is_variable: true
  parallelism_config_sp_attn_implementation: sdpa

Daha önce belirtildiği gibi Ulysses sequence parallelism genellikle data parallelism ile birlikte kullanılır - benzersiz veri akışlarını beslemek için aynı rank'lar kullanılır ve ayrıca Ulysses Sequence Parallelism de gerçekleştirir. Ancak replikaları şu şekilde oluşturabilirsiniz:

# Example: 4 GPUs with 2D parallelism (SP=2, DP=2)
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 2 × 1 × 2 = 4 GPUs
parallelism_config = ParallelismConfig(
    dp_replicate_size=2,
    dp_shard_size=1,  # Explicit: no sharding within replicas
    sp_size=2,
    sp_backend="deepspeed",
    sp_handler=DeepSpeedSequenceParallelConfig(...),
)

Burada 4 GPU kullanıyoruz ve 2 dizi paralelliği replikası var. Veri paralelliğini burada sürükleyen Deepspeed-ZeRO'dur.

Lütfen UlyssesSPDataLoaderAdapter içinde çok fazla sihirli gizli olduğunu unutmayın. Arka planda, orijinal DataLoader nesnenizi sarmalayarak kullanılır, ancak herhangi bir sorunla karşılaşırsanız bunun farkında olmalısınız. Ayrıca, batch katılımcı ranklar arasında parçalanmadan önce, batch sözlüğüne doğru shift_labels değerini otomatik olarak enjekte eder.

ALST/UlyssesSP kullanmaya başlamak için geriye kalan tek parça, gradyanları doğru almak için kayıpları farklılaştırılabilir bir all_gather ile ranklar arasında toplamaktır. Aşağıdaki kod bunu yaparken, doğru ortalama değerini almak için -100 ile maskelenmiş tokenları da hariç tutar:

sp_size = parallelism_config.sp_size if parallelism_config is not None else 1
if sp_size > 1:
    sp_group = accelerator.torch_device_mesh["sp"].get_group()
    sp_world_size = parallelism_config.sp_size

# Normal training loop
for iter, batch in enumerate(dl):
    optimizer.zero_grad()

    batch = move_to_device(batch, model.device)

    # The model automatically receives shift_labels via **kwargs and uses it for loss computation.
    # Both standard transformers models and Liger-patched models handle this correctly.
    outputs = model(**batch)
    loss = outputs.loss
    shift_labels = batch["shift_labels"]

    if sp_size > 1:
        # differentiable weighted per-shard-loss aggregation across ranks
        losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
        # special dealing with SFT that has prompt tokens that aren't used in loss computation
        good_tokens = (shift_labels != -100).view(-1).sum()
        good_tokens_per_rank = torch.distributed.nn.functional.all_gather(
            good_tokens, group=sp_group
        )
        # Skip ranks with zero valid tokens to avoid NaN contamination (NaN * 0 = NaN)
        total_loss = sum(
            losses_per_rank[rank] * good_tokens_per_rank[rank]
            for rank in range(sp_world_size)
            if good_tokens_per_rank[rank] > 0
        )
        total_good_tokens = sum(good_tokens_per_rank)
        loss = total_loss / max(total_good_tokens, 1)

    if rank == 0: accelerator.print(f"{iter}: {loss=}")
    accelerator.log(dict(train_loss=loss, step=iter))

    accelerator.backward(loss)
    optimizer.step()

Partilerde shift_labels mevcut olduğunda modellerin bunu otomatik olarak ele aldığını unutmayın. Modelin ileri geçişi, shift_labels değerini **kwargs üzerinden alır ve bunu kayıp fonksiyonuna iletir; kayıp fonksiyonu da dizisel paralellik için kaybı doğru şekilde hesaplar. Liger Kernel kullanıyorsanız, o da shift_labels değerini sorunsuz bir şekilde ele alır ve kaybı çok bellek verimli bir şekilde hesaplar. Liger, uzun dizi uzunlukları için şiddetle önerilir; çünkü füzyonlu işlemler (örneğin, bellekte tam logits tensörünü asla maddi olarak oluşturmayan füzyonlu logit-kayıp hesabı) kullanarak GPU belleğini serbest bırakır.

HF Accelerate'in arka planda ne yaptığını görmek istiyorsanız lütfen bu tam entegratör kılavuzunu okuyun.

ALST/UlyssesSP etkinleştirilmiş bir Accelerate eğitim döngüsünün örneği için examples/alst_ulysses_sequence_parallelism dosyasına bakın.

[!Uyarı]

Bu API oldukça yenidir ve hâlâ deneysel aşamasındadır. Kararlı bir API sunmaya çalışsak da, genel API'nin bazı küçük kısımları gelecekte değişebilir.

Bu bir DeepSpeed arka ucu olduğu için standart DeepSpeed yapılandırması geçerlidir; bu nedenle daha fazla GPU belleği açığa çıkarmak ve daha uzun bir dizi uzunluğunu mümkün kılmak için dizi paralelliğini, optimizör durumları ve/veya ağırlıkların dışa aktarılmasıyla da birleştirebilirsiniz. Bu teknolojinin DeepSpeed ZeRO 2. ve 3. aşamalarıyla çalıştığı test edilmiştir.