i14 ve AWS, Laplace 3B Doğrusal Dikkat Modelini Trainium2'de Eğitti

i14 ve AWS, transformer mimarilerindeki anahtar-değer önbellek büyüme sorununa çözüm sunan Laplace 3B doğrusal dikkat modelini AWS Trainium2 üzerinde, yerel PyTorch desteğiyle eğitti ve çalıştırdı.

Çoğu büyük dil modeli, bağlam uzunluğu arttıkça anahtar-değer (KV) önbelleğinin büyüdüğü transformer mimarilerini ve softmax dikkat mekanizmalarını kullanmaktadır. Bu sorunu aşmak için i14 tarafından tasarlanan Laplace 3B, yinelemeli durumu dizi başına sabit 29,5 MB kalan bir doğrusal dikkat alternatifi olarak geliştirilmiştir. AWS ile yürütülen ortak çalışma sonucunda bu özel mimari, TorchNeuron Native aracılığıyla sunulan yerel PyTorch desteğiyle AWS Trainium2 donanımına uyarlanmış ve Gated Linear Attention tasarımı ile Muon optimize edicisi korunmuştur.

Tek bir trn2.48xlarge örneği üzerinde gerçekleştirilen testlerde Laplace 3B, saniyede yaklaşık 120.000 token hızında eğitim gerçekleştirmiş ve 1.280 eş zamanlı dizi boyunca saniyede 11.544 çıkış tokeni sunmuştur. Bu entegrasyon sürecinde GLA dikkat çekirdeğinin standart PyTorch operatörlerine uyarlanması, ağ trafiğini azaltmak için Muon iletişiminin all_to_all kolektifleriyle yeniden tasarlanması ve derleyici sınırlarına uyum sağlanması için geriye doğru geçişin statik olarak açılması gerektiği belirtilmiştir.