ZAYA1: przełom w trenowaniu modeli AI na GPU AMD

Współpraca giganta technologicznego AMD z firmami Zyphra i IBM zaowocowała powstaniem ZAYA1 – pierwszego dużego modelu opartego na architekturze Mixture-of-Experts (MoE), który został wytrenowany wyłącznie na procesorach graficznych i systemach sieciowych AMD. To przełomowe osiągnięcie podważa dotychczasową dominację jednego dostawcy w sektorze zaawansowanego sprzętu do AI i otwiera nowe perspektywy dla rozwoju sztucznej inteligencji.

Przez ostatni rok konsorcjum testowało możliwości GPU AMD oraz ich platformy w kontekście trenowania modeli AI na dużą skalę. Wynikiem tych intensywnych prac jest ZAYA1, która zdaniem twórców dorównuje, a w niektórych aspektach przewyższa, ugruntowane otwarte modele w zakresie rozumowania, matematyki i kodowania. Model był szkolony na chipach AMD Instinct MI300X, wykorzystując technologię sieciową Pensando oraz oprogramowanie ROCm, wszystko to w ramach infrastruktury IBM Cloud.

AMD w roli alternatywy

Co istotne, zastosowana konfiguracja jest zaskakująco konwencjonalna. Zamiast eksperymentalnych rozwiązań Zyphra zbudowała system przypominający standardowe klastry korporacyjne, jednak bez komponentów dominującego konkurenta. Dla przedsiębiorstw zmagających się z ograniczeniami dostaw lub rosnącymi cenami procesorów graficznych, ZAYA1 stanowi rzadką drugą opcję, która nie wymaga kompromisów w kwestii wydajności.

Kluczowym elementem sukcesu było efektywne wykorzystanie pamięci o dużej przepustowości (HBM) dostępnej w układach MI300X. Każdy procesor graficzny MI300X oferuje 192 GB HBM, co daje inżynierom znaczną swobodę, umożliwiając wykonywanie wczesnych etapów treningu bez natychmiastowego uciekania się do złożonego paralelizmu. Upraszcza to projekty, które w innym przypadku byłyby kruche i czasochłonne w optymalizacji. Każdy węzeł klastra został zbudowany z ośmiu procesorów graficznych MI300X połączonych przez InfinityFabric, z oddzielnym układem sieciowym Pollara do obsługi odczytu danych i punktów kontrolnych. Prosta i efektywna architektura minimalizuje koszty przełączników i utrzymuje stabilne czasy iteracji.

ZAYA1: efektywność i architektura

ZAYA1-base aktywuje 760 milionów parametrów z puli 8,3 miliarda, trenując na 12 bilionach tokenów w trzech etapach. Architektura modelu opiera się na skompresowanej uwadze (compressed attention), zaawansowanym systemie routingu do kierowania tokenów do odpowiednich „ekspertów” oraz delikatniejszym skalowaniu resztkowym, zapewniającym stabilność głębszych warstw. Model wykorzystuje mieszankę optymalizatorów Muon i AdamW. Aby zoptymalizować Muon pod kątem sprzętu AMD, Zyphra scaliła jądra i zredukowała zbędny ruch pamięci, co zapobiegło dominacji optymalizatora w każdej iteracji. Rozmiary partii treningowych były zwiększane stopniowo, co wymagało efektywnych potoków danych.

To wszystko doprowadziło do stworzenia modelu AI, który mimo mniejszych rozmiarów jest w stanie konkurować z większymi odpowiednikami, takimi jak Qwen3-4B, Gemma3-12B, Llama-3-8B i OLMoE. Jedną z zalet struktury MoE jest to, że w danym momencie działa tylko niewielka część modelu, co pozwala na efektywne zarządzanie pamięcią podczas wnioskowania (inference) i redukuje koszty obsługi.

Wyzwania i optymalizacja oprogramowania ROCm

Przeniesienie dojrzałego środowiska pracy opartego na rozwiązaniach konkurenta na platformę ROCm nie było pozbawione wyzwań. Zespół Zyphra poświęcił czas na analizę zachowania sprzętu AMD i odpowiednie przekształcanie wymiarów modelu, wzorców GEMM oraz innych zmiennych, aby dopasować je do preferowanych zakresów obliczeniowych MI300X. Zoptymalizowano wykorzystanie InfinityFabric, aby wszystkie osiem procesorów graficznych w węźle uczestniczyło w operacjach zbiorowych, oraz Pollara, która osiąga szczytową przepustowość przy większych wiadomościach. Intensywne treningi z długimi kontekstami (od 4k do 32k tokenów) wykorzystywały mechanizm ring attention dla dzielonych sekwencji i tree attention podczas dekodowania, aby uniknąć wąskich gardeł.

Również kwestie pamięci masowej zostały rozwiązane praktycznie. Mniejsze modele obciążają operacje IOPS, podczas gdy większe wymagają stałej przepustowości. Zyphra zgrupowała fragmenty zbiorów danych, aby zredukować rozproszone odczyty, i zwiększyła pamięć podręczną stron per-węzeł, aby przyspieszyć odzyskiwanie punktów kontrolnych – kluczowe podczas długich sesji treningowych, gdzie cofanie się jest nieuniknione.

Stabilność klastra i zarządzanie błędami

Długotrwałe zadania treningowe, trwające tygodniami, rzadko przebiegają bez zakłóceń. Usługa Aegis firmy Zyphra monitoruje logi i metryki systemowe, identyfikuje awarie, takie jak usterki NIC czy błędy ECC, i automatycznie podejmuje proste działania naprawcze. Zespół zwiększył również limity czasu RCCL, aby krótkie przerwy w sieci nie przerywały całych zadań. Punktowanie kontrolne (checkpointing) jest rozproszone na wszystkie procesory graficzne, zamiast przechodzić przez jedno wąskie gardło.

Wnioski dla zamówień publicznych AI

Raport ZAYA1 wyraźnie wskazuje na istnienie ugruntowanych odpowiedników ekosystemu AMD dla rozwiązań konkurencji: NVLINK kontra InfinityFabric, NCCL kontra RCCL, cuBLASLt kontra hipBLASLt. Autorzy projektu twierdzą, że stos technologiczny AMD jest wystarczająco dojrzały do poważnego rozwoju modeli na dużą skalę. Nie oznacza to jednak, że przedsiębiorstwa powinny natychmiast usuwać istniejące klastry oparte na konkurencyjnych rozwiązaniach. Bardziej realistyczne jest utrzymanie tych systemów do zadań produkcyjnych, jednocześnie wykorzystując AMD do etapów, które czerpią korzyści z pojemności pamięci procesorów graficznych MI300X i otwartości ROCm. Takie podejście rozkłada ryzyko dostawców i zwiększa ogólną zdolność treningową bez większych zakłóceń.

Praktyczne wnioski z prac Zyphry, AMD i IBM to: traktowanie kształtu modelu jako zmiennego, a nie stałego; projektowanie sieci wokół kolektywnych operacji, które będą faktycznie używane w treningu; budowanie tolerancji na błędy, która chroni godziny pracy procesorów graficznych, a nie tylko rejestruje awarie; oraz modernizacja punktów kontrolnych, aby nie zakłócały rytmu treningu. Dla organizacji pragnących zwiększyć swoje możliwości w zakresie AI bez polegania wyłącznie na jednym dostawcy, model ZAYA1 i podejście partnerów stanowią cenny schemat działania.