Automatisiertes Teilen mit PyTorch Manuelles Teilen mit PyTorch Überlegungen Nicht unterstützte Framework-Funktionen

Ein PyTorch Trainingsskript ändern

In diesem Abschnitt erfahren Sie, wie Sie PyTorch Trainingsskripte ändern, um die SageMaker Modellparallelitätsbibliothek für automatische Partitionierung und manuelle Partitionierung zu konfigurieren.

Anmerkung

Informationen darüber, welche PyTorch Versionen von der Bibliothek unterstützt werden, finden Sie unter. Unterstützte Frameworks und AWS-Regionen

Tipp

end-to-end Notebook-Beispiele, die veranschaulichen, wie ein PyTorch Trainingsskript mit der SageMaker Modellparallelitätsbibliothek verwendet wird, finden Sie unter. Beispiele für die Amazon SageMaker Model Parallelism Library v1

Beachten Sie, dass die automatische Partitionierung standardmäßig aktiviert ist. Sofern nicht anders angegeben, verwenden die folgenden Skripten automatische Partitionierung.

Automatisiertes Teilen mit PyTorch

Die folgenden Änderungen am Trainingsskript sind erforderlich, um ein PyTorch Trainingsskript mit SageMaker der Modellparallelismus-Bibliothek auszuführen:

Importieren und initialisieren Sie die Bibliothek mit smdistributed.modelparallel.torch.init().
Schließen Sie das Modell mit smdistributed.modelparallel.torch.DistributedModel um. Beachten Sie, dass alle Tensoren, die von der forward Methode des zugrunde liegenden nn.Module Objekts zurückgegeben werden, über modellparallele Geräte übertragen werden, was zu Kommunikationsaufwand führt. Daher sollten alle Tensoren, die außerhalb der Aufrufmethode nicht benötigt werden (z. B. Zwischenaktivierungen), nicht zurückgegeben werden.

Anmerkung
Für das FP16-Training müssen Sie den Kontextmanager smdistributed.modelparallel.torch.model_creation()verwenden, um das Modell zu umschließen. Weitere Informationen finden Sie unter FP16Training mit Modellparallelität.
Umschließen Sie den Optimierer mit smdistributed.modelparallel.torch.DistributedOptimizer.

Anmerkung
Für das FP16-Training müssen Sie eine statische oder dynamische Verlust-Scaling einrichten. Weitere Informationen finden Sie unter FP16Training mit Modellparallelität.
Verwenden Sie das zurückgegebene DistributedModel Objekt anstelle eines Benutzermodells.
Fügen Sie die Vorwärts- und Rückwärtslogik in eine Schrittfunktion ein und dekorieren Sie sie mit smdistributed.modelparallel.torch.step.
Beschränken Sie jeden Prozess auf sein eigenes Gerät durch torch.cuda.set_device(smp.local_rank()).
Verschieben Sie die Eingangstensoren mithilfe der .to() API vor dem smp.step Aufruf auf die GPU (siehe Beispiel unten).
Ersetzen Sie torch.Tensor.backward und torch.autograd.backward mit DistributedModel.backward.
Führen Sie die Nachbearbeitung der Ausgaben für alle Mikrobatches mithilfe von StepOutput Methoden wie reduce_mean durch.
Wenn es einen Bewertungsschritt gibt, platzieren Sie die Vorwärtslogik auf ähnliche Weise in einer mit -smp.stepdekorierten Funktionen und bearbeiten Sie die Ausgaben mithilfe der StepOutputAPI nach.
drop_last=True in DataLoader einstellen. Alternativ können Sie einen Stapel in der Trainingsschleife manuell überspringen, wenn die Batchgröße nicht durch die Anzahl der Mikrobatches teilbar ist.

Weitere Informationen über die API SageMaker der Modellparallelismus-Bibliothek finden Sie in der API-Dokumentation.


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchnet.dataset import SplitDataset
from torchvision import datasets

import smdistributed.modelparallel.torch as smp

class GroupedNet(nn.Module):
    def __init__(self):
        super(GroupedNet, self).__init__()
        # define layers

    def forward(self, x):
        # define forward pass and return model outputs


# smdistributed: Define smp.step. Return any tensors needed outside.
@smp.step
def train_step(model, data, target):
    output = model(data)
    loss = F.nll_loss(output, target, reduction="mean")
    model.backward(loss)
    return output, loss


