🤗 accelerate'ta Bağlam Paralelliği

Bu kılavuz, 🤗accelerate içinde bağlam paralelliğinin kullanımının temellerini ele alacaktır; meraklı okuyucular için sonraki bölümlerde bazı teknik detaylara da değineceğiz.

Çok ilgili olan Sıralı Paralellik Kılavuzu'na da bakın.

Neden bağlam 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.

Bağlam paralelliği, dikkat hesabına giren girdileri sıra boyutu boyunca parçalara ayırıp dikkati birden fazla GPU üzerinde paralel olarak hesaplamamızı sağlar. Bu sayede, potansiyel olarak 1M+ sıra uzunluğuna ölçeklenebilen uzun sıralı modeller eğitebiliriz.

Bağlam paralelliği nasıl kullanılır?

from accelerate.utils import ParallelismConfig, TorchContextParallelConfig

+ cp_config = TorchContextParallelConfig(
+       cp_comm_strategy="alltoall", # no need to use cp_config at all, if you want to use the default "allgather"
+ )

+ parallelism_config = ParallelismConfig(
+     cp_size=8,
+     cp_handler=cp_config,  # or just cp_size=8, if you want to use the default "allgather"
+ )

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

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

accelerate launch --parallelism-config-cp-size 8 --parallelism-config-cp-comm-strategy [allgather|alltoall] ...

Ayrıca cp_size ve cp_comm_strategy değerlerini accelerate config komutunda da ayarlayabilirsiniz; bu değerler accelerate yapılandırma dosyanızda saklanır, böylece her betiğinizi başlattığınızda bunları tekrar iletmenize gerek kalmaz.

Bağlam paralelliği, veri paralelliği, tensör paralelliği ve FSDP2 gibi diğer paralelleştirme stratejileriyle uyumludur. Paralelleştirme boyutlarını istenen değerlere ayarlayarak bunları kolayca birleştirebilirsiniz, örneğin --parallelism-config-dp-size 8 --parallelism-config-tp-size 2 --parallelism-config-cp-size 8. Veya bunları programatik olarak ayarlamak için ParallelismConfig sınıfını kullanabilirsiniz.

Bağlam paralelliği, FSDP2 tanıtımı bölümünde daha fazla bilgi edinebileceğiniz FSDP2 ile sıkı bir şekilde ilişkilidir. Bu, bağlam paralelliğinin yalnızca programınıza FullyShardedDataParallelPlugin veya sürüm 2 olarak ayarlanmış --use-fsdp ile kullanıldığında çalıştığı anlamına gelir. FSDP2 kullanılmıyorsa bir hata oluşturulur.

Bağlam paralelliği yalnızca SDPA ile ve maske veya nedensel maske olmadan çalışır. Bunu sizin için doğru bir şekilde algılayamayız, bu nedenle SDPA’yi maske veya nedensel maske olmadan kullandığınızdan emin olmanız sizin sorumluluğunuzdadır. Başka bir dikkat uygulaması kullanırsanız bir hata oluşturulur.

Yukarıda bahsedilen yöntemlerle bağlam paralelliğini etkinleştirdikten sonra, bunu eğitim döngünüze uygulayabilirsiniz. Eğitim döngünüzde kullanabileceğiniz, kullanımındaki bazı karmaşıklıkları soyutlayan torch.distributed.tensor.experimental.context_parallel etrafında ince bir sarmalayıcı sunuyoruz (bununla ilgili daha fazla bilgiyi daha sonra paylaşacağız). Eğitim döngünüzde yapmanız gereken değişiklikleri en aza indirmek için, bağlam paralelliği etkinleştirilmemişse noop olan, etkinleştirilmişse bağlam paralelliğini uygulayan bir bağlam yöneticisi sunuyoruz. Bu sayede, paralelleştirme yapılandırmanıza bağlı olarak herhangi bir kod değişikliği yapmadan eğitim döngünüzde kullanabilirsiniz. Bunu şu şekilde kullanabilirsiniz:

for batch in dataloader:
    with accelerator.maybe_context_parallel(
        buffers=[batch["input_ids"], batch["attention_mask"]],
        buffer_seq_dims=[1, 1],
        no_restore_buffers={batch["input_ids"], batch["labels"]},
    ):
        outputs = model(**batch)
        ...

