Trenuj gigantyczne modele z niemal liniowym skalowaniem przy użyciu równoległości danych sharded w Amazon SageMaker

W dążeniu do najwyższej dokładności modele uczenia głębokiego w obszarach takich jak przetwarzanie języka naturalnego i widzenie komputerowe znacznie się rozrosły w ciągu ostatnich kilku lat, często liczone w dziesiątkach do setek miliardów parametrów. Szkolenie tych gigantycznych modeli jest trudne i wymaga złożonych strategii dystrybucji. Analitycy danych i inżynierowie zajmujący się uczeniem maszynowym nieustannie szukają najlepszego sposobu na optymalizację swoich obliczeń szkoleniowych, ale zmagają się z obciążeniem komunikacyjnym, które może rosnąć wraz z ogólnym rozmiarem klastra.

Dlatego niedawno uruchomiliśmy sutwardzona równoległość danych on Amazon Sage Maker, nowa technika rozproszonego treningu oszczędzająca pamięć w Biblioteka modeli równoległych SageMaker (SMP). Równoległość danych podzielonych na fragmenty została stworzona specjalnie z myślą o modelach o ekstremalnej skali i wykorzystuje własną firmę Amazon MiCS technologia pod maską, wysiłek naukowy mający na celu zminimalizowanie skali komunikacji poprzez obniżenie kosztownych narzutów komunikacyjnych zakorzenionych w zbieraniu parametrów i synchronizacji gradientów. Dzięki modelowi GPT-30 o parametrach 2B i długości sekwencji 2048, ta nowa funkcja osiągnęła 141 TFLOP, co oznacza 39.7% przyspieszenie w porównaniu z DeepSpeed ​​ZeRO-3. W przypadku modelu 10B GPT-2 o długości sekwencji 512, ta nowa funkcja osiągnęła również 564 próbki na sekundę, co oznacza 13.9% przyspieszenie w porównaniu do Fully Sharded Data Parallel (FSDP) firmy PyTorch. Pamiętaj, że w gigantycznych szkoleniach modelowych każdy procent przyspieszenia przekłada się na zaoszczędzone pieniądze i osiągniętą produktywność w Twoim zespole.

W tym poście na blogu najpierw przyjrzymy się bliżej kluczowym wyróżnikom równoległości danych podzielonych na fragmenty i kiedy z nich korzystać. Następnie dowiesz się, jak z łatwością trenować model GPT-30 o parametrach 2B w programie SageMaker dzięki tej nowej funkcji. Na koniec porównamy wydajność z innymi opcjami open source, w szczególności przewyższając DeepSpeed ​​ZeRO nawet o 39.7% na 256 procesorach graficznych.

Jak działa równoległość danych podzielonych na fragmenty i kiedy z niego korzystać

Zanim wprowadzimy równoległość danych podzielonych na fragmenty, przyjrzyjmy się jego szerszej rodzinie technik. Ostatnie podejścia do rozproszonego szkolenia dla dużych modeli przeszły do ​​paradygmatu, w którym parametry modelu, gradienty i stany optymalizatora są współdzielone przez węzły równoległe do danych. W przeciwieństwie do paralelizmu potokowego, który ma wrodzoną złożoność polegającą na wyborze warstw do partycjonowania między urządzeniami, zwłaszcza gdy twoja platforma nie obsługuje automatyczne dzielenie modeli, ten paradygmat elegancko zachowuje prostotę równoległości danych, usuwając jednocześnie ograniczenie równoległości danych, w którym model musi mieścić się w pojedynczym procesorze graficznym.

W istniejących frameworkach, które mieszczą się w tym paradygmacie, w szczególności DeepSpeed ​​ZeRO-3 i FSDP PyTorch's upstream z FairScale, stany modeli są dzielone na części cała kolekcja GPU, strategia, która zmniejsza zużycie pamięci na każdym GPU kosztem ponoszenia dużych kosztów komunikacji, które zwiększają się wraz z rozmiarem klastra, a tym samym powodują znaczny spadek skalowalności przy skalowaniu. W przeciwieństwie do tego, równoległość danych podzielonych na fragmenty w modelu partycji biblioteki SMP stany w a świadomy skali sposób, dzieląc każdą replikę stanów modelu tylko w obrębie podzbiór GPU.

