Biblioteka równoległa modeli Amazon SageMaker przyspiesza teraz obciążenia PyTorch FSDP nawet o 20% | Usługi internetowe Amazona

Biblioteka równoległa modeli Amazon SageMaker przyspiesza teraz obciążenia PyTorch FSDP nawet o 20% | Usługi internetowe Amazona

Węzeł źródłowy: 3033172

Popularność szkoleń w oparciu o modele dużych języków (LLM) wzrosła w ciągu ostatniego roku wraz z wypuszczeniem kilku popularnych modeli, takich jak Llama 2, Falcon i Mistral. Klienci przeprowadzają obecnie wstępne szkolenie i dostrajanie LLM w zakresie od 1 miliarda do ponad 175 miliardów parametrów, aby zoptymalizować wydajność modelu dla aplikacji w różnych branżach, od opieki zdrowotnej po finanse i marketing.

Szkolenie wydajnych modeli na taką skalę może być wyzwaniem. Bardzo dokładne LLM mogą wymagać terabajtów danych szkoleniowych i tysięcy, a nawet milionów godzin czasu obliczeniowego akceleratora, aby osiągnąć docelową dokładność. Aby ukończyć szkolenia i wprowadzić produkty na rynek w odpowiednim czasie, klienci polegają na technikach równoległości w celu rozłożenia tego ogromnego obciążenia na nawet tysiące urządzeń akceleracyjnych. Jednakże te techniki równoległości mogą być trudne w użyciu: różne techniki i biblioteki są kompatybilne tylko z określonymi obciążeniami lub ograniczone do określonych architektur modeli, wydajność uczenia może być bardzo wrażliwa na niejasne konfiguracje, a stan wiedzy szybko się rozwija. W rezultacie praktycy uczenia maszynowego muszą spędzić tygodnie na przygotowaniach do skalowania obciążeń LLM do dużych klastrów procesorów graficznych.

W tym poście podkreślamy nowe funkcje programu Amazon Sage Maker bibliotekę modeli równoległych (SMP), która upraszcza proces uczenia dużych modeli i pomaga szybciej szkolić LLM. W szczególności omawiamy nowe, uproszczone środowisko użytkownika biblioteki SMP, które opiera się na interfejsach API PyTorch Fully Sharded Data Parallel (FSDP) typu open source, rozszerzoną funkcjonalność równoległego tensora, która umożliwia trenowanie modeli z setkami miliardów parametrów, oraz optymalizacje wydajności, które skracają czas szkolenia modelu i koszt nawet o 20%.

Aby dowiedzieć się więcej na temat biblioteki równoległej modeli SageMaker, zobacz Biblioteka równoległości modeli SageMaker v2 dokumentacja. Możesz także zapoznać się z naszym przykładowe zeszyty zacząć.

Nowe funkcje upraszczające i przyspieszające uczenie dużych modeli

W tym poście omówiono najnowsze funkcje zawarte w wersji 2.0 biblioteki równoległej modelu SageMaker. Funkcje te poprawiają użyteczność biblioteki, rozszerzają funkcjonalność i przyspieszają szkolenie. W poniższych sekcjach podsumowujemy nowe funkcje i omawiamy, w jaki sposób można wykorzystać bibliotekę do przyspieszenia uczenia dużych modeli.

Dopasowanie SMP do PyTorch typu open source

Od chwili premiery w 2020 r. rozwiązanie SMP umożliwiło wysokowydajne szkolenia na dużą skalę na instancjach obliczeniowych SageMaker. W najnowszej wersji głównej SMP biblioteka upraszcza obsługę użytkownika, dopasowując swoje interfejsy API do oprogramowania PyTorch typu open source.

Oferta PyTorch W pełni pofragmentowana równoległość danych (FSDP) jako główna metoda obsługi dużych obciążeń szkoleniowych na wielu urządzeniach obliczeniowych. Jak pokazano w poniższym fragmencie kodu, zaktualizowane interfejsy API firmy SMP dla technik takich jak równoległość danych podzielonych na fragmenty odzwierciedlają interfejsy PyTorch. Możesz po prostu biegać import torch.sagemaker i użyj go zamiast torch.

