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 ein Modul mit einem Checkpoint versehen ist, bleiben am Ende eines Vorwärtsdurchlaufs nur die anfänglichen 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 während des Rückwärtsdurchlaufs 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-Checkpoint-Modul,. PyTorch apply_activation_checkpointing Im Folgenden finden Sie Beispiele für das Aktivierungs-Checkpointing des GPT-NeoX-Modells von Hugging Face.

Checkpointing der Transformer-Schichten des GPT-NeoX-Modells von Hugging Face

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)

Checkpointing jeder anderen Transformer-Schicht des GPT-NeoX-Modells von Hugging Face

# 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.