Przyspieszenie wnioskowania AI w tworzeniu treści 3D na platformie Roblox

• Firma Roblox wdrożyła CUDA Graphs i buforowanie KV w celu przyspieszenia generowania siatki 3D, co zapewnia bardziej responsywne iteracje.
• W momencie premiery model Cube 3D mógł generować tokeny w 7,8 milisekundy (w porównaniu z 60,5 milisekundy) oraz pełne obiekty w 4 sekundy (w porównaniu z 31 sekund). 

Na początku tego roku firma Roblox udostępniła pierwszą wersję naszego modelu podstawowego Cube 3D. Dzięki Cube 3D twórcy mogą generować modele 3D i środowiska bezpośrednio na podstawie poleceń tekstowych. Od samego początku priorytetowo traktowaliśmy optymalizację opóźnień, zdając sobie sprawę, że długi czas generowania zakłóca proces, który z natury jest iteracyjny. Jeszcze przed premierą Cube 3D w marcu udało nam się przyspieszyć etap wnioskowania 7,8-krotnie, zapewniając większą responsywność zarówno dla programistów, jak i użytkowników. 

Od momentu uruchomienia wygenerowano ponad 578 000 obiektów w ramach kilku znaczących doświadczeń. Deweloperzy wyrazili również zainteresowanie umożliwieniem użytkownikom generowania obiektów 3D w ramach doświadczeń za pomocą poleceń tekstowych, takich jak „koty”, „burgery” itp. W szczególności Mic Up, popularna gra towarzyska wykorzystująca czat głosowy, wykorzystała Cube 3D, aby zapewnić graczom zabawny i interaktywny sposób generowania obiektów. W tej implementacji gracze mogą otworzyć menu po lewej stronie z dodatkowymi funkcjami, w tym ikoną AI. Po kliknięciu tej ikony gracze mogą wprowadzić polecenie tekstowe, aby wygenerować obiekt 3D. Dla użytkowników dłuższy czas generowania powoduje utrudnienia, pozbawiając ich magii obserwowania, jak ich pomysły przekształcają się w 3D w czasie rzeczywistym. 

Chcieliśmy przekształcić proces generowania 3D z interakcji typu „zatrzymaj się i czekaj” w coś, co wydaje się responsywne i naturalne, umożliwiając szybkie eksperymentowanie. Możliwość szybkiego dodawania obiektów do sceny ma kluczowe znaczenie dla programistów. Aby przyspieszyć działanie Cube 3D, najpierw przeanalizowaliśmy potok wnioskowania w celu zidentyfikowania wąskich gardeł wydajności. Pomimo użycia wydajnych procesorów graficznych (GPU) stwierdziliśmy znaczny czas bezczynności między operacjami.

Nowoczesne frameworki głębokiego uczenia się polegają na procesorze CPU w zakresie planowania i uruchamiania operacji (lub jąder) na procesorze GPU. Procesor CPU przygotowuje każdą operację, wysyła ją do procesora GPU i oczekuje na potwierdzenie przed przygotowaniem kolejnej operacji. To oczekiwanie powoduje wąskie gardło w planowaniu, w wyniku którego procesor GPU może pozostawać bezczynny, podczas gdy procesor CPU przygotowuje kolejną partię zadań. W idealnym przypadku chcielibyśmy, aby procesor CPU wyprzedzał procesor GPU, przygotowując i umieszczając operacje w kolejce, tak aby procesor GPU zawsze miał zadania do wykonania.

Jest to szczególnie problematyczne w przypadku dekoderów autoregresyjnych w modelach typu transformator, takich jak Cube 3D, które muszą przetwarzać dane wejściowe i generować tokeny sekwencyjnie. Modele te wymagają tysięcy pojedynczych operacji dla jednego generowania, a obciążenie obliczeniowe kumuluje się z każdym krokiem w sekwencji.

„Chcieliśmy stworzyć coś, co umożliwi czterowymiarową interakcję” – powiedział wiceprezes ds. inżynierii Anupam Singh, wyjaśniając, dlaczego firma Roblox wybrała podejście autoregresyjne. „Nie chcemy tylko zbudować samochodu; chcemy również móc otworzyć drzwi samochodu i wsiąść do niego”.

Każda operacja wymagała:

• Czas procesora potrzebny na przygotowanie każdego jądra
• Obciążenie związane z uruchomieniem jądra
• Czas wykonania na GPU (rzeczywiste obliczenia)
• Obciążenie związane z synchronizacją podczas sprawdzania zakończenia

W przypadku małych operacji, które szybko się wykonują na GPU, to obciążenie może zdominować czas wnioskowania. GPU może aktywnie obliczać tylko przez niewielką część całkowitego czasu wnioskowania.

Aby rozwiązać ten problem, wykorzystaliśmy CUDA Graphs — funkcję, która pozwala na rejestrowanie i odtwarzanie sekwencji operacji GPU bez interwencji procesora. Komponent dekodera autoregresywnego w architekturze Cube 3D przetwarza podpowiedzi tekstowe i generuje tokeny kształtów za pomocą wektora o stałej długości. 

