Ein PyTorch Trainingsskript ändern

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

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.