## training_script.py
import torch.sagemaker as tsm
tsm.init()

# Set up a PyTorch model
model = ...

# Wrap the PyTorch model using the PyTorch FSDP module
model = FSDP(
model,
...
)

optimizer = ...
...

Dzięki tym aktualizacjom interfejsów API SMP możesz teraz wykorzystać zalety wydajności SageMaker i biblioteki SMP bez konieczności zmiany istniejących skryptów szkoleniowych PyTorch FSDP. Ten paradygmat pozwala również na użycie tego samego kodu bazowego podczas szkoleń w siedzibie firmy, co w SageMaker, upraszczając doświadczenie użytkownika dla klientów, którzy szkolą się w wielu środowiskach.

Aby uzyskać więcej informacji na temat włączania protokołu SMP w istniejących skryptach szkoleniowych PyTorch FSDP, zobacz Rozpocznij pracę z SMP.

Integracja równoległości tensorów w celu umożliwienia uczenia na ogromnych klastrach

To wydanie SMP rozszerza także możliwości PyTorch FSDP o techniki równoległości tensorowej. Jednym z problemów związanych z używaniem samej równoległości danych podzielonych na fragmenty jest to, że w miarę zwiększania rozmiaru klastra mogą wystąpić problemy z konwergencją. Dzieje się tak, ponieważ parametry fragmentowania, gradienty i stan optymalizatora w szeregach równoległych danych również zwiększają globalny rozmiar wsadu; w przypadku dużych klastrów ten globalny rozmiar partii można przesunąć poza próg, poniżej którego model byłby zbieżny. Należy zastosować dodatkową technikę równoległości, która nie wymaga zwiększania globalnego rozmiaru partii podczas skalowania klastra.

Aby złagodzić ten problem, SMP v2.0 wprowadza możliwość tworzenia równoległości danych podzielonych na fragmenty z równoległością tensora. Równoległość tensorów umożliwia zwiększenie rozmiaru klastra bez zmiany globalnego rozmiaru wsadu lub wpływu na zbieżność modelu. Dzięki tej funkcji można bezpiecznie zwiększyć przepustowość uczenia, udostępniając klastry zawierające co najmniej 256 węzłów.

Obecnie równoległość tensorów z PyTorch FSDP jest dostępna tylko w SMP v2. SMP v2 umożliwia włączenie tej techniki za pomocą kilku linijek kodu i odblokowanie stabilnego uczenia nawet w dużych klastrach. SMP v2 integruje się z Silnik transformatorowy za implementację równoległości tensorów i czyni go kompatybilnym z interfejsami API PyTorch FSDP. Możesz włączyć jednocześnie równoległość tensora PyTorch FSDP i SMP bez wprowadzania jakichkolwiek zmian w modelu PyTorch lub konfiguracji PyTorch FSDP. Poniższe fragmenty kodu pokazują, jak skonfigurować słownik konfiguracyjny SMP w formacie JSON i dodać moduł inicjujący SMP torch.sagemaker.init(), który akceptuje słownik konfiguracji w zapleczu po uruchomieniu zadania szkoleniowego, do skryptu szkoleniowego.

Konfiguracja SMP wygląda następująco:

{
"tensor_parallel_degree": 8,
"tensor_parallel_seed": 0
}

W skrypcie szkoleniowym użyj następującego kodu:

import torch.sagemaker as tsm
tsm.init()

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model)

Aby dowiedzieć się więcej na temat używania równoległości tensorów w SMP, zobacz równoległość tensorowa część naszej dokumentacji.

Skorzystaj z zaawansowanych funkcji, aby przyspieszyć uczenie modelu nawet o 20%

Oprócz umożliwienia rozproszonego szkolenia w klastrach zawierających setki instancji, SMP oferuje również techniki optymalizacji, które mogą przyspieszyć uczenie modeli nawet o 20%. W tej sekcji przedstawiamy kilka z tych optymalizacji. Aby dowiedzieć się więcej, zapoznaj się z podstawowe cechy część naszej dokumentacji.

Fragmentowanie hybrydowe

