Ändern Sie ein PyTorch Trainingsskript

Wenn Sie eine Optimierungstechnik verwenden, die auf globalen Gradientennormen basiert, z. B. der Gradientennorm aus dem gesamten Modell, wie z. B. einige Varianten von LAMB Optimizer oder Global Gradient Clipping, 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 verschieben GPU (z. B. verwendenmodel.to(device)), wenn Sie die Bibliothek verwenden. Wenn Sie versuchen, das Modell GPU vor der Partitionierung des Modells (vor dem ersten smp.step Aufruf) zu verschieben, wird der Move-Aufruf ignoriert. Die Bibliothek verschiebt den Teil des Modells, dem eine Rangfolge zugewiesen wurde, automatisch in ihrenGPU. Sobald das Training mit der Bibliothek begonnen hat, sollten Sie das Modell nicht mehr dorthin verschieben CPU 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 Modellparallelitätsbibliothek trainiert wurde, wird empfohlen, das vollständige Modell mithilfe unseres Checkpoints zu speichern API 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 zu verwendenden Eingangstensoren smp.step müssen auf das aktuelle Gerät verschoben werden .to()API, was nach dem Aufruf erfolgen muss. torch.cuda.set_device(local_rank())

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

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 wurde, automatisch an seinenGPU.

@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 der systemeigenen Version verwenden PyTorch DDP, 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 tun torch.distributedAPI, 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 deren Bibliotheksimplementierungen durch smp.optimizers und smp.nn APIs verwenden.