Nowy natywny backend i integracja z bibliotekami
Dotychczas integracja bibliotek z układami AWS Trainium oraz Inferentia wymagała korzystania z pakietu optimum-neuron oraz PyTorch XLA. Podejście to wymuszało stosowanie leniwych tensorów oraz statycznych kształtów, co komplikowało proces tworzenia kodu. Nowe rozwiązanie integruje się bezpośrednio z PyTorch i udostępnia zasoby Trainium jako natywne urządzenia z obsługą trybu eager oraz funkcji torch.compile.
Dzięki temu biblioteki takie jak Transformers, Diffusers czy Accelerate współpracują z tymi procesorami bezpośrednio, bez konieczności stosowania dedykowanych klas pomocniczych czy zmodyfikowanych forków kodu modeli. Publiczna wersja beta pakietu zostanie udostępniona dwudziestego października dwa tysiące dwudziestego szóstego roku z możliwością instalacji za pomocą polecenia pip install torch-neuronx.
Generowanie tekstu i obsługa modeli w Transformers
W ramach biblioteki Transformers generowanie tekstu jest teraz obsługiwane natywnie dla szerokiej gamy modeli. Rozwiązanie to zapewnia wsparcie dla osiemdziesięciu dziewięciu procent ze stu najpopularniejszych architektur służących do generowania tekstu, w porównaniu do czterdziestu sześciu procent w starszej wersji optimum-neuron.
Uruchomienie modelu na nowym urządzeniu wymaga przekazania argumentu device_map=neuron. Ze względu na specyfikację procesorów zaleca się stosowanie statycznej pamięci podręcznej w procesie generowania, co zapobiega częstym rekompilacjam wykresów obliczeniowych.
Trenowanie modeli za pomocą TRL
Proces trenowania na procesorach Trainium korzysta obecnie z tych samych standardowych bibliotek, co w przypadku innych akceleratorów. Skrypty dostrajania oparte na SFTTrainer działają bezpośrednio po wprowadzeniu odpowiednich ustawień konfiguracyjnych.
Zarządzanie zrównoległowaniem odbywa się natywnie w bibliotece Transformers. Wystarczy przekazać strukturę DistributedConfig zawierającą parametry zrównoleglenia tensorowego oraz mechanizmu FSDP, aby waga modelu została podzielona podczas wczytywania. Podczas uruchamiania skryptów szkoleniowych konieczne jest wyłączenie automatycznych dynamicznych kształtów oraz asynchronicznego wczytywania wag za pomocą zmiennej środowiskowej HF_DEACTIVATE_ASYNC_LOAD=1.
Wsparcie dla biblioteki Diffusers i optymalizacja jądra
Wiele potoków z biblioteki Diffusers wykazuje pełną kompatybilność z nowym backendem w trybie eager oraz po kompilacji. Modele takie jak FLUX.1-dev, FLUX.2 czy Qwen-Image obsługują zrównoleglenie tensorowe, co pozwala w pełni wykorzystać moc obliczeniową urządzeń Trainium.
Użytkownicy mogą również korzystać z dedykowanych jąder optymalizacyjnych tworzonych przy użyciu interfejsu Neuron Kernel Interface. Pozwala to na pisanie własnych operacji niestandardowych i osiąganie znacznego przyspieszenia działania modeli wizualnych bezpośrednio na platformie Hugging Face Hub.
