Verzögerte Parameterinitialisierung

Die Initialisierung eines großen Modells für das Training ist aufgrund des begrenzten GPU Speichers nicht immer möglich. Um dieses Problem des unzureichenden GPU Speichers zu lösen, können Sie das Modell im Arbeitsspeicher initialisieren. CPU Bei größeren Modellen mit mehr als 20 oder 40 Milliarden Parametern reicht jedoch möglicherweise selbst der CPU Arbeitsspeicher nicht aus. In einem solchen Fall empfehlen wir, dass Sie das Modell auf einem sogenannten Metagerät PyTorch initialisieren, das die Erstellung von Tensoren ermöglicht, ohne dass Daten an sie angehängt sind. Ein Tensor auf einem Metagerät benötigt nur die Forminformationen, was es ermöglicht, ein großes Modell mit seinen Parametern auf Metageräten zu erstellen. Hugging Face Accelerate bietet den Kontextmanagerinit_empty_weights, mit dem Sie ein solches Modell auf Metageräten erstellen und gleichzeitig die Puffer auf einem normalen Gerät initialisieren können. PyTorch FSDPInitialisiert die Modellparameter, bevor das Training beginnt. Diese Funktion zur verzögerten Parameterinitialisierung von SMP v2 verzögert die Erstellung von Modellparametern, sodass sie erst nach dem Parameter-Sharding erfolgt PyTorch FSDP. PyTorch FSDPakzeptiert beim Sharding der Module eine Parameterinitialisierungsfunktion (param_init_fn) und ruft param_init_fn jedes Modul auf. Das param_init_fn API nimmt ein Modul als Argument und initialisiert alle darin enthaltenen Parameter, ohne die Parameter eines untergeordneten Moduls. Beachten Sie, dass sich dieses Verhalten von der nativen Version PyTorch 2.0.1 unterscheidet, die einen Fehler aufweist, der dazu führt, dass die Parameter mehrfach initialisiert werden.

SMPv2 bietet die Möglichkeit, eine verzögerte torch.sagemaker.delayed_param.DelayedParamIniter API Parameterinitialisierung anzuwenden.

Die folgenden Codefragmente zeigen, wie Sie das auf Ihr Trainingsskript torch.sagemaker.delayed_param.DelayedParamIniter API anwenden können.

Gehen Sie davon aus, dass Sie ein PyTorch FSDP Trainingsskript wie folgt haben.

# Creation of model on meta device
from accelerate import init_empty_weights
with init_empty_weights():
    model = create_model()

# Define a param init fn, below is an example for Hugging Face GPTNeoX.
def init_weights(module):
    d = torch.cuda.current_device()
    # Note that below doesn't work if you have buffers in the model
    # buffers will need to reinitialized after this call
    module.to_empty(device=d, recurse=False)
    if isinstance(module, (nn.Linear, Conv1D)):
        module.weight.data.normal_(mean=0.0, std=args.initializer_range)
        if module.bias:
            module.bias.data.zero_()
    elif isinstance(module, nn.Embedding):
        module.weight.data.normal_(mean=0.0, std=args.initializer_range)
        if module.padding_idx:
            module.weight.data[module.padding_idx].zero_()
    elif isinstance(module, nn.LayerNorm):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)

# Changes to FSDP wrapper.
model = FSDP(
    model,
    ...,
    param_init_fn=init_weights
)

# At this point model is initialized and sharded for sharded data parallelism.

Beachten Sie, dass der Ansatz der verzögerten Parameterinitialisierung nicht modellunabhängig ist. Um dieses Problem zu lösen, müssen Sie, wie im vorherigen Beispiel gezeigt, eine init_weights Funktion schreiben, die der Initialisierung in der ursprünglichen Modelldefinition entspricht und alle Parameter des Modells abdecken sollte. Um diesen Prozess der Vorbereitung einer solchen init_weights Funktion zu vereinfachen, implementiert SMP v2 diese Initialisierungsfunktion für die folgenden Modelle: GPT -2, GPT -J, GPT -NeoX und Llama von Hugging Face Transformers. Das funktioniert torch.sagemaker.delayed_param.DelayedParamIniter API auch mit dem torch.sagemaker.tensor_parallel.transformer.TransformerLMHead Modell der parallel SMP Tensor-Implementierung, das Sie nach dem torch.sagemaker.transform API Aufruf aufrufen können.

Mit dem torch.sagemaker.delayed_param.DelayedParamIniter API können Sie Ihr PyTorch FSDP Skript wie folgt anpassen. Nachdem Sie ein Modell mit leeren Gewichten erstellt haben, registrieren Sie es für das Modell und definieren Sie ein Objekt daraus. torch.sagemaker.delayed_param.DelayedParamIniter API Übergeben Sie das Objekt an das param_init_fn der PyTorch FSDP Klasse.

from torch.sagemaker.delayed_param import DelayedParamIniter
from accelerate import init_empty_weights

with init_empty_weights():
    model = create_model()
    
delayed_initer = DelayedParamIniter(model)

with delayed_initer.validate_params_and_buffers_inited():
    model = FSDP(
        model,
        ...,
        param_init_fn=delayed_initer.get_param_init_fn()
    )

Hinweise zu Gewichtsgleichgewichten

Beim Training von Modellen mit gebündelten Gewichten müssen wir besonders darauf achten, die Gewichte nach der Initialisierung der Gewichte mit verzögerter Parameterinitialisierung zu verknüpfen. PyTorchFSDPverfügt nicht über einen Mechanismus, mit param_init_fn dem die Gewichte nach der Initialisierung wie oben beschrieben miteinander verknüpft werden können. Um solchen Fällen API entgegenzuwirken, haben wir die Option a hinzugefügtpost_init_hook_fn, mit der die Gewichte verknüpft werden können. Sie können dort jede Funktion übergeben, die das Modul als Argument akzeptiert, aber wir haben auch eine vordefinierte Funktion post_param_init_fn definiert, in der DelayedParamIniter die tie_weights Methode des Moduls aufgerufen wird, falls sie existiert. Beachten Sie, dass es sicher ist, immer etwas zu übergeben, post_param_init_fn auch wenn es keine tie_weights Methode für das Modul gibt.

with delayed_initer.validate_params_and_buffers_inited():
    model = FSDP(
        model,
        ...,
        param_init_fn=delayed_initer.get_param_init_fn(),
        post_param_init_fn=delayed_initer.get_post_param_init_fn()
    )