Trainieren Sie gigantische Modelle mit nahezu linearer Skalierung unter Verwendung von Sharded-Data-Parallelität auf Amazon SageMaker

Im Streben nach überlegener Genauigkeit sind Deep-Learning-Modelle in Bereichen wie der Verarbeitung natürlicher Sprache und Computer Vision in den letzten Jahren erheblich an Größe gewachsen und umfassen häufig Zehn- bis Hundertmilliarden von Parametern. Das Training dieser gigantischen Modelle ist anspruchsvoll und erfordert komplexe Vertriebsstrategien. Datenwissenschaftler und Ingenieure für maschinelles Lernen suchen ständig nach der besten Möglichkeit, ihre Trainingsberechnungen zu optimieren, kämpfen jedoch mit dem Kommunikationsaufwand, der mit der Gesamtclustergröße zunehmen kann.

Aus diesem Grund haben wir kürzlich gestartet sgehärtete Datenparallelität on Amazon Sage Maker, eine neue speichersparende verteilte Trainingstechnik in der SageMaker Model Parallel (SMP)-Bibliothek. Sharded Data Parallelism wurde speziell für Modelle im extremen Maßstab entwickelt und verwendet Amazon intern MiCS Technologie unter der Haube, eine wissenschaftliche Anstrengung zur Minimierung des Kommunikationsumfangs durch Reduzierung des teuren Kommunikationsaufwands, der in der Parametererfassung und Gradientensynchronisierung verwurzelt ist. Mit einem GPT-30-Modell mit 2 B Parametern und einer Sequenzlänge von 2048 erreichte diese neue Funktion 141 TFLOPs, eine Geschwindigkeitssteigerung von 39.7 % im Vergleich zu DeepSpeed ​​ZeRO-3. Bei einem 10B-GPT-2-Modell mit einer Sequenzlänge von 512 erreichte diese neue Funktion auch 564 Samples pro Sekunde, eine Geschwindigkeitssteigerung von 13.9 % im Vergleich zu PyTorchs Fully Sharded Data Parallel (FSDP). Denken Sie daran, dass sich beim gigantischen Modelltraining jeder Prozentsatz der Beschleunigung in eingesparte Dollar und gesteigerte Produktivität in Ihrem Team niederschlägt.

In diesem Blog-Beitrag werfen wir zunächst einen genaueren Blick auf die Hauptunterscheidungsmerkmale der Sharded-Data-Parallelität und wann sie verwendet werden sollte. Anschließend erfahren Sie, wie Sie mit dieser neuen Funktion ein GPT-30-Modell mit 2 B-Parametern auf SageMaker mühelos trainieren. Abschließend vergleichen wir die Leistung mit anderen Open-Source-Optionen und übertreffen insbesondere DeepSpeed ​​ZeRO um bis zu 39.7 % auf 256-GPUs.

Funktionsweise der Sharded-Data-Parallelität und wann sie zu verwenden ist

Bevor wir die Sharded-Data-Parallelität einführen, werfen wir einen Blick auf ihre breitere Technikfamilie. Neuere verteilte Trainingsansätze für große Modelle sind zu einem Paradigma übergegangen, bei dem Modellparameter, Gradienten und Optimiererzustände über datenparallele Knoten hinweg geteilt werden. Im Gegensatz zur Pipeline-Parallelität, die die angeborene Komplexität der Auswahl von Ebenen zur Partitionierung über Geräte hinweg aufweist, insbesondere wenn Ihr Framework dies nicht unterstützt automatisiertes Modellsplitting, bewahrt dieses Paradigma auf elegante Weise die Einfachheit der Datenparallelität und entfernt gleichzeitig die Beschränkung der Datenparallelität, wonach ein Modell in eine einzelne GPU passen muss.

In bestehenden Frameworks, die unter dieses Paradigma fallen, insbesondere DeepSpeed ​​ZeRO-3 und PyTorchs FSDP, das von FairScale vorgelagert ist, werden Modellzustände aufgeteilt alle GPUs, eine Strategie, die den Speicherverbrauch auf jeder GPU auf Kosten eines hohen Kommunikations-Overheads senkt, der mit der Clustergröße zunimmt und daher dazu führt, dass die Skalierbarkeit bei Skalierung erheblich abnimmt. Im Gegensatz dazu gibt es in a skalenbewusst Weise durch Partitionieren jeder Kopie von Modellzuständen nur innerhalb Eine Teilmenge von GPUs.