def train(model, device, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # smdistributed: Move input tensors to the GPU ID used by the current process,
        # based on the set_device call.
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        # Return value, loss_mb is a StepOutput object
        _, loss_mb = train_step(model, data, target)

        # smdistributed: Average the loss across microbatches.
        loss = loss_mb.reduce_mean()

        optimizer.step()

# smdistributed: initialize the backend
smp.init()

# smdistributed: Set the device to the GPU ID used by the current process.
# Input tensors should be transferred to this device.
torch.cuda.set_device(smp.local_rank())
device = torch.device("cuda")

# smdistributed: Download only on a single process per instance.
# When this is not present, the file is corrupted by multiple processes trying
# to download and extract at the same time
dataset = datasets.MNIST("../data", train=True, download=False)

# smdistributed: Shard the dataset based on data-parallel ranks
if smp.dp_size() > 1:
    partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())}
    dataset = SplitDataset(dataset, partitions=partitions_dict)
    dataset.select(f"{smp.dp_rank()}")

# smdistributed: Set drop_last=True to ensure that batch size is always divisible
# by the number of microbatches
train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True)

model = GroupedNet()
optimizer = optim.Adadelta(model.parameters(), lr=4.0)

# smdistributed: Use the DistributedModel container to provide the model
# to be partitioned across different ranks. For the rest of the script,
# the returned DistributedModel object should be used in place of
# the model provided for DistributedModel class instantiation.
model = smp.DistributedModel(model)
optimizer = smp.DistributedOptimizer(optimizer)

train(model, device, train_loader, optimizer)

Manuelles Teilen mit PyTorch

Verwenden Sie smp.partition Kontextmanager, um Module auf bestimmten Geräten zu platzieren. Jedes Modul, das sich nicht in einem smp.partition Kontext befindet, wird in den default_partition platziert. Das default_partition muss angegeben werden, wenn auto_partition auf False gesetzt ist. Die Module, die in einem bestimmten smp.partition Kontext erstellt werden, werden auf der entsprechenden Partition platziert.

Weitere Informationen zur API SageMaker der Modellparallelismus-Bibliothek finden Sie in der API-Dokumentation.


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchnet.dataset import SplitDataset
from torchvision import datasets

import smdistributed.modelparallel.torch as smp

class GroupedNet(nn.Module):
    def __init__(self):
        super(GroupedNet, self).__init__()
        with smp.partition(0):
            # define child modules on device 0
        with smp.partition(1):
            # define child modules on device 1

    def forward(self, x):
        # define forward pass and return model outputs


# smdistributed: Define smp.step. Return any tensors needed outside.
@smp.step
def train_step(model, data, target):
    output = model(data)
    loss = F.nll_loss(output, target, reduction="mean")
    model.backward(loss)
    return output, loss


def train(model, device, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # smdistributed: Move input tensors to the GPU ID used by the current process,
        # based on the set_device call.
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        # Return value, loss_mb is a StepOutput object
        _, loss_mb = train_step(model, data, target)

        # smdistributed: Average the loss across microbatches.
        loss = loss_mb.reduce_mean()

        optimizer.step()

# smdistributed: initialize the backend
smp.init()

# smdistributed: Set the device to the GPU ID used by the current process.
# Input tensors should be transferred to this device.
torch.cuda.set_device(smp.local_rank())
device = torch.device("cuda")

# smdistributed: Download only on a single process per instance.
# When this is not present, the file is corrupted by multiple processes trying
# to download and extract at the same time
dataset = datasets.MNIST("../data", train=True, download=False)

# smdistributed: Shard the dataset based on data-parallel ranks
if smp.dp_size() > 1:
    partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())}
    dataset = SplitDataset(dataset, partitions=partitions_dict)
    dataset.select(f"{smp.dp_rank()}")

# smdistributed: Set drop_last=True to ensure that batch size is always divisible
# by the number of microbatches
train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True)

model = GroupedNet()
optimizer = optim.Adadelta(model.parameters(), lr=4.0)

# smdistributed: Use the DistributedModel container to provide the model
# to be partitioned across different ranks. For the rest of the script,
# the returned DistributedModel object should be used in place of
# the model provided for DistributedModel class instantiation.
model = smp.DistributedModel(model)
optimizer = smp.DistributedOptimizer(optimizer)

train(model, device, train_loader, optimizer)

Überlegungen