Równoległość danych fragmentowanych to oszczędzająca pamięć technika rozproszonego uczenia, która dzieli stan modelu (parametry modelu, gradienty i stany optymalizatora) na urządzenia. Ta mniejsza ilość pamięci umożliwia dopasowanie większego modelu do klastra lub zwiększenie rozmiaru partii. Jednak równoległość danych podzielonych na fragmenty zwiększa również wymagania komunikacyjne zadania szkoleniowego, ponieważ artefakty modelu podzielonego na fragmenty są często zbierane podczas szkolenia z różnych urządzeń. W ten sposób stopień fragmentowania jest ważną konfiguracją, która pozwala zaoszczędzić zużycie pamięci i obciążenie komunikacyjne.

Domyślnie PyTorch FSDP fragmentuje artefakty modelu we wszystkich urządzeniach akceleracyjnych w klastrze. W zależności od zadania szkoleniowego ta metoda fragmentowania może zwiększyć obciążenie komunikacyjne i stworzyć wąskie gardło. Aby w tym pomóc, biblioteka SMP oferuje konfigurowalną hybrydową równoległość danych w formacie Shard na bazie PyTorch FSDP. Ta funkcja umożliwia ustawienie stopnia fragmentowania optymalnego dla obciążenia szkoleniowego. Wystarczy określić stopień fragmentowania w konfiguracyjnym obiekcie JSON i uwzględnić go w skrypcie szkoleniowym SMP.

Konfiguracja SMP wygląda następująco:

{ "hybrid_shard_degree": 16 }

Aby dowiedzieć się więcej o zaletach hybrydowej równoległości danych, zobacz Niemal liniowe skalowanie treningu na gigantycznych modelach na AWS. Aby uzyskać więcej informacji na temat wdrażania hybrydowego fragmentowania w istniejącym skrypcie szkoleniowym FSDP, zobacz hybrydowa równoległość współdzielonych danych w naszej dokumentacji.

Korzystaj z operacji komunikacji zbiorowej SMDDP zoptymalizowanych dla infrastruktury AWS

Możesz używać biblioteki SMP z Biblioteka rozproszonego równoległości danych (SMDDP) SageMaker aby przyspieszyć rozproszone obciążenia szkoleniowe. SMDDP zawiera zoptymalizowany AllGather operacja komunikacji zbiorowej zaprojektowana z myślą o najlepszej wydajności w przyspieszonych instancjach SageMaker p4d i p4de. W szkoleniu rozproszonym operacje komunikacji zbiorowej służą do synchronizowania informacji pomiędzy pracownikami GPU. AllGather to jedna z podstawowych operacji komunikacji zbiorowej, zwykle stosowana w przypadku równoległości danych podzielonych na fragmenty w celu materializacji parametrów warstwy przed etapami obliczeń w przód i w tył. W przypadku zadań szkoleniowych, które są wąskie gardła z powodu komunikacji, szybsze operacje zbiorowe mogą skrócić czas i koszty szkolenia bez skutków ubocznych dla konwergencji.

Aby skorzystać z biblioteki SMDDP wystarczy dodać do skryptu szkoleniowego jedynie dwie linijki kodu:

import torch.distributed as dist

# Initialize with SMDDP
import smdistributed.dataparallel.torch.torch_smddp
dist.init_process_group(backend="smddp") # Replacing "nccl"

# Initialize with SMP
import torch.sagemaker as tsm
tsm.init()

Oprócz SMP, SMDDP obsługuje open source PyTorch FSDP i DeepSpeed. Aby dowiedzieć się więcej o bibliotece SMDDP, zob Przeprowadzaj szkolenia rozproszone za pomocą biblioteki równoległości danych rozproszonych SageMaker.

Odciążanie aktywacji

Zazwyczaj przebieg do przodu szkolenia modelu oblicza aktywacje w każdej warstwie i przechowuje je w pamięci procesora graficznego do czasu zakończenia przebiegu wstecznego dla odpowiedniej warstwy. Te przechowywane aktywacje mogą zużywać znaczną ilość pamięci GPU podczas treningu. Odciążanie aktywacji to technika, która zamiast tego przenosi te tensory do pamięci procesora po przejściu w przód, a później pobiera je z powrotem do procesora graficznego, gdy są potrzebne. Takie podejście może znacznie zmniejszyć zużycie pamięci GPU podczas treningu.