Schauen wir uns die näher an Skalierungsbewusste Modellpartitionierung in MiCS, der Kerntechnologie hinter Sharded Data Parallel. Die Intuition hinter diesem Design ist, dass das Partitionieren von Trainingszuständen über die gesamte datenparallele Gruppe möglicherweise nicht erforderlich ist, um ein Modell mit mehreren zehn Milliarden Parametern zu trainieren. Beispielsweise reichen 8 V100-GPUs (jeweils 32 GB) aus, um die Modellzustandsreplik eines 10B-Parametermodells zu speichern, das beim Training mit dem Adam-Optimierer mit gemischter Genauigkeit etwa 200 GB Speicher benötigt. Durch die Einschränkung einer vollständigen Nachbildung von Modellzuständen in der kleinste Untergruppe von GPUs können wir den Umfang des Kommunikations-Overheads im Vergleich zu DeepSpeed ​​und PyTorch FSDP effektiv reduzieren. Sharded Data Parallel nutzt auch andere Techniken in MiCS, wie z. B. hierarchische Kommunikation und 2-Hop-Gradientensynchronisierung. Weitere Informationen finden Sie unter Nahezu lineare Skalierung des gigantischen Modelltrainings auf AWS or MiCS: Nahezu lineare Skalierung zum Trainieren gigantischer Modelle in der Public Cloud.

Woher wissen Sie nun, wann Sie aufgeteilte Daten parallel gegenüber anderen verteilten Trainingsmethoden wählen sollten? Die allgemeine Regel lautet: Wenn Ihr Modell weniger als 1 Milliarde Parameter hat und in den GPU-Speicher passt, Parallelbibliothek für SageMaker-Daten or SageMaker-Trainingscompiler kann für dich ausreichen. Wenn Sie über größere Sprach- oder Computer-Vision-Modelle verfügen, schlagen wir vor, diese mit der Sharded-Data-Parallelismus-Technik in Kombination mit zu trainieren Aktivierungs-Checkpointing und Aktivierung auslagern in der parallelen SageMaker-Modellbibliothek zuerst, bevor andere Techniken wie z Tensorparallelität oder Pipeline-Parallelität.

Verwenden von Sharded-Data-Parallelität zum Trainieren von GPT-2 auf Amazon SageMaker

Lassen Sie uns nun lernen, wie man ein GPT-2-Modell parallel mit fragmentierten Daten trainiert, wobei SMP die Komplexität für Sie kapselt. Dies komplettes Tutorial-Notizbuch führt Sie durch den gesamten Prozess, von der Datenverarbeitung über das Definieren und Senden von Trainingsaufträgen bis hin zur Überwachung der Trainingsprotokolle. Im Folgenden finden Sie eine kurze Übersicht, in der die wichtigsten Schritte zur Verwendung dieser Funktion hervorgehoben werden.

1. Loslegen

Sharded Data Parallelism ist in PyTorch v1.12.0+ verfügbar und funktioniert sowohl mit FP16 als auch mit BF16. Der einfachste Weg, die SMP-Bibliothek zu verwenden, ist ein vorgefertigter AWS Deep Learning Container für PyTorch. Wenn Sie jedoch Ihren eigenen Docker-Container mitbringen möchten, können Sie auf verweisen Erstellen Sie Ihren eigenen Docker-Container mit der SageMaker Distributed Model Parallel Library. Um loszulegen, folgen Sie Ändern Sie ein PyTorch-Trainingsskript um die APIs von SMPs in Ihrem Trainingsskript anzupassen. In diesem Abschnitt nennen wir nur einige Hauptschritte mit Codeschnipseln aus dem gebrauchsfertigen Trainingsskript train_gpt_simple.py. Sie können den Kommentaren im Skript folgen und API-Dokument um mehr darüber zu erfahren, wo SMP-APIs verwendet werden.

Importieren und initialisieren Sie zunächst die Bibliothek durch Aufrufen smdistributed.modelparallel.torch.init() am Anfang des Trainingsskripts:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

