Portowanie modelu Google Gemma-4 31B na AWS Inferentia2 inf2.24xlarge dało idealne dopasowanie token po tokenie względem referencji CPU — jednak każde wygenerowane zdanie było bełkotem. Luka między „dopasowaniem” a „działaniem” służy teraz jako przestroga dla każdego, kto próbuje upchnąć ogromne modele LLM na dedykowanych chipach inferencyjnych Amazon.

Dlaczego dopasowanie token po tokenie to za mało

Deweloper porównał każdy token wyjściowy z urządzenia Inferentia z tokenem wygenerowanym przez model na procesorze CPU. Strumienie były identyczne, więc sprzęt wydawał się dokładnie odtworzyć implementację referencyjną. W rzeczywistości oba strumienie przekazywały błędnie sformatowany prompt do modelu pozbawionego szablonu czatu (chat template) i wyposażonego w niewłaściwe znaczniki tur (turn markers). Brakujący szablon wprowadził model w nieskończoną pętlę, wyrzucając bezsensowne dane. Sprzęt wykonał swoje zadanie — odtworzył błąd, który istniał w kodzie referencyjnym.

Lekcja jest prosta: SEQ_MATCH (sekwencyjna równość tokenów) nie oznacza poprawności. Jeśli implementacja referencyjna jest wadliwa, wierna replika sprzętowa dziedziczy tę samą usterkę. Walidacja musi wykraczać poza parzystość na poziomie tokenów; wymaga ona testów funkcjonalnych end-to-end z poprawnie sformatowanymi danymi wejściowymi.

Bufory podszywające się pod parametry

Podczas fazy ładowania loader modelu pominął komponent o nazwie layer_scalar. Kod zarejestrował ten obiekt jako bufor (buffer), a nie jako parametr (parameter) w definicji modelu PyTorch. Bufory to statyczne tensory, których proces uczenia nie aktualizuje, a wiele loaderów ignoruje je podczas konwersji na formaty kompatybilne z Neuron. Pominięcie go spowodowało, że czynniki skalujące dla kilku warstw pozostały na wartościach domyślnych, co zniekształciło obliczenia w całej sieci. Nie zgłoszono żadnego błędu; model został skompilowany, a potok inferencyjny uruchomiony, ale wyniki numeryczne były błędne.

Dla każdego, kto przenosi duże modele na Inferentia, zaleca się audyt każdego tensora niebędącego parametrem. Nawet jeśli tensor nie ma być uczony, może być niezbędny do poprawnego obliczenia przejścia w przód (forward-pass). Ręczna weryfikacja uwzględnienia buforów może zapobiec cichym błędom skalowania, które trudno jest inaczej zdiagnozować.

Zmienność instancji typu Spot i 39-minutowa kompilacja

Uruchamianie modelu o 31 miliardach parametrów na instancji typu spot wydaje się tanie, ale oszczędności wiążą się z nieprzewidywalnymi zdarzeniami odzyskiwania zasobów. Czas kompilacji dewelopera — około 39 minut na przetłumaczenie modelu na kod kompatybilny z Neuron — przepadł, gdy AWS odzyskał instancję. Aby przetrwać przerwy, zbudowano trójstopniową siatkę bezpieczeństwa:

  • ModelBuilder utrzymywał zużycie pamięci w granicach limitu 384 GB hosta, unikając awarii wymuszających restart.
  • Natychmiastowe lustrzane odbicie w S3 zarówno surowych plików wag, jak i skompilowanych plików „neffs” (pliki wykonywalne Neuron) pozwalało nowej instancji kontynuować pracę dokładnie tam, gdzie poprzednia przerwała.
  • Wieloregionowy poller skanował regiony AWS w poszukiwaniu dostępnej pojemności spot i uruchamiał nową instancję, gdy tylko się pojawiła.

Te kroki zmieniły kruchą, jednopunktową kompilację w odporny potok, który radzi sobie z wahaniami na rynkach spot.

Pułapki shardingu przy mieszanych układach attention

Gemma-4 31B używa dwóch konfiguracji attention. Niektóre warstwy wykorzystują cztery głowice klucz-wartość (KV heads), inne inną liczbę. Równomierne rozdzielenie modelu na osiem równoległych rang (ranks) kończy się niepowodzeniem, gdy liczba głowic KV warstwy nie jest podzielna przez osiem. Próba shardingu warstwy z 4 głowicami na osiem rang wymusiłaby na każdej randzie obsługę połowy głowicy — co jest matematycznie niemożliwe i powoduje niedopasowania kształtów (shape mismatches) oraz błędy czasu wykonania (runtime errors).

Rozwiązaniem było replikowanie warstw shardowanych globalnie (tych z kompatybilną liczbą głowic) we wszystkich rangach i shardowanie tylko warstw „sliding”, których liczba głowic pozwalała na równy podział. Ta hybrydowa strategia zachowała wydajność tensor-parallel, unikając jednocześnie nielegalnego dzielenia głowic KV, co wyeliminowało błędy tensor-parallelization, które nękały wcześniejsze próby.

Podsumowanie

Przenoszenie gigantycznego modelu LLM na Inferentia to coś więcej niż tylko ćwiczenie typu „skompiluj i uruchom”. Wymaga ono rygorystycznych testów funkcjonalnych wykraczających poza równość tokenów, skrupulatnej weryfikacji, czy każdy tensor — parametr lub bufor — jest poprawnie obsługiwany, oraz strategii wdrożenia uwzględniającej odzyskiwanie instancji spot. Na koniec, sharding musi respektować wewnętrzną geometrię attention modelu; w przeciwnym razie paralelizacja, która obiecuje szybkość, stanie się źródłem cichych awarii.