Chociaż pod względem funkcjonalnym przypomina to tradycyjny duży model językowy (LLM), nasza architektura dekodera z podwójnym strumieniem ma jedną istotną różnicę — wykorzystuje dwa równoległe strumienie uwagi. Jeden strumień jest przeznaczony dla tokenów warunkowych, a drugi dla tokenów kształtów. Gotowe silniki wnioskowania LLM nie odpowiadały naszym potrzebom i potrzebowaliśmy niestandardowej implementacji dostosowanej do naszej konkretnej architektury.

Warto wyobrazić sobie wykresy CUDA jako nagrywanie makra dla procesora graficznego. Zamiast wydawania każdego polecenia osobno przez procesor centralny, rejestruje on całą sekwencję operacji procesora graficznego (wykres) i uruchamia cały wykres za pomocą jednej instrukcji procesora centralnego. Takie podejście znacznie zmniejsza obciążenie związane z uruchamianiem jądra, eliminując konieczność indywidualnego planowania każdej operacji przez procesor centralny podczas wnioskowania. Po uruchomieniu wykresu procesor graficzny wykonuje całą sekwencję autonomicznie, bez oczekiwania na dalsze instrukcje.

Wykresy CUDA mają pewne ograniczenia. Ponieważ struktura wykresu musi być ustalona z góry, wymagają one stałego rozmiaru partii i wymiarów danych wejściowych. Oznacza to, że dla każdego rozmiaru partii lub kształtu danych wejściowych trzeba tworzyć osobne wykresy. W naszym przypadku z Cube 3D to ograniczenie było do zaakceptowania, bo mogliśmy ujednolicić proces wnioskowania wokół typowych wymiarów danych wejściowych.

Musieliśmy dostosować nasze podejście, aby zaimplementować wykresy CUDA w naszym modelu Cube 3D. W tradycyjnych modelach LLM operacje uwagi są zawsze wykonywane przy tej samej długości sekwencji, co zapewnia statyczny kształt do pracy. Jednak w naszej niestandardowej architekturze dwustrumieniowej niektóre warstwy uwagi działają wyłącznie na długości sekwencji, podczas gdy inne działają na kombinacji długości sekwencji i warunku.

Pomimo tych wyzwań, po wdrożeniu CUDA Graphs zaobserwowaliśmy niezwykłe wyniki. Używamy czasu na token wyjściowy (TPOT) do pomiaru czasu generowania każdego tokenu podczas wnioskowania. Po wdrożeniu CUDA Graphs nasz TPOT poprawił się z 60,5 milisekund do 20,5 milisekund, co stanowi 2,9-krotną poprawę. Ogólny czas generowania spadł o 66%, z 31 sekund do 10,5 sekundy.

Aby jeszcze bardziej skrócić opóźnienia, wdrożyliśmy buforowanie KV, standardową praktykę w wnioskowaniu LLM, która okazała się bardzo skuteczna w całej branży. 

W modelach opartych na transformatorach, takich jak Cube 3D, generowanie każdego tokenu wymaga obliczenia macierzy klucza (K) i wartości (V) na podstawie wszystkich wcześniej wygenerowanych tokenów. W miarę wydłużania się sekwencji ponowne obliczanie tych macierzy dla każdego tokenu staje się coraz mniej wydajne.

Buforowanie KV rozwiązuje ten problem poprzez:

1. Przechowywanie macierzy K i V dla wszystkich wcześniej wygenerowanych tokenów
2. Obliczanie macierzy K i V tylko dla nowych tokenów
3. Dodawanie tych nowych macierzy do wartości w pamięci podręcznej

Takie podejście eliminuje zbędne obliczenia, zmniejszając nakład pracy wymagany dla każdego nowego tokenu. Ma to szczególne znaczenie w miarę wydłużania się generowanej sekwencji.

Nasze podejście do integracji buforowania KV z implementacją CUDA Graph było podobne do tradycyjnego wnioskowania LLM. Dodanie buforowania KV zmniejszyło nasz TPOT do zaledwie 7,8 milisekundy. Ogólny czas generowania zmniejszył się o 87%, z pierwotnych 31 sekund do zaledwie 4 sekund. To znaczące skrócenie czasu sprawia, że narzędzie to jest znacznie bardziej efektywne dla twórców, którzy z niego korzystają.

Badamy techniki, które jeszcze bardziej zmniejszają opóźnienia i poprawiają komfort użytkowania, w tym zoptymalizowane jądra, kwantyzację modeli dla jeszcze szybszego wnioskowania, optymalizacje specyficzne dla sprzętu oraz równoległe generowanie tokenów.

Praca ta staje się bardziej złożona, gdy rozszerzamy ją na generowanie i rozumienie pełnych scen, gdzie wiele elementów 3D musi współpracować ze sobą w kontekście układu. Chcemy również, aby tworzone przez nas obiekty i światy 3D były w pełni funkcjonalne, tak aby drzwi się otwierały i zamykały, koła się obracały itp. Aby to osiągnąć, potrzebujemy szybkiego generowania i iteracji, aby skalować całe sceny, w pełni funkcjonalne obiekty i awatary. Cieszymy się, że możemy podzielić się dalszymi ulepszeniami i nowymi funkcjami w miarę rozbudowywania naszego modelu podstawowego Cube 3D — oraz że możemy zobaczyć wciągające światy, które tworzy nasza społeczność twórców.