Als zweites wickeln Sie das zu partitionierende Modell ein smdistributed.modelparallel.torch.DistributedModel und verwenden Sie die zurückgegebenen DistributedModel Objekt geht voran:

from transformers import AutoModelForCausalLM

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

Wickeln Sie den Optimierer mit ein smdistributed.modelparallel.torch.DistributedOptimizer zum Speichern und Laden von Optimiererzuständen.

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},
        )

Setzen Sie die Vorwärts- und Rückwärtslogik in eine Schrittfunktion und dekorieren Sie sie mit smdistributed.modelparallel.torch.step.  Jede Berechnung, die innerhalb der definiert ist smp.step-decorated Die Funktion wird verteilt ausgeführt.

@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. Bereiten Sie den Datensatz vor

Wir nutzen die openwebtext ist der Datensatz, den wir in diesem Beispiel verwenden. Das Notebook verwendet das Skript data_prep_512.py um den Datensatz herunterzuladen und vorzuverarbeiten. Sie können auch mit anderen Datensätzen trainieren, indem Sie sie ändern data_pipeline.py. Beim Umgang mit großen Datensätzen und Modellen können Sie den Trainingsjob beschleunigen, indem Sie Daten verwenden, die in gespeichert sind Amazon FSx für Lustre, das ein nativ integriertes Hochleistungsdateisystem bereitstellt Amazon Simple Storage-Service (S3). Bitte beachten Sie die Anweisungen von Konfigurieren Sie den Dateneingabekanal für die Verwendung von Amazon FSx für Lustre für Anleitungen zum Festlegen eines FSx-Luster-Dateisystems als Dateneingabekanal.

3. Starten Sie die Trainingsjobs

Dieser Schritt setzt voraus, dass Sie dies bereits getan haben Ihr Trainingsskript geändert und den Datensatz wie in den vorangegangenen Abschnitten erwähnt vorbereitet. Zu Sharded Data Parallelism aktivieren, setzen Sie einfach die sharded_data_parallel_degree der PyTorch-Schätzer. In diesem Tutorial legen wir fest sharded_data_parallel_degree=128 und instace_count=32 für p4d.24xlarge-Knoten, was darauf hinweist, dass die Modellzustände auf 128 GPUs unter den insgesamt 256 GPUs aufgeteilt werden. Basierend auf diesem ausgewählten Wert setzt SMP dann automatisch den Grad der Datenparallelität auf 2 (weil 256/128 = 2), was bedeutet, dass wir zwei Replikate für die Datenparallelität haben. Eine allgemeine Regel für die Auswahl eines idealen Werts für sharded_data_parallel_degree besteht darin, pro 3B Modellparameter einen weiteren Knoten zur gemeinsamen Nutzung hinzuzufügen. In diesem Tutorial beträgt unsere Modellgröße 30 B, daher sollten wir mindestens 10 Knoten für das Sharding verwenden. Und weil 16 Knoten (128 GPUs) die kleinste Potenz von 2 über dem Schwellenwert ist, legen wir fest sharded_data_parallel_degree=128.

Für das Checkpointing bieten wir auch eine Reihe von Checkpointing-Dienstprogrammen in sharded_data_parallel_checkpoint.py , einschließlich eines Dienstprogramms zum Rekonstruieren des vollständigen state_dict für fortgeschrittene Anwendungsfälle. Schließlich können wir einen verteilten Trainingsjob starten, indem wir fit() auf dem Estimator aufrufen.

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. Überwachen Sie die Trainingsjobs

Sie können auf die Trainingsprotokolle zugreifen und die GPU- und Speicherauslastung verfolgen Amazon CloudWatch. Schauen Sie sich unbedingt die Protokolle von „algo-1“ an da dies der Hauptknoten ist, dessen Ausgabestrom die Trainingsjobprotokolle von allen Instanzen enthält.

Benchmarking-Leistung

