Entraînement FP16 avec parallélisme des modèles - Amazon SageMaker AI

Les traductions sont fournies par des outils de traduction automatique. En cas de conflit entre le contenu d'une traduction et celui de la version originale en anglais, la version anglaise prévaudra.

Entraînement FP16 avec parallélisme des modèles

Pour l'entraînement FP16, appliquez les modifications suivantes à votre script de formation et à votre estimateur.

Note

Cette fonctionnalité est disponible dans la bibliothèque de parallélisme de modèles SageMaker version v1.10.0 et ultérieure.

Adaptation du script d'entraînement PyTorch

  1. Enveloppez votre modèle en utilisant le gestionnaire de contexte smdistributed.modelparallel.torch.model_creation ().

    # fp16_training_script.py

import torch
import smdistributed.modelparallel.torch as smp

with smp.model_creation(
    dtype=torch.float16 if args.fp16 else torch.get_default_dtype()
):
    model = ...
    Astuce

    Si vous utilisez le parallélisme des tenseurs, ajoutez tensor_parallelism=smp.tp_size() > 1 au gestionnaire de contexte smp.model_creation. L'ajout de cette ligne aide également à détecter automatiquement si le parallélisme des tenseurs est activé ou non.

    with smp.model_creation(
    ... ,
    tensor_parallelism=smp.tp_size() > 1
):
    model = ...

  2. Lorsque vous enveloppez l'optimiseur avec smdistributed.modelparallel.torch.DistributedOptimizer, définissez l'argument static_loss_scaling ou dynamic_loss_scaling. Par défaut, static_loss_scaling a la valeur de 1.0, et dynamic_loss_scaling a la valeur False. Si vous définissez dynamic_loss_scale=True, vous pouvez introduire les options de mise à l'échelle dynamique des pertes sous forme de dictionnaire via l'argument dynamic_loss_args. Dans la plupart des cas, nous vous recommandons d'utiliser l'échelle dynamique de perte avec les options par défaut. Pour plus d'informations, d'options et d'exemples de la fonction d'encapsulation de l'optimiseur, consultez l'API smdistributed.modelparallel.torch.DistributedOptimizer.

    Le code suivant est un exemple d'encapsulation d'un objet optimiseur Adadelta avec mise à l'échelle dynamique des pertes pour l'entraînement FP16.

    optimizer = torch.optim.Adadelta(...)
optimizer = smp.DistributedOptimizer(
    optimizer,
    static_loss_scale=None,
    dynamic_loss_scale=True,
    dynamic_loss_args={
        "scale_window": 1000,
        "min_scale": 1,
        "delayed_shift": 2
    }
)

Configuration d'un estimateur PyTorch SageMaker

Ajoutez le paramètre FP16 ("fp16") à la configuration de distribution pour le parallélisme du modèle lors de la création d'un objet estimateur SageMaker PyTorch. Pour trouver une liste complète de paramètres de configuration pour le parallélisme de modèle, consultez les paramètres pour smdistributed.

from sagemaker.pytorch import PyTorch

smp_options = {
    "enabled": True,
    "parameters":  {
        "microbatches":  4,
        "pipeline_parallel_degree":  2,
        "tensor_parallel_degree":  2,
        ...,

        "fp16": True
    }
}

fp16_estimator = PyTorch(
    entry_point="fp16_training_script.py", # Specify your train script
    ...,

    distribution={
        "smdistributed": {"modelparallel": smp_options},
        "mpi": {...}
    }
)

fp16_estimator.fit(...)

Lorsque l’entraînement FP16 commence, le modèle et l'optimiseur sont enveloppés par FP16_Module et FP16_Optimizer respectivement, qui sont des versions smdistributed modifiées des utilitaires Apex. FP16_Module convertit le modèle en dtype FP16 et traite de la transmission vers l'avant dans FP16.

Astuce

Vous pouvez appliquer un écrêtage de gradient en appelant clip_master_grads avant optimizer.step.

optimizer.clip_master_grads(max_norm)     # max_norm(float or int): max norm of the gradients
Astuce

Lorsque vous utilisez torch.optim.lr_scheduler et l'entraînement FP16, vous devez transmettre optimizer.optimizer au planificateur LR plutôt qu'à l'optimiseur. Voici l'exemple de code suivant :

from torch.optim.lr_scheduler import StepLR

scheduler = StepLR(
    optimizer.optimizer if smp.state.cfg.fp16 else optimizer,
    step_size=1,
    gamma=args.gamma
)