Przyjrzyjmy się bliżej partycjonowanie modelu uwzględniające skalę w MiCS, podstawowa technologia stojąca za równoległymi danymi podzielonymi na fragmenty. Intuicja stojąca za tym projektem polega na tym, że partycjonowanie stanów uczenia w całej grupie danych równoległych może nie być wymagane do uczenia modelu z dziesiątkami miliardów parametrów. Na przykład 8 procesorów graficznych V100 (po 32 GB) wystarcza do przechowywania repliki stanów modelu modelu o parametrach 10B, który wymaga około 200 GB pamięci podczas trenowania z optymalizatorem Adama przy użyciu precyzji mieszanej. Ograniczając pełną replikę stanów modelu w najmniejszy podzbiór procesorów graficznych, możemy skutecznie zmniejszyć skalę obciążenia komunikacyjnego w porównaniu z DeepSpeed ​​i PyTorch FSDP. Równoległe dane podzielone na fragmenty wykorzystują również inne techniki w MiCS, takie jak komunikacja hierarchiczna i 2-skokowa synchronizacja gradientowa. Aby uzyskać więcej informacji, sprawdź Niemal liniowe skalowanie treningu na gigantycznych modelach na AWS or MiCS: prawie liniowe skalowanie do trenowania gigantycznego modelu w chmurze publicznej.

Skąd wiesz, kiedy wybrać dane podzielone równolegle, a nie inne rozproszone techniki szkoleniowe? Ogólna zasada jest taka, że ​​jeśli Twój model ma mniej niż 1 miliard parametrów i może zmieścić się w pamięci GPU, Biblioteka równoległa danych SageMaker or Kompilator szkoleń SageMaker może ci wystarczyć. Jeśli masz większe modele językowe lub komputerowe, naszą sugestią jest wyszkolenie go za pomocą techniki równoległości danych sharded w połączeniu z punkt kontrolny aktywacji i odciążenie aktywacji najpierw w bibliotece równoległej modelu SageMaker, przed innymi technikami, takimi jak równoległość tensorowa lub równoległość rurociągu.

Używanie równoległości danych podzielonych do trenowania GPT-2 na Amazon SageMaker

Nauczmy się teraz, jak trenować model GPT-2 z równoległymi danymi podzielonymi na fragmenty, z SMP hermetyzującą złożoność dla Ciebie. Ten kompletny zeszyt samouczka przeprowadzi Cię przez cały proces, od przetwarzania danych, definiowania i przesyłania zadań szkoleniowych, po monitorowanie dzienników szkoleń. Poniżej znajduje się krótki przegląd podkreślający kluczowe kroki korzystania z tej funkcji.

1. Zaczynać

Równoległość danych podzielonych na fragmenty jest dostępna w PyTorch v1.12.0+ i działa zarówno z FP16, jak i BF16. Najłatwiejszym sposobem korzystania z biblioteki SMP jest użycie gotowego kontenera AWS Deep Learning dla PyTorch. Jeśli jednak chcesz przywieźć własny kontener Docker, możesz odwołać się do Utwórz własny kontener Dockera za pomocą biblioteki równoległej modeli rozproszonych SageMaker. Aby rozpocząć, wykonaj Zmodyfikuj skrypt szkoleniowy PyTorch aby dostosować interfejsy API SMP w swoim skrypcie szkoleniowym. W tej sekcji przedstawiamy tylko kilka głównych kroków z fragmentami kodu z gotowego do użycia skryptu szkoleniowego train_gpt_simple.py. Możesz śledzić komentarze w skrypcie i Dokument API aby dowiedzieć się więcej o tym, gdzie używane są interfejsy API SMP.

Najpierw zaimportuj i zainicjuj bibliotekę, wywołując smdistributed.modelparallel.torch.init() na początku skryptu szkoleniowego:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

Po drugie, owiń model, który ma być podzielony smdistributed.modelparallel.torch.DistributedModel i skorzystaj ze zwróconego DistributedModel obiekt w przyszłości:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_config(model_config)
model = smp.DistributedModel(model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation)

Owiń optymalizator za pomocą smdistributed.modelparallel.torch.DistributedOptimizer do zapisywania i ładowania stanów optymalizatora.

from torch import optim

optimizer = optim.Adam(
    param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
)

optimizer = smp.DistributedOptimizer(
        optimizer, 
        static_loss_scale=None, 
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
        )