Wir haben die Sharded-Daten-Parallelität in der SMP-Bibliothek sowohl auf 16 als auch auf 32 p4d.24xlarge-Knoten für die Sequenzlänge 512 bzw. 2048 getestet. Das GPT30-Modell mit 2B-Parametern ist so konfiguriert, dass es eine verborgene Breite von 7168, 48 Schichten und 64 Köpfen verwendet. Sie können genau dieselbe Konfiguration übernehmen, bei der die Sequenzlänge 2048 beträgt, indem Sie die Einstellung vornehmen model_config = "gpt2-30b" im Tutorial-Notizbuch. Mit dieser Einstellung erreichte SMP 73.52 Samples pro Sekunde, eine Beschleunigung von 39.7 % im Vergleich zu DeepSpeed ​​ZeRO-3. Wenn Ihre Token-Größe 500 Milliarden beträgt, bedeutet diese Beschleunigung fast 367 Stunden Einsparungen auf p4d.24xlarge-Knoten, was einem eingesparten Budget von mehr als 12,000 US-Dollar pro Training entspricht! Die folgende Tabelle fasst unsere Benchmark-Ergebnisse zusammen.

1/ Für jede Modellkonfiguration haben wir verschiedene Funktionen, Stufen und Konfigurationen in DeepSpeed ​​ZeRO getestet und diejenige ausgewählt, die den besten Durchsatz als DeepSpeed-Basislinie bietet. Der Benchmark wurde ausgeführt Amazon Elastic Compute-Cloud (Amazon EC2). 2/ Diese Ergebnisse beruhen auf verbesserten, für AWS optimierten Kommunikationskollektiven, die in Kürze verfügbar sein werden. 3/ Die Zeit zum Trainieren wird aus der Geschwindigkeit basierend auf der Anzahl der verarbeiteten Token hochgerechnet.

Zusammenfassend haben wir bei einer Reihe von Modellen und Konfigurationen einen durchweg höheren Durchsatz mit Parallelität von Shard-Daten in SMP im Vergleich zu DeepSpeed ​​beobachtet. Diese neue Funktion zeigte auch eine bessere Speichereffizienz im Vergleich zu DeepSpeed, wodurch SMP eine größere Stapelgröße anpassen und den Grad der Gradientenakkumulation reduzieren kann, der erforderlich ist, um einer bestimmten globalen Stapelgröße gerecht zu werden.

Zusammenfassung

In diesem Beitrag haben wir eine neue verteilte Trainingstechnik vorgestellt – Sharded Data Parallelism – und wie sie das gigantische Modelltraining mit nahezu linearer Skalierung auf Amazon SageMaker beschleunigt. Wir haben auch erklärt, wie man ein GPT-2-Modell mit der neuen Technik im Anschluss daran trainiert komplettes Beispiel. Sie können dem folgen Amazon SageMaker-Beispiele GitHub-Repository um alle parallelen Beispiele für SageMaker-Modelle zu verfolgen oder an unserer nächsten teilzunehmen verteilte Schulungsworkshops. Um mehr über die Parallelität von Sharding-Daten zu erfahren, lesen Sie bitte die Dokumentation.

Über die Autoren

Emily Weber trat AWS kurz nach dem Start von SageMaker bei und versucht seitdem, der Welt davon zu erzählen! Abgesehen davon, dass sie neue ML-Erfahrungen für Kunden entwickelt, meditiert Emily gerne und studiert den tibetischen Buddhismus.

Kann Karakus ist Senior Applied Scientist bei AWS und optimiert verteiltes Deep Learning in großem Maßstab auf AWS. Seine Forschungsinteressen umfassen Deep Learning, verteilte Optimierung, verteilte Systeme und Informationstheorie. Außerhalb der Arbeit fährt er gerne Fahrrad, reist, liest und lernt.

Rahul Huilgol ist Senior Software Engineer bei AWS. Er arbeitet an verteilten Deep-Learning-Systemen, um das Training großer Deep-Learning-Modelle in der Cloud einfach und performant zu gestalten. In seiner Freizeit fotografiert er gerne, fährt Rad und arbeitet im Garten.

Suhit Kodgule ist ein Softwareentwicklungsingenieur in der AWS-Gruppe für künstliche Intelligenz, der an Deep-Learning-Frameworks arbeitet. In seiner Freizeit wandert, reist und kocht er gerne.

Erin Ho ist Produktmanager für AWS Deep Learning. Sie arbeitet an Produkten, die es Kunden erleichtern, Deep-Learning-Modelle auf AWS zu trainieren. Zum Spaß außerhalb der Arbeit geht sie gerne wandern und Ski fahren.