Checkpointing bei der Aktivierung

Beim Aktivierungs-Checkpointing handelt es sich um eine Technik zur Reduzierung der Speicherbelegung, indem Aktivierungen bestimmter Layers gelöscht und bei dem Rücklauf neu berechnet werden. Dadurch wird zusätzliche Datenverarbeitungszeit effektiv gegen eine verringerte Speicherauslastung eingetauscht. Wenn bei einem Modul ein Checkpoint gesetzt wird, bleiben am Ende eines Vorwärtsdurchlaufs nur die ersten Eingänge des Moduls und die letzten Ausgänge des Moduls im Speicher. PyTorch gibt während des Vorwärtsdurchlaufs alle Zwischentensoren frei, die Teil der Berechnung innerhalb dieses Moduls sind. Berechnet diese Tensoren beim Rückwärtsdurchlauf der Checkpoint-Module neu. PyTorch Zu diesem Zeitpunkt haben die Ebenen hinter diesem Prüfpunktmodul ihren Rückwärtsdurchlauf abgeschlossen, sodass die maximale Speichernutzung mit Checkpointing geringer wird.

SMP v2 unterstützt das Aktivierungs-Checkpointing-Modul,. PyTorch apply_activation_checkpointing Im Folgenden finden Sie Beispiele für Aktivierungs-Checkpoints des Modells Hugging Face. GPT-NeoX

Überprüfen Sie die Transformer-Ebenen des Modells „Hugging Face“ GPT-NeoX

from transformers.models.gpt_neox import GPTNeoXLayer
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing
)
    
# check_fn receives a module as the arg, 
# and it needs to return whether the module is to be checkpointed
def is_transformer_layer(module):
    from transformers.models.gpt_neox import GPTNeoXLayer
    return isinstance(submodule, GPTNeoXLayer)
    
apply_activation_checkpointing(model, check_fn=is_transformer_layer)

Überprüfen Sie alle anderen Transformer-Ebenen des Hugging Face-Modells GPT-NeoX

# check_fn receives a module as arg, 
# and it needs to return whether the module is to be checkpointed
# here we define that function based on global variable (transformer_layers)
from transformers.models.gpt_neox import GPTNeoXLayer
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing
)

transformer_layers = [
    m for m model.modules() if isinstance(m, GPTNeoXLayer)
]

def is_odd_transformer_layer(module):
    return transformer_layers.index(module) % 2 == 0
    
apply_activation_checkpointing(model, check_fn=is_odd_transformer_layer)

Alternativ gibt es PyTorch auch das torch.utils.checkpoint Modul für Checkpointing, das von einer Untergruppe der Hugging Face Transformers-Modelle verwendet wird. Dieses Modul funktioniert auch mit SMP v2. Sie benötigen jedoch Zugriff auf die Modelldefinition, um den Prüfpunkt-Wrapper hinzufügen zu können. Daher empfehlen wir, die apply_activation_checkpointing-Methode zu verwenden.