Wenn Sie ein PyTorch Trainingsskript mithilfe SageMaker der Modellparallelismus-Bibliothek konfigurieren, sollten Sie Folgendes beachten:

Wenn Sie eine Optimierungstechnik verwenden, die auf globalen Gradientennormen basiert, z. B. der Gradientennorm aus dem gesamten Modell, wie z. B. einige Varianten des LAMB-Optimizers oder des globalen Gradientenclippings, müssen Sie alle Normen aus den Modellpartitionen zusammenstellen, um ihre Richtigkeit zu überprüfen. Zu diesem Zweck können Sie die grundlegenden Kommunikationsdatentypen der Bibliothek verwenden.
Alle torch.Tensor Argumente für die Vorwärtsmethoden von nn.Modules in Ihrem Modell müssen bei der Berechnung der Modulausgabe verwendet werden. Mit anderen Worten, die Bibliothek unterstützt nicht den Fall, dass es ein torch.Tensor Argument für ein Modul gibt, von dem die Modulausgabe nicht abhängt.
Das Argument für den smp.DistributedModel.backward() Aufruf muss von allen Modellausgaben abhängen. Mit anderen Worten, es darf keine Ausgabe des smp.DistributedModel.forward Aufrufs geben, die nicht bei der Berechnung des Tensors verwendet wird, der in den smp.DistributedModel.backward Aufruf eingespeist wird.
Wenn Ihr Code torch.cuda.synchronize() Aufrufe enthält, müssen Sie möglicherweise torch.cuda.set_device(smp.local_rank()) unmittelbar vor dem Synchronisierungsaufruf anrufen. Andernfalls könnten in Gerät 0 unnötige CUDA-Kontexte erstellt werden, die unnötig Speicherplatz verbrauchen.
Da sich die Bibliothek nn.Modules auf unterschiedlichen Geräten befindet, dürfen die Module im Modell nicht von einem globalen Status abhängen, der im Inneren von smp.step geändert wird. Jeder Status, der während des gesamten Trainings unverändert bleibt oder außerhald von smp.step so verändert wird, dass er für alle Prozesse sichtbar ist, ist zulässig.
Sie müssen das Modell nicht auf die GPU verschieben (z. B. verwenden von model.to(device)), wenn Sie die Bibliothek verwenden. Wenn Sie versuchen, das Modell auf die GPU zu verschieben, bevor das Modell partitioniert wurde (vor dem ersten smp.step Aufruf), wird der Move-Aufruf ignoriert. Die Bibliothek verschiebt den Teil des Modells, der einem Rang zugewiesen wurde, automatisch auf ihre GPU. Sobald das Training mit der Bibliothek begonnen hat, sollten Sie das Modell nicht auf die CPU verschieben und es verwenden, da es sonst keine korrekten Parameter für Module enthält, die nicht der vom Prozess gespeicherten Partition zugewiesen sind. Wenn Sie ein Modell neu trainieren oder es ohne die Bibliothek für Inferenz verwenden möchten, nachdem es mit der Modellparallelismus-Bibliothek trainiert wurde, empfiehlt es sich, das vollständige Modell mithilfe unserer Checkpoint-API zu speichern und es wieder in ein reguläres Modul zu laden. PyTorch
Wenn Sie eine Liste von Modulen haben, bei denen die Ausgabe eines Moduls in ein anderes einfließen kann, kann das Ersetzen dieser Liste durch die Leistung erheblich verbessern. nn.Sequential
Das Gewichtsupdate (optimizer.step()) muss außerhalb von smp.step erfolgen, da dann der gesamte Rückwärtsdurchlauf abgeschlossen ist und die Farbverläufe bereit sind. Wenn Sie ein Hybridmodell mit Modell- und Datenparallelität verwenden, ist zu diesem Zeitpunkt auch garantiert, dass die Gradienten beendet AllReduce sind.
Wenn Sie die Bibliothek in Kombination mit Datenparallelität verwenden, stellen Sie sicher, dass die Anzahl der Batches auf allen datenparallelen Rängen gleich ist, damit Sie AllReduce nicht auf einen Rang warten, der nicht am Schritt teilnimmt.

Wenn Sie einen Trainingsjob mit einem ml.p4d-Instance-Typ (z. B. ml.p4d.24xlarge) starten, müssen Sie die Dataloader-Variable num_workers=0 festlegen. Sie können DataLoader Ihren beispielsweise wie folgt definieren:


dataloader = torch.utils.data.DataLoader(
            data,
            batch_size=batch_size,
            num_workers=0,
            pin_memory=True,
            drop_last=True,
            shuffle=shuffle,
        )

Die Eingaben für smp.step müssen die Modelleingaben sein, die von DataLoader generiert wurden. Der Grund dafür ist, dass smp.step die Eingabetensoren intern entlang der Stapeldimension aufteilt und sie in eine Pipeline einfügt. Dies bedeutet, dass es nicht funktioniert, DataLoader sich selbst an die smp.step Funktion zur Generierung der darin enthaltenen Modelleingaben zu übergeben.

Wenn Sie beispielsweise a DataLoader wie folgt definieren:
```
train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True)
```
Sie sollten auf die Modelleingaben zugreifen, die von generiert wurden, train_loader und diese an eine smp.step dekorierte Funktion übergeben. Übergeben Sie train_loader nicht direkt an die smp.step Funktion.
```
def train(model, device, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        ...
        _, loss_mb = train_step(model, data, target)
        ...

@smp.step
def train_step(model, data, target):
    ...
    return output, loss
```

Die Eingangstensoren für smp.step müssen mithilfe der .to() API auf das aktuelle Gerät verschoben werden, was nach dem torch.cuda.set_device(local_rank()) Aufruf erfolgen muss.

Sie können z. B. wie folgt die Funktion train definieren. Diese Funktion fügt dem aktuellen Gerät mithilfe der .to() API data und target hinzu, bevor diese Eingangstensoren zum Aufrufen von train_step verwendet werden.


def train(model, device, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # smdistributed: Move input tensors to the GPU ID used by the current process,
        # based on the set_device call.
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        # Return value, loss_mb is a StepOutput object
        _, loss_mb = train_step(model, data, target)

        # smdistributed: Average the loss across microbatches.
        loss = loss_mb.reduce_mean()

        optimizer.step()

Die Eingangstensoren für diese smp.set dekorierte Funktion wurden in der obigen train Funktion auf das aktuelle Gerät verschoben. Das Modell muss nicht auf das aktuelle Gerät verschoben werden. Die Bibliothek verschiebt den Teil des Modells, der einem Rang zugewiesen ist, automatisch auf ihre GPU.


@smp.step
def train_step(model, data, target):
    output = model(data)
    loss = F.nll_loss(output, target, reduction="mean")
    model.backward(loss)
    return output, loss

Nicht unterstützte Framework-Funktionen

Die folgenden PyTorch Funktionen werden von der Modellparallelitätsbibliothek nicht unterstützt SageMaker:

Wenn Sie Datenparallelität mit dem nativen PyTorch DDP verwenden, wird das torch.nn.parallel.DistributedDataParallelWrapper-Modul von der Bibliothek nicht unterstützt. Die Bibliothek verwaltet intern die Integration mit PyTorch DDP, einschließlich Parameterübertragung und Gradient. AllReduce Bei Verwendung der Bibliothek werden Modulpuffer zu Beginn des Trainings nur einmal übertragen. Wenn Ihr Modell über Modulpuffer verfügt, die bei jedem Schritt über datenparallele Gruppen hinweg synchronisiert werden müssen, können Sie dies über die torch.distributed API tun, indem Sie die Prozessgruppe verwenden, die über smp.get_dp_process_group() abgerufen werden kann.
Für gemischtes Präzisionstraining wird das apex.amp Modul nicht unterstützt. Es wird empfohlen, die Bibliothek mit automatischer Mixed-Precision torch.cuda.amp zu verwenden, mit der Ausnahme, dass smp.amp.GradScaler anstelle der Implementierung in Torch verwendet wird.
torch.jit.ScriptModules und ScriptFunctions werden von smp.DistributedModel nicht unterstützt.
apex : FusedLayerNorm, FusedAdam FusedLAMB, und FusedNovoGrad von apex werden nicht unterstützt. Sie können stattdessen die Bibliotheksimplementierungen dieser durch smp.optimizers und smp.nn-APIs verwenden.

Warnung JavaScript ist in Ihrem Browser nicht verfügbar oder deaktiviert.

Zur Nutzung der AWS-Dokumentation muss JavaScript aktiviert sein. Weitere Informationen finden auf den Hilfe-Seiten Ihres Browsers.

Dokumentkonventionen

TensorFlow

Schritt 2: Starten eines Trainingsjobs