Initialisation différée des paramètres

L'initialisation d'un grand modèle à des fins d'entraînement n'est pas toujours possible avec une mémoire GPU limitée. Pour résoudre ce problème de mémoire GPU insuffisante, vous pouvez initialiser le modèle sur la mémoire du processeur. Cependant, pour les modèles plus grands avec plus de 20 ou 40 milliards de paramètres, même la mémoire du processeur peut ne pas être suffisante. Dans ce cas, nous vous recommandons d'initialiser le modèle sur ce qui PyTorch appelle un méta-périphérique, ce qui permet de créer des tenseurs sans qu'aucune donnée ne leur soit attachée. Un tenseur sur un méta-dispositif n'a besoin que des informations de forme, ce qui permet de créer un grand modèle avec ses paramètres sur des méta-périphériques. Hugging Face Accelerate fournit le init_empty_weights gestionnaire de contexte qui permet de créer un tel modèle sur des méta-appareils tout en initialisant les tampons sur un appareil normal. Avant le début de l'entraînement, PyTorch FSDP initialise les paramètres du modèle. Cette fonctionnalité d'initialisation différée des paramètres de SMP v2 retarde la création des paramètres du modèle une fois que PyTorch FSDP a effectué le partitionnement des paramètres. PyTorch Le FSDP accepte une fonction d'initialisation des paramètres (param_init_fn) lors du partitionnement des modules, et il appelle param_init_fn chaque module. L'param_init_fnAPI prend un module comme argument et initialise tous les paramètres qu'il contient, à l'exclusion des paramètres d'un module enfant. Notez que ce comportement est différent de la PyTorch version native 2.0.1 qui présente un bogue entraînant l'initialisation des paramètres plusieurs fois.

SMP v2 fournit l'torch.sagemaker.delayed_param.DelayedParamIniterAPI permettant d'appliquer l'initialisation différée des paramètres.

Les extraits de code suivants montrent comment appliquer l'torch.sagemaker.delayed_param.DelayedParamIniterAPI à votre script d'entraînement.

Supposons que vous disposiez d'un script d'entraînement PyTorch FSDP comme suit.

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

Notez que l'approche d'initialisation différée des paramètres n'est pas indépendante du modèle. Pour résoudre ce problème, vous devez écrire une init_weights fonction, comme indiqué dans l'exemple précédent, afin qu'elle corresponde à l'initialisation de la définition du modèle d'origine et qu'elle couvre tous les paramètres du modèle. Pour simplifier le processus de préparation de cette init_weights fonction, SMP v2 implémente cette fonction d'initialisation pour les modèles suivants : GPT-2, GPT-J, GPT-Neox et Llama de Hugging Face Transformers. L'torch.sagemaker.delayed_param.DelayedParamIniterAPI fonctionne également avec l'implémentation parallèle du tenseur SMP, torch.sagemaker.tensor_parallel.transformer.TransformerLMHead model, que vous pouvez appeler après l'appel de l'torch.sagemaker.transformAPI.

À l'aide de l'torch.sagemaker.delayed_param.DelayedParamIniterAPI, vous pouvez adapter votre script PyTorch FSDP comme suit. Après avoir créé un modèle avec des poids vides, enregistrez l'torch.sagemaker.delayed_param.DelayedParamIniterAPI dans le modèle et définissez-en un objet. Passez l'objet à param_init_fn la classe PyTorch FSDP.

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()
    )

Remarques sur les poids liés

Lorsque nous entraînons des modèles avec des poids liés, nous devons faire particulièrement attention à lier les poids après les avoir initialisés avec une initialisation différée des paramètres. PyTorchLe FSDP ne dispose pas d'un mécanisme pour lier les poids après les avoir initialisés comme ci-dessus. param_init_fn Pour résoudre de tels cas, nous avons ajouté une API pour autoriser apost_init_hook_fn, qui peut être utilisée pour lier les poids. Vous pouvez y transmettre n'importe quelle fonction qui accepte le module comme argument, mais nous avons également une tie_weights méthode post_param_init_fn prédéfinie définie dans DelayedParamIniter laquelle appelle le module s'il existe. Notez qu'il est prudent de toujours le transmettre post_param_init_fn même s'il n'existe aucune tie_weights méthode pour le module.

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()
    )