Bu bağlam yöneticisi, yukarıdaki örnekte gösterildiği gibi her eğitim adımıyla yeniden oluşturulmalıdır. Bunu yapmak son derece önemlidir.

Bu, bağlam boyutunu potansiyel olarak 1M+ dizi uzunluğuna kadar ölçeklendirebilir. Aşağıda, 256k'ye kadar bağlam boyutu için bağlam paralelliğinin hızını ve bellek kullanımını gösteriyoruz. Bağlam boyutunu ve GPU sayısını ikiye katladığımızda tutarlı bir bellek kullanımı elde edebildiğimizi görüyoruz; bu da potansiyel olarak sınırsız bağlam uzunluğu ölçeklendirmesini mümkün kılabilir.

Şekil 1: 256k'ye kadar bağlam boyutu için bağlam paralelliğinin bellek kullanımı ve hızı.

Bu örnekler, örnekler klasöründe bulabileceğiniz bir betik ile oluşturulmuştur. Örneği 8 H100 GPU üzerinde (128k dizi uzunluğu) çalıştırmak için aşağıdaki komutu kullanabilirsiniz:

Kopyalandı

accelerate launch --use-fsdp --fsdp-activation-checkpointing=TRUE examples/torch_native_parallelism/nd_parallel.py --cp-size=8 --sequence-length=128000

Accelerate arayüzü

Bağlam yöneticisi, bağlam paralelliğini yapılandırmak için kullanılan birkaç argüman alır.

• buffers: Sıralama boyutu boyunca parçalanacak tensörlerin listesi. Bu tensörler genellikle girdi kimlikleri, etiketler ve dikkat maskesidir.
• buffer_seq_dims: buffers listesinin sırasına göre tamponların dizi boyutunu belirleyen tamsayı listesidir. buffers=[input_ids, shift_labels] parametresini, her ikisinin de [batch_size, sequence_length] şekline sahip olduğu durumda, tensörlerin dizi boyutu ikinci boyut olduğu için buffer_seq_dims=[1, 1] olarak iletirsiniz. Bu, model çıktılarının doğru hesaplanması için gereklidir.
• no_restore_buffers: Bağlam paralelliğinin uygulaması, tamponları yerinde değiştirerek onları torch.distributed.tensor.Dtensor türüne dönüştürür. Bağlam yöneticisi sona erdikten sonra, tamponların orijinal durumuna (genellikle all-gather) geri yüklenmesi için bir iletişim çekirdeğinin başlatılması gerekir. Bu işlem biraz zaman aldığından, bağlam yöneticisi sona erdikten sonra tamponları kullanmanız gerektiğinden emin olmadıkça, gereksiz iletişimi önlemek için buffers argümanındakiyle aynı tensörleri geçmeniz önerilir.

Bağlam paralelliği, 🤗 transformers modellerinin nedensel dil modellemeyi etkinleştirmek için kendilerinin kaydırabildiği, input_ids'in bir kopyası olan labels ile uyumlu değildir. Bu durumu hayal edin: labels = [l1, l2, l3, l4, … li] bağlam paralelliği uygulandığında, her rank etiketlerin bir kısmına sahip olur, örneğin: labels_rank_0 = [l1, l2], labels_rank_1 = [l3, l4], … transformers modelleme kodu etiketleri kaydırdıktan sonra, sonuç şu olur: labels_rank_0 = [l2, PAD], labels_rank_1 = [l3, PAD], … burada PAD bir dolgu işaretidir. Bu, etiketlerin artık girdilerle hizalanmamış olması nedeniyle hatalı kayıp hesabına yol açar. Bu nedenle, etiketleri modele geçirmeden önce manuel olarak kaydırmanız gerekir

Yapılandırılabilir seçenekler

Accelerate, bağlam paralelliğini yapılandırmak için yalnızca tek bir seçenek sunar (cp_size hariç)

• cp_comm_strategy: Parçalar için kullanılacak rotasyon yöntemi. Çoğu durumda "alltoall" performansını büyük olasılıkla aşacağından, bu değeri "allgather" olarak bırakmanızı şiddetle öneririz.

Context parallel size oldukça açıklayıcıdır; girdilerin shard edileceği rank sayısıdır. Context parallel shard rotation, girdilerin shard'larının rank'lar arasında nasıl döndürüleceğini tanımlar. 2 seçeneği bir sonraki bölümde daha ayrıntılı olarak ele alacağız.

