Verzögerte Parameterinitialisierung

Die Initialisierung eines großen Modells für das Training ist mit dem begrenzten GPU-Speicher nicht immer möglich. Um dieses Problem des unzureichenden GPU-Speichers zu beheben, können Sie das Modell im CPU-Speicher initialisieren. Bei größeren Modellen mit mehr als 20 oder 40 Milliarden Parametern reicht jedoch möglicherweise nicht einmal der CPU-Speicher aus. In einem solchen Fall empfehlen wir, das Modell auf einem sogenannten PyTorch Metagerät zu initialisieren, das die Erstellung von Tensoren ermöglicht, ohne dass Daten an sie angehängt werden. 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. Bevor das Training beginnt, initialisiert PyTorch FSDP die Modellparameter. Diese Funktion zur verzögerten Parameterinitialisierung von SMP v2 verzögert die Erstellung von Modellparametern, sodass sie erst erfolgt, nachdem PyTorch FSDP das Parameter-Sharding durchgeführt hat. PyTorch FSDP akzeptiert beim Sharding der Module eine Parameterinitialisierungsfunktion (param_init_fn) und ruft jedes Modul auf. param_init_fn Die param_init_fn API verwendet 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.

SMP v2 stellt die torch.sagemaker.delayed_param.DelayedParamIniter API für die Anwendung der verzögerten Parameterinitialisierung bereit.

Die folgenden Codefragmente zeigen, wie Sie die torch.sagemaker.delayed_param.DelayedParamIniter API auf Ihr Trainingsskript anwenden.

Gehen Sie wie folgt davon aus, dass Sie über ein PyTorch FSDP-Trainingsskript verfügen.

# 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. Die torch.sagemaker.delayed_param.DelayedParamIniter API funktioniert 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.

Mithilfe der 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 die torch.sagemaker.delayed_param.DelayedParamIniter API für das Modell und definieren Sie ein Objekt daraus. Ü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. PyTorchFSDP verfügt nicht über einen Mechanismus, mit dem die Gewichte nach der Initialisierung wie oben beschrieben verknüpft werden können. param_init_fn Um solche Fälle zu lösen, haben wir eine API hinzugefügt, die eine erlaubtpost_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()
    )