Aktivierungs-Prüfpunkte

Bei den Aktivierungs-Prüfpunkten (oder Steigungs-Prüfpunkten) handelt es sich um eine Technik zur Reduzierung der Speicherbelegung, indem Aktivierungen bestimmter Layers gelöscht und bei einem Rücklauf neu berechnet werden. Dadurch wird zusätzliche Datenverarbeitungszeit effektiv gegen eine geringere Speicherauslastung eingetauscht. Wenn ein Modul mit einem Prüfpunkt versehen wird, bleiben am Ende eines Vorwärtsdurchlaufs die Ein- und Ausgaben des Moduls im Speicher. Alle Tensoren, die zwischenzeitlich Teil der Berechnung innerhalb dieses Moduls gewesen wären, werden während des Vorwärtsdurchlaufs wieder freigegeben. Beim Rückwärtsdurchlauf von Modulen mit Prüfpunkten werden diese Tensoren neu berechnet. Zu diesem Zeitpunkt haben die Layers hinter diesem Prüfpunkt-Modul ihren Rückwärtsdurchlauf abgeschlossen, so dass die maximale Speichernutzung mit Prüfpunkten geringer sein kann.

Anmerkung

Diese Funktion ist PyTorch in der SageMaker Modellparallelitätsbibliothek v1.6.0 und höher verfügbar.

So werden Aktivierungs-Prüfpunkte verwendet

Mit smdistributed.modelparallel können Sie Aktivierungs-Prüfpunkte bei der Granularität eines Moduls verwenden. Für alle torch.nn Module außer torch.nn.Sequential können Sie Prüfpunkte für einen Modulbaum nur verwenden, wenn er aus Sicht der Pipeline-Parallelität innerhalb einer Partition liegt. Im Fall des torch.nn.Sequential Moduls muss jeder Modulbaum innerhalb des sequentiellen Moduls vollständig innerhalb einer Partition liegen, damit die Aktivierungs-Prüfpunkte funktionieren. Diese Einschränkungen sollten Sie berücksichtigen, wenn Sie die manuelle Partitionierung verwenden.

Wenn Sie die automatisierte Modellpartitionierung verwenden, finden Sie die Protokolle der Partitionierungszuweisungen, beginnend mit Partition assignments: in den Protokollen der Trainingsaufträge. Wenn ein Modul über mehrere Ränge partitioniert ist (z. B. mit einem abstammenden Element auf einem Rang und einem anderen auf einem anderen Rang), ignoriert die Bibliothek den Versuch, für das Modul einen Checkpoint zu setzen, und gibt eine Warnmeldung aus, dass für das Modul kein Prüfpunkt verwendet wird.

Anmerkung

Die SageMaker Modellparallelitätsbibliothek unterstützt sowohl überlappende als auch nicht allreduce überlappende Operationen in Kombination mit Checkpointing.

Anmerkung

PyTorchDie native Checkpointing-API ist nicht kompatibel mit. smdistributed.modelparallel

Beispiel 1: Der folgende Beispielcode zeigt, wie Sie Aktivierungsprüfpunkte verwenden, wenn Sie in Ihrem Skript eine Modelldefinition haben.

import torch.nn as nn
import torch.nn.functional as F

from smdistributed.modelparallel.torch.patches.checkpoint import checkpoint

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = F.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        # This call of fc1 will be checkpointed
        x = checkpoint(self.fc1, x)
        x = self.fc2(x)
        return F.log_softmax(x, 1)

Beispiel 2: Der folgende Beispielcode zeigt, wie Sie Aktivierungs-Prüfpunkte verwenden, wenn Ihr Skript ein sequentielles Modell enthält.

import torch.nn as nn
from smdistributed.modelparallel.torch.patches.checkpoint import checkpoint_sequential

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.seq = nn.Sequential(
            nn.Conv2d(1,20,5),
            nn.ReLU(),
            nn.Conv2d(20,64,5),
            nn.ReLU()
        )

    def forward(self, x):
        # This call of self.seq will be checkpointed
        x = checkpoint_sequential(self.seq, x)
        return F.log_softmax(x, 1)

Beispiel 3: Der folgende Beispielcode zeigt, wie Aktivierungsprüfpunkte verwendet werden, wenn Sie ein vorgefertigtes Modell aus einer Bibliothek importieren, z. B. Hugging PyTorch Face Transformers. Gehen Sie wie folgt vor, unabhängig davon, ob Sie sequentielle Module mit Prüfpunkten versehen oder nicht:

1. Umschließen Sie das Modell mit smp.DistributedModel().

2. Definieren Sie ein Objekt für sequenzielle Ebenen.

3. Umschließen Sie das sequentielle Layer-Objekt mit smp.set_activation_checkpointig().

import smdistributed.modelparallel.torch as smp
from transformers import AutoModelForCausalLM

smp.init()
model = AutoModelForCausalLM(*args, **kwargs)
model = smp.DistributedModel(model)

# Call set_activation_checkpointing API
transformer_layers = model.module.module.module.transformer.seq_layers
smp.set_activation_checkpointing(
    transformer_layers, pack_args_as_tuple=True, strategy='each')