ND parallel example dosyasında uçtan uca bir örneği görebilirsiniz; bu örnekte tek bir 8xH100 düğümünde, 128k'ye kadar bağlam uzunluğuna sahip bir 8B modeli eğitebilirsiniz. Çoklu düğüm eğitimi kullanarak bu ölçeği, birden çok GPU üzerinde 1M+ dizi uzunluğuna çıkarabilirsiniz. Ayrıca ihtiyaçlarınıza göre bunu diğer paralelleştirme stratejileriyle sorunsuz bir şekilde birleştirebilirsiniz.

Teknik ayrıntılar

Bu bölüm oldukça teknik olduğu için bağlam paralelliğinin iç yapısını öğrenmeniz gerekmiyorsa, bu bölümü atlayıp inşa etmeye başlayabilirsiniz 🚀

Aşağıdaki bölümlerde shard kelimesini yoğun şekilde kullanacağız, bu yüzden önce onu tanımlayalım. Bir tensörün Dth boyutu boyunca, N rank üzerinde sharded (parçalara ayrılmış) olduğunu söylediğimizde, bu tensörün N parçaya bölündüğünü, tensörün her bir parçasının şeklinin ise [..., D//N, ...] olduğunu kastediyoruz.

Peki nasıl çalışır?

Bağlam paralelliği, Q, K and V matrislerinin dizi boyutu boyunca parçalanarak dağıtılması üzerinde çalışır. Her rank, kendisine atanmış Q parçasına sahiptir; buna Q_i diyelim. Bu matris, hesaplamanın tamamı boyunca yalnızca bu rank üzerinde kalır. Benzer şekilde, her rank kendi K ve V parçalarına sahiptir; bunlara K_i ve V_i diyelim. Ardından, her rank kendi Q_i, K_i ve V_i parçalarıyla dikkat (attention) hesaplar; buna attn_i diyelim. Bu hesaplama sırasında, diğer tüm ranklardan Ks ve Vs değerlerini toplamak için bir iletişim çekirdeği başlatılır. Hangi iletişim ilkelinin kullanılacağı, context_parallel_shard_rotation seçeneğine bağlıdır. Bu sayede, her rank önce Q_i, K_i ve V_i ile, ardından diğer tüm ranklardan gelen K_j ve V_j ile yerel dikkat hesaplaması yapar. Her rank, dizi boyutu boyunca parçalanmış Q, K and V matrislerini barındırdığından, sonuç matrisleri daha küçük olur ve tek bir GPU’ya sığabilir.

Bunu aşağıdaki sözde kodda formalize edebiliriz:

comm_kernel = {"allgather": allgather, "alltoall": alltoall}[context_parallel_shard_rotation]
Qi, Ki, Vi = shard(Q, K, V, seq_dim)
attn[i] = attn(Qi, Ki, Vi)
for j in range(context_parallel_size):
    Kj, Vj = comm_kernel()
    attn[j] = attn(Qi, Kj, Vj) # [batch, num_heads, seq_len // context_parallel_size, head_dim]

final_attn = combine(attn)

all-to-all ve all-gather karşılaştırması

all-gather

Peki, all-to-all ile all-gather arasındaki fark nedir? All-gather'da iletişim çok basittir. (Aslında, genellikle daha uzun sürdüğü için) yerel dikkat attn_i hesabını tamamladıktan (daha doğrusu, başlamadan önce) sonra, diğer tüm rank'lardan gelen diğer Ks ve Vs değerlerini toplamak için bir all-gather başlatırız. Bu iletişim tamamlandığında, her rank diğer tüm rank'lardan gelen Ks ve Vs değerlerine sahip olur ve bunlarla dikkati sıralı olarak hesaplayabilir. İdeal senaryoda, all-gather, attn_i hesabının bittiği tam anda tamamlanır. Ancak pratikte bu asla gerçekleşmez; bu nedenle ideal gerçek örtüşme, tam attn_i hesabının iletişimin bir kısmıyla örtüşmesi ve ardından K_j ile V_j kullanılarak hesaplamanın başlatılması için all-gather'ın bitmesini beklememizle sağlanır.

all-to-all

All-to-all, bazen ring-rotation olarak da adlandırılır, halka benzeri bir iletişim deseni kullanır. attn_i hesabı tamamlandıktan sonra, K_i ve V_i değerlerini komşu sıralara göndermek için bir all-to-all başlatılır. Ardından bunu context_parallel_size-1 kez tekrarlarız, böylece her sıra, diğer tüm sıralardan gelen K ve V parçacıklarının tamamını bir kez görür. İdeal senaryoda, komşu sıradan K_i+1 ve V_i+1 parçacıklarını önceden getiririz ve bu iletişim, mevcut attn_i hesabıyla tam olarak örtüşür. Yine de gerçekçi olarak, bu mükemmel örtüşme asla gerçekleşmez. Yaklaşımın doğası gereği, mükemmel örtüşme sağlanamazsa, ceza all-gather'a kıyasla çok daha büyüktür.

Doğru döndürme yöntemini nasıl seçersiniz?

Teoride, all-to-all daha iyi bir seçim olmalıdır. Ancak pratikte bu nadiren gerçekleşir. Bu nedenle, genellikle daha iyi performans elde etme olasılığı daha yüksek olan all-gather varsayılan olarak kullanılır. torchtitan ekibinden gelen kapsamlı benchmarks da all-to-all'ın nadiren all-gather'dan daha iyi performans gösterdiğini göstermektedir. Yine de, kullanım senaryonuza göre birinin daha iyi olabileceğini düşündüğümüz için her iki seçeneği de sunuyoruz.

Aşağıdaki görüntüdeki profiler çıktısında bu sorunu doğrudan görebilirsiniz:

Şekil 1: Kırmızıyla gösterilen alan, all-to-all çekirdeğinin bitmesini beklediğimiz boşta geçen süredir. İlk mavi çubukta vurgulanan kısımda, bu işlemin tamamlanmasının ~250us sürdüğü görülmektedir; bu süre, N bağlam paralel boyutu olduğunda her dikkat çağrısı için N-1 kez tekrarlanır.

Neden yalnızca FSDP2?

Bağlam paralelliğini yalnızca FSDP2 ile destekliyoruz, çünkü tam potansiyelini kullanmak için context_parallel_size ve dp_shard_size için birleşik bir mesh oluşturuyoruz. Çalışma şekli şöyledir: modeli, bellek tasarrufunu en üst düzeye çıkaran cp_size*dp_shard_size boyutundaki birleşik mesh üzerinde parçalara ayırıyoruz. Bu, bir tür “bedava öğle yemeği”dir; çünkü aşağıdaki görsellerde gösterildiği gibi FSDP iletişimi, dikkat (attention) hesabıyla tamamen örtüşür.

Şekil 2: Mavi dikdörtgenlerde (Stream 23), `FSDP` shard'inin ön-alımının (pre-fetch), attention hesabıyla (Stream 7) tamamen örtüştüğünü görebilirsiniz; kırmızı dikdörtgenlerde (Stream 24) ise all-gather çekirdeğinin, hesaplama akışımızın (7) boşta kaldığı bir boş zaman baloncuklarına (bubble) yol açtığını görebilirsiniz.

Yukarıdaki şekilde, all-to-all ile all-gather arasındaki farkı da görebilirsiniz. All-to-all'da (Şekil 1) her dikkat çağrısı için N-1 kez bir iletişim çekirdeği başlatılırken, all-gather'da (Şekil 2) bir iletişim çekirdeği yalnızca bir kez başlatılır. Bu, daha büyük bir baloncuk (bubble) oluşmasına neden olur, ancak bu durum her dikkat çağrısında yalnızca bir kez gerçekleşir; all-to-all'da ise N-1 kez gerçekleşir.

Ortak mesh'te veri dağıtımı

Tüm cp alt grubuna aynı veri grubunu yönlendirdiğimizden emin oluruz, böylece sonuçlar doğru olur. (Yani cp alt grubundaki her sıra aynı veri grubunu alır.) Ancak dp_shard grubundaki her sıraya farklı veri grupları yönlendiririz. Bunu şöyle hayal edin:

# 8 GPUS, --dp_shard_size 4, --cp_size 2
# mesh = [[0, 1], [2, 3], [4, 5], [6, 7]]
# model is sharded across the whole mesh (each GPU holds 1/8 of the model)
# GPUs 0,1 = batch 0
# GPUs 2,3 = batch 1
... and so on.