Ring Attention

Halka dikkati, tek bir hızlandırıcıya sığmayacak kadar uzun bir diziyi halka biçiminde düzenlenmiş cihazlar arasında bloklara ayırarak tam öz-dikkat hesaplar. Her cihaz sorguların bir bloğunu tutar ve anahtar ile değer bloklarını komşusuna aktarırken bu iletişimi elinde tuttuğu blok üzerindeki hesaplamayla eşzamanlı yürütür; böylece bir modelin işleyebileceği en uzun dizi cihaz sayısıyla orantılı olarak uzar. UC Berkeley'den Hao Liu, Matei Zaharia ve Pieter Abbeel yöntemi Ekim 2023'te Sonsuza Yakın Bağlam için Bloklu Transformers ile Halka Dikkat başlıklı makalelerinde, daha önceki bloklar hâlinde paralel transformer yaklaşımını temel alarak tanıttılar ve dikkat hesaplamasında yaklaşık yöntemlere başvurmadan, bağlam pencerelerini milyonlarca belirteç uzunluğunda kullanarak eğitim ve çıkarım yaptıklarını bildirdiler.

Halka Dikkati Nasıl Çalışır#

Bir dizi, her bir GPU veya TPU için bir tane olmak üzere N bloğa bölünür. Her cihaz, tüm hesaplama boyunca sorgu bloğunu elinde tutar. N adımın her birinde sorgularıyla o anda tuttuğu anahtar-değer bloğu arasındaki dikkati hesaplar, bloğu sonraki cihaza gönderir ve önceki cihazdan başka bir blok alır:

N adımın ardından her sorgu bloğu her anahtar ve değer bloğuyla eşleşmiş olur. Hiçbir cihaz dikkat matrisinin tamamını veya anahtar ve değerlerin tümünü saklamaz; bu nedenle cihaz başına bellek kullanımı dizi uzunluğuna değil blok boyutuna bağlıdır. Bir blok üzerindeki dikkat hesaplaması o bloğu göndermekten uzun sürdüğü sürece aktarım hesaplamanın ardında gizlenir; bu da belirli donanım için gereken en küçük blok boyutunu belirler. Kısmi sonuçlar, tek bir GPU içindeki flash attention yönteminin kullandığı aynı devam eden softmax ile birleştirilir: her cihaz, üstel puanların ve ağırlıklı değerlerin devam eden toplamını tutar ve sonunda birini diğerine böler. Tek bir makinede dört cihazı simüle etmek sonucun tam olduğunu gösterir:

import torch

torch.manual_seed(0)
n_devices, block, dim = 4, 8, 16
q, k, v = torch.randn(3, n_devices * block, dim)
q_blocks, k_blocks, v_blocks = q.chunk(n_devices), k.chunk(n_devices), v.chunk(n_devices)

outputs = []
for i in range(n_devices):  # device i keeps query block i for the whole computation
    numerator, denominator = torch.zeros(block, dim), torch.zeros(block, 1)
    for step in range(n_devices):
        j = (i - step) % n_devices  # key-value block that has arrived from the previous device
        p = torch.exp(q_blocks[i] @ k_blocks[j].T / dim**0.5)
        numerator += p @ v_blocks[j]
        denominator += p.sum(dim=-1, keepdim=True)
    outputs.append(numerator / denominator)

full = torch.softmax(q @ k.T / dim**0.5, dim=-1) @ v
print(f"Largest score block per device: {block}x{block}, full matrix: {len(q)}x{len(k)}")
print("Matches full attention:", torch.allclose(torch.cat(outputs), full, atol=1e-5))

Matches full attention: True çıktısını verir: Her cihaz yalnızca 8 × 8 puanlık bir blok tuttu ve birlikte 32 × 32 matrisin tamamı üzerinde dikkati hesapladılar. Üretim çekirdekleri, sayıları kararlı tutmak için üstel alma işleminden önce devam eden bir maksimum değeri de çıkarır. Nedensel maske kullanıldığında, erken sorgu bloklarını tutan cihazlar anahtar bloklarının çoğunu atlar ve boşta kalır; Şeritli Dikkat, her cihaza dizi boyunca eşit dağıtılmış belirteçler vererek bunu düzeltir ve 256 bin belirteçte uçtan uca eğitim verimini 1,45 kata kadar artırır.

Uygulamalar#

Halka dikkati, tek bir aygıta sığmayacak kadar uzun bir girdi olduğunda önem kazanır; bunun en açık örneği videodur. Aynı Berkeley grubu bu yöntemi, bağlamı aşamalı olarak 4 binden 1 milyon metin ve video belirtecine çıkarılan, 7 milyar parametreli bir görüntü-dil modeli (VLM) ailesi olan Large World Model'i (LWM) eğitmek için kullandı. LWM her kareyi 256 belirteç olarak kodlar; bu nedenle saniyede bir kare örneklenen bir saatlik video 921.600 belirteç eder. Model, uzun süreli bir video anlama görevi olan bir saatlik YouTube derlemeleri hakkındaki soruları yanıtladı. Tüm kod tabanlarını veya belge koleksiyonlarını tek bir istemde okuyan büyük dil modellerini (LLM) eğitmek de aynı sınıra takılır. PyTorch'un context_parallel() API'si, kendi bağlamında ölçeklendirilmiş nokta çarpımı dikkatini halka dikkatiyle değiştirir; TorchTitan ekibi bu yöntemi 32 H100 GPU üzerinde 1 milyon belirteçlik dizilerle Llama 3 8B'yi eğitmek için kullandı.

Halka Dikkati ve Diğer Paralellik Türleri#

• Bağlam paralelliği: Tek bir diziyi aygıtlar arasında bölmeye yönelik genel strateji. Halka dikkati, PyTorch uygulamasının ardındaki algoritmadır. DeepSpeed Ulysses ise her aygıta dikkat başlıklarının bir alt kümesi için tüm belirteçleri veren, tüm aygıtlar arası değiş tokuşları kullanan bir alternatiftir.
• Tensör paralelliği ve iş hattı paralelliği: Modelin ağırlıklarını katmanların içinde veya katmanlar arasında böler; parametreler tek bir aygıtın belleğine sığmayacak kadar büyüdüğünde yarar sağlar.
• Flash attention: Tek aygıtlı bir çekirdek. Halka dikkati, her aygıtta yerel bloğu için bu türden bloklar hâlinde çalışan bir çekirdek kullanır.

Veri paralelliği, dizi yerine yığını böler. Ultralytics YOLO26, çoklu GPU eğitiminde bu modu kullanır: device=[0, 1] değerini iletmek PyTorch DistributedDataParallel'ı başlatır. Bu yöntem her GPU'da modelin tam bir kopyasını tutar, her yığın görüntüyü GPU'lar arasında eşit olarak dağıtır ve her adımdan sonra gradyanların ortalamasını alır.