DeepSeek-V3 od zera : budowa architektury w PyTorch

Serwis PyImageSearch publikuje tutorial poświęcony budowie architektury DeepSeek-V3 przy użyciu PyTorch, biblioteki do głębokiego uczenia. To piąta część z serii sześciu lekcji. Ten tutorial, autorstwa Puneeta Mangli, składa w całość elementy przedstawione w poprzednich lekcjach, tworząc pełny model językowy, który przekształca tokeny – jednostki tekstu przetwarzane przez model – w logity, czyli wyniki przypisywane każdemu możliwemu tokenowi.

Poprzednie lekcje obejmowały konfigurację modelu, Rotary Position Embeddings (RoPE), które kodują pozycję tokenów, Multihead Latent Attention (MLA), mechanizm uwagi, oraz Mixture of Experts (MoE), czyli mieszaninę ekspertów, która rozdziela obliczenia między kilka wyspecjalizowanych podsieci. Multi-Token Prediction (MTP) pozwala natomiast przewidywać kilka tokenów jednocześnie. Według tutorialu poprawia to efektywność treningu i może przyspieszyć wnioskowanie dzięki dekodowaniu spekulatywnemu.

Każdy blok transformera łączy MLA i MoE według schematu zwanego « pre-norm », w którym normalizacja następuje przed każdą podwarstwą. Połączenia rezydualne przesyłają sygnały bezpośrednio z jednej warstwy do drugiej, co, jak wyjaśnia autor, pomaga rozwiązać problem zanikania gradientów w głębokich sieciach. Pełny model składa się kolejno z osadzeń tokenów i pozycji, warstwy dropout, bloków, końcowej normalizacji, głowy predykcyjnej słownika oraz opcjonalnie głów MTP.

Współdzielenie wag i predykcja wielu tokenów

Tutorial przedstawia również współdzielenie wag między macierzą osadzania tokenów a głową predykcyjną. Dla słownika o rozmiarze 50 259 tokenów i wymiarze osadzania 256, ta technika pozwala zaoszczędzić około 12,9 miliona parametrów. Zdaniem autora często poprawia to wydajność, ponieważ przestrzenie wejściowe i wyjściowe pozostają wyrównane.

Podczas treningu wyjścia głów MTP są obliczane dla kilku przyszłych kroków, a następnie ich straty są łączone z główną stratą oraz z pomocniczą stratą MoE, która ma na celu zrównoważenie obciążenia między ekspertami. W przedstawionym kodzie średnia strata MTP jest ważona współczynnikiem 0,3 w stosunku do straty głównej. Głowy MTP nie są używane podczas wnioskowania, czyli generowania tekstu.

Rozmiar modelu i zużycie pamięci

Według autora model zawiera około 30 do 35 milionów parametrów rozłożonych na sześć warstw, czyli 2 do 3 milionów na warstwę, do których dodaje się 2 do 3 milionów parametrów na każdą głowę MTP. Przy precyzji FP32, na 32 bitach, same parametry zajmują 120 do 140 megabajtów. Trening, wraz z gradientami i stanami optymalizatora AdamW, wymaga około 500 do 600 MB, do których należy dodać aktywacje. Dla partii 4 przykładów po 1024 tokeny, aktywacje zajmują 200 do 300 MB, co daje blisko 800 MB łącznie – ilość, którą tutorial uznaje za zgodną z możliwościami współczesnych procesorów graficznych (GPU).

Podczas wnioskowania precyzja FP16 zmniejsza zużycie pamięci o około połowę. Implementacja nie korzysta z pamięci podręcznej klucz-wartość, więc kontekst jest przeliczany na każdym kroku generowania. Kolejna lekcja z serii będzie dotyczyć trenowania tego modelu na rzeczywistych danych.