Chociaż PyTorch obsługuje odciążanie aktywacji, jego implementacja jest nieefektywna i może powodować bezczynność procesorów graficznych podczas pobierania aktywacji z procesora podczas przebiegu wstecznego. Może to spowodować znaczny spadek wydajności podczas korzystania z odciążania aktywacji.

SMP v2 oferuje zoptymalizowany algorytm odciążania aktywacji, który może poprawić wydajność treningu. Implementacja SMP pobiera z wyprzedzeniem aktywacje, zanim będą potrzebne na GPU, skracając czas bezczynności.

Ponieważ SMP opiera się na interfejsach API PyTorch, umożliwienie zoptymalizowanego odciążania aktywacji wymaga tylko kilku linijek zmiany kodu. Po prostu dodaj powiązane konfiguracje (sm_activation_offloading i activation_loading_horizon parametry) i uwzględnij je w swoim skrypcie szkoleniowym.

Konfiguracja SMP wygląda następująco:

{
"activation_loading_horizon": 2,
"sm_activation_offloading": True
}

W skrypcie szkoleniowym użyj następującego kodu:

import torch.sagemaker as tsm
tsm.init()

# Native PyTorch module for activation offloading
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
apply_activation_checkpointing,
offload_wrapper,
)

model = FSDP(...)

# Activation offloading requires activation checkpointing.
apply_activation_checkpointing(
model,
check_fn=checkpoint_tformer_layers_policy,
)

model = offload_wrapper(model)

Aby dowiedzieć się więcej na temat narzędzi punktu kontrolnego PyTorch typu open source do odciążania aktywacji, zobacz checkpoint_wrapper.py skrypt w repozytorium PyTorch GitHub i Punkt kontrolny aktywacji w poście na blogu PyTorch Skalowanie multimodalnych modeli fundamentów w TorchMultimodal z Pytorch Distributed. Aby dowiedzieć się więcej na temat zoptymalizowanej implementacji odciążania aktywacji firmy SMP, zobacz odciążenie aktywacji część naszej dokumentacji.

Oprócz hybrydowego fragmentowania, SMDDP i odciążania aktywacji, SMP oferuje dodatkowe optymalizacje, które mogą przyspieszyć obciążenie szkoleniem dużych modeli. Obejmuje to zoptymalizowane punkty kontrolne aktywacji, opóźnioną inicjalizację parametrów i inne. Aby dowiedzieć się więcej, zapoznaj się z podstawowe cechy część naszej dokumentacji.

Wnioski

W miarę ciągłego zwiększania się zbiorów danych, rozmiarów modeli i klastrów szkoleniowych, wydajne szkolenia rozproszone mają coraz większe znaczenie dla terminowego i niedrogiego dostarczania modeli i produktów. Najnowsza wersja biblioteki równoległej modeli SageMaker pomaga to osiągnąć, ograniczając liczbę zmian kodu i dostosowując się do interfejsów API PyTorch FSDP, umożliwiając szkolenie na ogromnych klastrach poprzez równoległość tensorów i optymalizacje, które mogą skrócić czas szkolenia nawet o 20%.

Aby rozpocząć pracę z SMP v2, zapoznaj się z naszym dokumentacja i nasz przykładowe zeszyty.


O autorach

Roberta Van Dusena jest starszym menedżerem produktu w Amazon SageMaker. Prowadzi frameworki, kompilatory i techniki optymalizacji do szkoleń głębokiego uczenia się.

Luisa Quintelę jest Software Developer Manager dla biblioteki równoległej modelu AWS SageMaker. W wolnym czasie można go spotkać jadącego na Harleyu w rejonie SF Bay Area.

Gautama Kumara jest inżynierem oprogramowania z AWS AI Deep Learning. Pasjonuje się budowaniem narzędzi i systemów AI. W wolnym czasie lubi jeździć na rowerze i czytać książki.

Rahula Huilgola jest starszym inżynierem ds. rozwoju oprogramowania w dziale Distributed Deep Learning w Amazon Web Services.

Znak czasu:

Więcej z Uczenie maszynowe AWS