Initialisation différée des paramètres

L'initialisation d'un grand modèle pour l'entraînement n'est pas toujours possible avec une GPU mémoire limitée. Pour résoudre ce problème de GPU mémoire insuffisante, vous pouvez initialiser le modèle en CPU mémoire. Cependant, pour les modèles plus grands comportant plus de 20 ou 40 milliards de paramètres, même CPU la mémoire 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 la SMP v2 retarde la création des paramètres du modèle après avoir effectué le PyTorch FSDP sharding des paramètres. PyTorch FSDPaccepte une fonction d'initialisation des paramètres (param_init_fn) lors du sharding des modules, et elle appelle param_init_fn chaque module. param_init_fnAPIprend 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.

SMPv2 permet d'appliquer torch.sagemaker.delayed_param.DelayedParamIniter API une initialisation différée des paramètres.

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

Supposons que vous disposiez d'un script d' PyTorch FSDPentraînement 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, la 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. Cela fonctionne torch.sagemaker.delayed_param.DelayedParamIniter API également avec l'implémentation parallèle du SMP tenseur, torch.sagemaker.tensor_parallel.transformer.TransformerLMHead model, que vous pouvez appeler après l'torch.sagemaker.transformAPIappel.

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

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 avoir initialisé les poids avec une initialisation différée des paramètres. PyTorchFSDPn'a pas de mécanisme pour lier les poids après les avoir initialisés param_init_fn comme ci-dessus. Pour résoudre de tels cas, nous avons ajouté API unpost_init_hook_fn, qui peut être utilisé 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()
    )