Umieść logikę naprzód i wstecz w funkcji kroku i udekoruj ją smdistributed.modelparallel.torch.step.  Wszelkie obliczenia zdefiniowane wewnątrz smp.step-decorated funkcja jest wykonywana w sposób rozproszony.

@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    model.backward(loss)

    return loss

@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    
    return loss

2. Przygotuj zbiór danych

Używamy otwarty tekst internetowy to zbiór danych, którego używamy w tym przykładzie. Notatnik używa skryptu data_prep_512.py aby pobrać i wstępnie przetworzyć zbiór danych. Możesz także trenować z innymi zbiorami danych, modyfikując data_pipeline.py. Mając do czynienia z dużym zbiorem danych i modelem, możesz przyspieszyć zadanie szkoleniowe, korzystając z danych przechowywanych w Amazon FSx dla Luster, który zapewnia wysokowydajny system plików natywnie zintegrowany z Usługa Amazon Simple Storage (S3). Proszę zapoznać się z instrukcjami od Skonfiguruj kanał wprowadzania danych, aby używać Amazon FSx dla Lustre aby uzyskać wskazówki dotyczące ustawiania systemu plików FSx Luster jako kanału wprowadzania danych.

3. Rozpocznij prace szkoleniowe

Ten krok zakłada, że ​​już zmodyfikowałeś swój skrypt szkoleniowy i przygotował zestaw danych, jak wspomniano w poprzednich sekcjach. Do włącz równoległość danych podzielonych na fragmenty, po prostu ustaw sharded_data_parallel_degree  Estymator PyTorch. W tym samouczku ustawiamy sharded_data_parallel_degree=128 i instace_count=32 dla węzłów p4d.24xlarge, co oznacza, że ​​stany modelu zostaną podzielone na 128 GPU spośród wszystkich 256 GPU. Na podstawie tej wybranej wartości SMP automatycznie ustawi stopień równoległości danych na 2 (ponieważ 256/128=2), co oznacza, że ​​będziemy mieć dwie repliki dla równoległości danych. Ogólna zasada wyboru idealnej wartości dla sharded_data_parallel_degree jest dodanie jednego dodatkowego węzła do grupy udostępniania na każde 3B parametrów modelu. W tym samouczku rozmiar naszego modelu to 30B, więc powinniśmy użyć co najmniej 10 węzłów do shardingu. A ponieważ 16 węzłów (128 GPU) to najmniejsza moc z 2 powyżej progu, ustawiliśmy sharded_data_parallel_degree=128.

Do wstawiania punktów kontrolnych zapewniamy również zestaw narzędzi do wstawiania punktów kontrolnych w sharded_data_parallel_checkpoint.py , w tym narzędzie do rekonstrukcji pełnego state_dict dla zaawansowanych przypadków użycia. Na koniec możemy uruchomić rozproszone zadanie szkoleniowe, wywołując fit() w estymatorze.

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    instance_type="ml.p4d.24xlarge",
    source_dir=os.getcwd(),
    volume_size=500,
    instance_count=32,
    distribution={
        "mpi": {
            "enabled": True,
            "processes_per_host": processes_per_host,
            "custom_mpi_options": mpioptions,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "ddp": True,
                    "skip_tracing": True,
                    "delayed_parameter_initialization": True,
                    "offload_activations": True,
                    "activation_loading_horizon": 4,
                    # To enable sharded data parallelism.
                    # Here we shard model states across 128 GPUs. 
                    "sharded_data_parallel_degree": 128, 
                    "fp16": False,
                    "bf16": True,
                    # This is to disable pipeline parallelism.
                    "partitions": 1,
                },
            }
        },
    },
    framework_version="1.12",
    py_version="py38",
    hyperparameters=hyperparameters,
    checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,
    checkpoint_local_path=hyperparameters["checkpoint-dir"] if use_fsx else None,
    ...
)

smp_estimator.fit(inputs=data_channels)

4. Monitoruj zadania szkoleniowe

Możesz uzyskać dostęp do dzienników treningowych i śledzić wykorzystanie GPU i pamięci na Amazon Cloud Watch. Koniecznie zajrzyj do dzienników „algo-1” ponieważ jest to główny węzeł, którego strumień wyjściowy zawiera dzienniki zadań szkoleniowych ze wszystkich wystąpień.

Wydajność benchmarkingowa

Porównaliśmy równoległość danych podzielonych na fragmenty w bibliotece SMP na 16 i 32 węzłach p4d.24xlarge odpowiednio dla długości sekwencji 512 i 2048. Model GPT30 o parametrach 2B jest skonfigurowany do korzystania z ukrytej szerokości 7168, 48 warstw i 64 głowic. Możesz przyjąć dokładnie tę samą konfigurację, w której długość sekwencji wynosi 2048, ustawiając model_config = "gpt2-30b" w zeszycie samouczka. Przy tym ustawieniu SMP osiągnął 73.52 próbki na sekundę, co oznacza 39.7% przyspieszenie w porównaniu z DeepSpeed ​​ZeRO-3. Jeśli rozmiar Twojego tokena wynosi 500 miliardów, to przyspieszenie oznacza prawie 367 godzin oszczędności na węzłach p4d.24xlarge, co odpowiada ponad 12,000 XNUMX USD budżetu zaoszczędzonego na szkolenie! Poniższa tabela podsumowuje nasze wyniki testów porównawczych.

1/ Dla każdej konfiguracji modelu przetestowaliśmy różne funkcje, etapy i konfiguracje w DeepSpeed ​​ZeRO i wybraliśmy tę, która zapewnia najlepszą przepustowość jako podstawę DeepSpeed. Test porównawczy został uruchomiony Elastyczna chmura obliczeniowa Amazon (Amazonka EC2). 2/ Wyniki te opierają się na ulepszonych zbiorach komunikacyjnych zoptymalizowanych pod kątem AWS, które zostaną wkrótce udostępnione. 3/ Czas na trenowanie jest prognozowany na podstawie prędkości na podstawie liczby przetworzonych tokenów.

Podsumowując, zaobserwowaliśmy konsekwentnie wyższą przepustowość z równoległością danych podzielonych na fragmenty w SMP w porównaniu z DeepSpeed ​​w wielu modelach i konfiguracjach. Ta nowa funkcja wykazała również lepszą wydajność pamięci w porównaniu z DeepSpeed, umożliwiając SMP dopasowanie większego rozmiaru partii i zmniejszenie poziomu akumulacji gradientu wymaganego do dopasowania do konkretnego globalnego rozmiaru partii.

Wnioski

W tym poście przedstawiliśmy nową technikę uczenia rozproszonego — równoległość danych podzielonych na fragmenty — i sposób, w jaki przyspiesza ona gigantyczne szkolenie modeli z niemal liniowym skalowaniem w Amazon SageMaker. Przeszliśmy również przez proces trenowania modelu GPT-2 za pomocą nowej techniki następującej po tym kompletny przykład. Możesz śledzić Przykłady Amazon SageMaker Repozytorium GitHub aby śledzić wszystkie równoległe przykłady modeli SageMaker lub wziąć udział w naszym następnym rozproszone warsztaty szkoleniowe. Aby dowiedzieć się więcej o równoległości danych podzielonych na fragmenty, zobacz dokumentacja.

O autorach

Emily Webber dołączył do AWS zaraz po uruchomieniu SageMaker i od tego czasu próbuje opowiedzieć o tym światu! Poza tworzeniem nowych doświadczeń związanych z ML dla klientów, Emily lubi medytować i studiować buddyzm tybetański.

Czy Karakuś jest starszym naukowcem stosowanym w AWS, optymalizującym rozproszone uczenie głębokie na dużą skalę w AWS. Jego zainteresowania badawcze obejmują głębokie uczenie, optymalizację rozproszoną, systemy rozproszone i teorię informacji. Poza pracą lubi jeździć na rowerze, podróżować, czytać i uczyć się.

Rahula Huilgola jest starszym inżynierem oprogramowania w AWS. Pracuje nad rozproszonymi systemami głębokiego uczenia się, aby ułatwić i wydajnie trenować duże modele uczenia głębokiego w chmurze. W wolnym czasie lubi fotografować, jeździć na rowerze i pracować w ogrodzie.

Suhit Kodgule jest inżynierem rozwoju oprogramowania w grupie AWS Artificial Intelligence pracującym nad frameworkami głębokiego uczenia się. W wolnym czasie lubi wędrówki, podróże i gotowanie.

Erin Ho jest Product Managerem ds. głębokiego uczenia AWS. Pracuje nad produktami, które ułatwiają klientom trenowanie modeli głębokiego uczenia się na AWS. Poza pracą lubi chodzić po górach i jeździć na nartach.