Précision Mixte Automatique avec PyTorch

Introduction

Les modèles profonds d’apprentissage automatique ont besoin de plus de puissance de calcul et de ressources de mémoire. La rapide formation des réseaux de neurones profonds a été réalisée grâce au développement de nouvelles techniques. Au lieu du format FP32 (nombres réels à précision complète), vous pouvez utiliser FP16 (nombres réels à précision半 complète), et les chercheurs ont découvert qu’utiliser les deux en même temps est une meilleure option.

La précision mixte permet une formation à précision semi-complète tout en conservant une grande partie de l’exactitude du réseau à précision complète. Le terme “technique de précision mixte” fait référence au fait que cette méthode utilise à la fois des représentations à précision complète et à précision semi-complète.

Dans cet aperçu de l’entraînement Automatique de Précision Mixte (AMP) avec PyTorch, nous expliquons comment fonctionne la technique, en passant pas à pas par le processus d’utilisation d’AMP, et discutons des applications plus avancées de cette technique avec des schémas de code pour que les utilisateurs puissent plus tard les intégrer à leur propre code.

Prérequis

Base de Connaissances de PyTorch: Connaissance de PyTorch, y compris ses concepts fondamentaux tels que les tenseurs, les modules et la boucle d’entraînement.

Compétences en Fondements de l’Apprentissage Automatique Profond: Concepts tels que les réseaux de neurones, la rétropropagation et l’optimisation.

Connaissance de l’entraînement à précision mixte : Connaissance des avantages et des inconvénients de l’entraînement à précision mixte, y compris la réduction de l’utilisation de la mémoire et la calculabilité accrue.

Accès à un Matériel Compatible : Une GPU supportant la précision mixte, comme les GPU NVIDIA avec des Cores Tensor (par exemple, les architectures Volta, Turing, Ampere).

Configuration de Python et CUDA : Un environnement Python fonctionnel avec PyTorch installé et CUDA configuré pour l’accélération par GPU.

Aperçu de la Précision Mixte

Comme la plupart des frameworks d’apprentissage profond sont normalement entraînés sur des données à virgule flottante de 32 bits (FP32). Par contre, FP32 n’est pas toujours nécessaire pour la réussite. Il est possible d’utiliser une virgule flottante de 16 bits pour certaines opérations, où FP32 consomme plus de temps et de mémoire.

Par conséquent, les ingénieurs NVIDIA ont développé une technique permettant l’entraînement à précision mixte en utilisant FP32 pour un petit nombre d’opérations tout en conservant la plupart de la structure de réseau en FP16.

  • Convertir le modèle pour utiliser le type de données float16 lorsque possible.
  • Garder les poids maîtres float32 pour accumuler les mises à jour de poids à chaque itération.
  • L’utilisation de l’escalade de perte pour préserver les petites valeurs de gradient.

Mixed-Precision dans PyTorch

Pour la formation mixte en précision, PyTorch offre déjà une quantité importante de fonctionnalités intégrées.
Les paramètres d’un module sont convertis en FP16 lorsque vous appelez la méthode .half(), et les données d’un tenseur sont converties en FP16 lorsque vous appelez .half(). Des opérations arithmétiques rapides en FP16 seront utilisées pour exécuter toutes les opérations sur ces modules ou tenseurs. Les bibliothèques mathématiques de NVIDIA (cuBLAS et cuDNN) sont bien supportées par PyTorch. Les données de la chaîne FP16 sont traitées à l’aide de Tensor Cores pour effectuer des GEMM et des convolutions. Pour employer les Tensor Cores dans cuBLAS, les dimensions d’une GEMM ([M, K] x [K, N] -> [M, N]) doivent être des multiples de 8.

Introducing Apex

Utilitaires de précision mixte d’ApexLes utilitaires de précision mixte d’Apex sont conçus pour accroître la vitesse de formation tout en conservant l’exactitude et la stabilité de la formation en précision simple. Apex peut effectuer des opérations en FP16 ou en FP32, gérer automatiquement la conversion des paramètres maîtres et automatiquement scaler les pertes.

Apex a été créé pour faciliter l’inclusion de la formation à précision mixte dans les modèles des chercheurs. Amp, abrégé de Automatic Mixed-Precision, est l’une des fonctionnalités d’Apex, une extension légère de PyTorch. Il suffit de quelques lignes supplémentaires sur leurs réseaux pour que les utilisateurs puissent profiter de la formation à précision mixte avec Amp. Apex a été lancé à CVPR 2018, et il convient de remarquer que la communauté PyTorch a montré un solide soutien pour Apex depuis son lancement.

En apportant uniquement des modifications mineures au modèle en cours d’exécution, Amp permet de ne pas s’occuper des types mixtes lors de la création ou de l’exécution de votre script. Les hypothèses d’Amp peuvent ne pas s’ajuster aussi bien dans les modèles qui utilisent PyTorch de manières inhabituelles, mais il y a des crochets pour ajuster ces hypothèses au besoin.

Amp offre tous les avantages de la formation à précision mixte sans la nécessité de scaler les pertes ou de convertir les types explicitement. Le site Web GitHub d’Apex contient des instructions pour le processus d’installation, et sa documentation API officielle peut être trouvée ici.

Comment fonctionnent les Amps

Amp utilise un paradigme de liste blanche/noire au niveau logique. Les opérations de tenseur dans PyTorch incluent des fonctions de réseau neuronal telles que torch.nn.functional.conv2d, des fonctions mathématiques simples telles que torch.log, et des méthodes de tenseur telles que torch.Tensor. add__. Dans ce univers, il existe trois catégories principales de fonctions :

  • Liste blanche : Fonctions qui pourraient tirer profit de la vitesse accrue des mathématiques en FP16. Des applications typiques incluent la multiplication de matrices et la convolution.
  • Liste noire : Les entrées doivent être en FP32 pour les fonctions où la précision de 16 bits peut ne pas être suffisante.
  • Tout le reste (toutes les fonctions qui restent) : Fonctions qui peuvent s’exécuter en FP16, mais le coût d’une conversion FP32 -> FP16 pour les exécuter en FP16 n’est pas rentable étant donné que l’accélération n’est pas significative.

La tâche d’Amp est aisée, du moins en théorie. Amp détermine si une fonction de PyTorch est présente sur la liste blanche, la liste noire ou ni l’une ni l’autre avant de l’appeler. Tous les arguments devraient être convertis en FP16 si sur la liste blanche ou en FP32 si sur la liste noire. Si ni l’un ni l’autre, il faut simplement s’assurer que tous les arguments sont du même type. Cette politique n’est pas aussi simple à appliquer en pratique que cela semble.

Utiliser Amp en conjonction avec un modèle PyTorch

Pour intégrer Amp dans un script PyTorch existant, suivez les étapes suivantes :

  • Utilisez la bibliothèque Apex pour importer Amp.
  • Initializez Amp afin qu’il puisse apporter les modifications nécessaires au modèle, à l’optimiseur et aux fonctions internes de PyTorch.
  • Notez où la rétropropagation (.backward()) est appliquée afin que Amp puisse simultanément scaler la perte et effacer l’état par itération.

Étape 1

Il n’y a qu’une seule ligne de code pour la première étape :

from apex import amp

Étape 2

Le modèle de réseau neural et l’optimiseur utilisés pour la formation doivent déjà être spécifiés pour terminer cette étape, qui est composée d’une seule ligne.

model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

Des paramètres supplémentaires permettent de régler les ajustements des tenseurs et des types d’opérations d’Amp. La fonction amp.initialize() accepte de nombreux paramètres, parmi lesquels nous ne spécifierons que trois :

  • models (torch.nn.Module ou liste de torch.nn.Modules) – Modèles à modifier/cast.
  • optimisateurs (facultatif, torch.optim.Optimizer ou liste de torch.optim.Optimizers) – Optimisateurs à modifier/cast. OBLIGATOIRE pour l’entraînement, facultatif pour l’inférence.
  • opt_level (str, facultatif, par défaut=“O1”) – Niveau d’optimisation de précision pure ou mixte. Les valeurs acceptées sont “O0”, “O1”, “O2” et “O3”, expliquées en détail ci-dessus. Il existe quatre niveaux d’optimisation :

O0 pour l’entraînement en FP32 : Il s’agit d’une opération sans effet. Il n’y a pas besoin de s’occuper de cela car votre modèle entrant devrait déjà être en FP32, et O0 pourrait aider à établir un basculement pour la précision.

O1 pour la Précision Mixte (recommandé pour l’utilisation typique) : Modifiez toutes les méthodes de Tensor et de Torch pour utiliser une méthode d’input avec liste blanche-noire. En FP16, les opérations de liste blanche telles que les opérations du noyau Tensor Core telles que les GEMMs et les convolutions sont exécutées. Par exemple, la softmax est une opération de liste noire qui nécessite de la précision FP32. A moins d’être explicitement indiqué autrement, O1 utilise également la mise à l’échelle dynamique des pertes.

O2 pour la Précision Mixte “Presque FP16” : O2 convertit les poids du modèle en FP16, patch le forward du modèle pour convertir les données d’entrée en FP16, conserve les batchnorms en FP32, maintient les poids maîtres en FP32, met à jour les param_groups de l’optimiseur de sorte que l’.step() agisse directement sur les poids en FP32 et implémente la mise à l’échelle dynamique des pertes (à moins d’être remplacé). Contrairement à O1, O2 ne patch pas les fonctions de Torch ou les méthodes de Tensor.

O3 pour le training en FP16 : O3 peut ne pas être aussi stable que O1 et O2 en termes de précision mixte véritable. Par conséquent, il pourrait être avantageux de définir une vitesse de base pour votre modèle, contre laquelle l’efficacité de O1 et O2 peut être évaluée.
L’override de propriété supplémentaire keep_batchnorm_fp32=True dans O3 pourrait vous aider à déterminer la « vitesse de la lumière » si votre modèle utilise la normalisation de batch, activant ainsi la normalisation de batch cudnn.

O0 et O3 ne sont pas véritablement en précision mixte, mais ils aident à définir des bornes de base pour la précision et la vitesse respectivement. Une implémentation de précision mixte est définie comme O1 et O2.
Vous pouvez essayer les deux et voir qui améliore le plus la performance et la précision pour votre modèle particulier.

Étape 3

Assurez-vous de déterminer où se produit la phase de recul dans votre code.
Certaines lignes de code ressemblent à ceci apparaîtront :

loss = criterion(…)
loss.backward()
optimizer.step()

Étape 4

En utilisant le gestionnaire de contexte Amp, vous pouvez activer le redimensionnement des pertes en simplement enveloppant la phase de recul :

loss = criterion(…)
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
optimizer.step()

C’est tout. Vous pouvez maintenant réexécuter votre script avec le training en précision mixte activé.

Capture d’appels de fonctions

PyTorch manque d’un objet de modèle statique ou de graphe pour s’accrocher et insérer les casts mentionnés ci-dessus, étant donné sa flexibilité et sa dynamique. En effectuant un « monkey patching » sur les fonctions requises, Amp peut intercepter et casté les paramètres dynamiquement.

Par exemple, vous pouvez utiliser le code ci-dessous pour s’assurer que les arguments de la méthode torch.nn.functional.linear sont toujours castés en fp16 :

orig_linear = torch.nn.functional.linear
def wrapped_linear(*args):
 casted_args = []
  for arg in args:
    if torch.is_tensor(arg) and torch.is_floating_point(arg):
      casted_args.append(torch.cast(arg, torch.float16))
    else:
      casted_args.append(arg)
  return orig_linear(*casted_args)
torch.nn.functional.linear = wrapped_linear

Bien que Amp puisse ajouter des raffinements pour rendre le code plus robuste, l’appel à Amp.init() entraîne effectivement l’insertion de patches de singe dans toutes les fonctions PyTorch pertinentes, ce qui permet de casté les arguments correctement au runtime.

Minimizing Casts

Chaque poids est casté une fois par itération de FP32 -> FP16 car Amp garde un cache interne de tous les casts de paramètres et les réutilise au besoin. A chaque itération, le gestionnaire de contexte pour le passage inverse indique à Amp quand nettoyer le cache.

Autocasting and Gradient Scaling Using PyTorch

“Entraînement automatique en précision mixte” fait référence à la combinaison de torch.cuda.amp.autocast et de torch.cuda.amp.GradScaler. En utilisant torch.cuda.amp.autocast, vous pouvez configurer l’autocast uniquement pour certaines parties. L’autocast sélectionne automatiquement la précision pour les opérations GPU pour optimiser l’efficacité tout en conservant l’exactitude.

Les instances de torch.cuda.amp.GradScaler facilitent la réalisation des étapes de scaling des gradients. Le scaling des gradients réduit l’underflow des gradients, ce qui aide les réseaux avec des gradients de float16 à atteindre une meilleure convergence.

Voici du code pour démontrer comment utiliser autocast() pour obtenir une précision mixte automatique dans PyTorch :

# Crée un modèle et un optimiseur en précision par défaut
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# Crée un GradScaler une fois au début de l'entraînement.
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # Exécute le passage en avant avec l'autocast.
        with autocast(device_type='cuda', dtype=torch.float16):
            output = model(input)
            loss = loss_fn(output, target)

        # Les opérations en arrière plan se déroulent dans le même type de données que celui choisi par l'autocast pour les opérations en avant.
        scaler.scale(loss).backward()

        # scaler.step() d'abord déséchele les gradients des paramètres assignés à l'optimiseur.
   
        scaler.step(optimizer)

        # Met à jour l'échelle pour l'itération suivante.
        scaler.update()

Si le passage en avant pour une opération spécifique a des entrées float16, alors le passage en arrière pour cette opération produit des gradientes float16, et float16 peut ne pas être capable d’expresser des gradientes avec de petites magnitudes.

Les mises à jour des paramètres associés seront perdues si ces valeurs sont vidées vers zéro (« arrondi inférieur »).

Le scaling des gradientes est une technique qui utilise un facteur d’échelle pour multiplier les pertes du réseau et effectue ensuite un passage en arrière sur la perte échelle pour éviter l’arrondi inférieur. Il est également nécessaire de scaler les gradientes en cours de propagation en arrière dans le réseau par ce même facteur. En conséquence, les valeurs des gradientes ont une magnitude plus grande, ce qui les empêche d’être vidées vers zéro.

Avant la mise à jour des paramètres, chaque gradient de paramètre (attribut .grad) devrait être désécalé pour que le facteur d’échelle ne perturbe pas la tension d’apprentissage. Les deux autocast et GradScaler peuvent être utilisés indépendamment car ils sont modulaires.

Travailler avec des Gradientes Non Scalées

Cliquer sur les Gradientes

Nous pouvons échellez tous les degrés de gradient en utilisant la méthode Scaler.scale(Loss).backward(). Les propriétés .grad des paramètres entre backward() et scaler.step(optimizer) doivent être déséchellez avant de les modifier ou d’y accéder. Si vous souhaitez limiter la norme globale (voir torch.nn.utils.clip_grad_norm_()) ou la magnitude maximale (voir torch.nn.utils.clip_grad_value_()) de votre ensemble de gradients à une certaine valeur inférieure ou égale à une certaine valeur (une limite imposée par l’utilisateur), vous pouvez utiliser une technique appelée “coupure des gradients.”

Une coupure sans déséchellement entraînerait une norme/magnitude maximale des gradients étant échelle, invalidant votre seuil demandé (qui était supposé être le seuil pour les gradients non échellés). Les gradients contenus par les paramètres donnés à l’optimiseur sont déséchellés par scaler.unscale (optimizer).
Vous pouvez désécheller les gradients d’autres paramètres qui ont été précédemment donnés à un autre optimiseur (comme optimizer1) en utilisant scaler.unscale (optimizer1). Nous pouvons illustrer ce concept en ajoutant deux lignes de code :

# Déséchelle les gradients des paramètres affectés par l'optimiseur en place
        scaler.unscale_(optimizer)
# Étant donné que les gradients des paramètres affectés par l'optimiseur sont déséchellés, ils sont couverts comme d'habitude : 
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

Travail avec des Gradients Scalés

Accumulation de Gradients

L’accumulation de gradients se base sur une idée extrêmement simple. Au lieu de mettre à jour les paramètres du modèle, il attend et accumule les gradients sur des lots successifs pour calculer la perte et les gradients.

Après un certain nombre de lots, les paramètres sont mis à jour en fonction de la gradient cumulative. Voici un extrait de code sur la manière d’utiliser l’accumulation de gradients en pytorch :

scaler = GradScaler()

for epoch in epochs:
    for i, (input, target) in enumerate(data):
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
            # Normalisez la perte 
            loss = loss / iters_to_accumulate

        # Accumule les gradients scalés.
        scaler.scale(loss).backward()
          # Mise à jour des poids
        if (i + 1) % iters_to_accumulate == 0:
            # Peut décaler ici si désiré 
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()
  • L’accumulation de gradient ajoute les gradients sur une taille de lot adéquate de batch_per_iter * iters_to_accumulate.
    La taille doit être réglée pour la taille effective de lot; il faut vérifier les valeurs inf/NaN des grades, sauter l’étape si des inf/NaN sont détectés et mettre à jour la taille à la granularité du lot effectif.
    Il est également crucial de maintenir les grads à une échelle scalaire et cohérente lorsque les grads pour un lot particulier sont additionnés.

Si les grads sont non réglés (ou le facteur de réglage change) avant que l’accumulation soit complète, le prochain passage arrière ajoutera les grads réglés aux grads non réglés (ou les grads réglés par un autre facteur) après quoi il est impossible de recupérer les grads non réglés accumulés. L’étape doit s’appliquer.

  • Vous pouvez dérégler les grads en utilisant unscale juste avant l’étape, après que tous les grads réglés pour l’étape à venir aient été accumulés.
    Pour s’assurer d’un lot complet et efficace, appelez simplement update à la fin de chaque itération où vous avez préalablement appelé la fonction step
  • enumerate(data) nous permet de suivre l’index du lot lorsque nous parcourons les données.
  • Divisez la perte en cours par iters_to_accumulate(loss / iters_to_accumulate). Cela réduit la contribution de chaque mini-lot que nous traitons en normalisant la perte. Si vous avez moyenne la perte à l’intérieur de chaque lot, la division est déjà correcte et aucune normalisation supplémentaire n’est requise. Cette étape peut ne pas être nécessaire en fonction de la manière dont vous calculez la perte.
  • Lorsque nous utilisons scaler.scale(loss).backward(), PyTorch accumule les gradientés échelle et les stocke jusqu’à ce que nous appelions optimizer.zero grad().

Pénalité de gradient

Lors de la mise en œuvre d’une pénalité de gradient, torch.autograd.grad() est utilisé pour construire les gradients, qui sont combinés pour former la valeur de la pénalité, puis ajoutés à la perte. Une pénalité L2 sans échelle ou sans automatisation est montrée dans l’exemple ci-dessous.

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)

        # Crée les gradients
        grad_prams = torch.autograd.grad(outputs=loss,
                                          inputs=model.parameters(),
                                          create_graph=True)

        # Calcule la composante de pénalité et l'ajoute à la perte
        grad_norm = 0
        for grad in grad_prams:
            grad_norm += grad.pow(2).sum()
        grad_norm = grad_norm.sqrt()
        loss = loss + grad_norm

        loss.backward()

        # Vous pouvez couper les gradients ici

        optimizer.step()

Les tenseurs fournis à torch.autograd.grad() devraient être échellés pour mettre en œuvre une pénalité de gradient. Il est nécessaire de désécheller les gradients avant de les combiner pour obtenir la valeur de la pénalité. Comme la computation de la composante de pénalité fait partie de l’appel avant, elle doit se produire à l’intérieur d’un contexte automatisé.
Pour la même pénalité L2, c’est comme suit :

scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # Effectuer l'échelle de perte pour la passée arrière de autograd.grad, ce qui entraîne #scaled_grad_prams
        scaled_grad_prams = torch.autograd.grad(outputs=scaler.scale(loss),
                                                 inputs=model.parameters(),
                                                 create_graph=True)

        # Crée grad_prams avant de calculer la pénalité (grad_prams doit être #déescalé). 
        # Comme aucun optimisateur n'appartient à scaled_grad_prams, la division conventionnelle est utilisée plutôt que scaler.unscale_ :
        inv_scaled = 1./scaler.get_scale()
        grad_prams = [p * inv_scaled for p in scaled_grad_prams]

        # La composante de pénalité est calculée et ajoutée à la perte. 
        with autocast():
            grad_norm = 0
            for grad in grad_prams:
                grad_norm += grad.pow(2).sum()
            grad_norm = grad_norm.sqrt()
            loss = loss + grad_norm

        # Applique l'échelle à l'appel arrière.
        # Accumule les gradients feuille correctement échellés.
        scaler.scale(loss).backward()

        # Vous pouvez déescaler ici 

        # step() et update() procèdent comme d'habitude.
        scaler.step(optimizer)
        scaler.update()

Travailler avec plusieurs modèles, pertes et optimisateurs

Scaler.scale doit être appelé sur chaque perte de votre réseau si vous en avez beaucoup.
Si vous avez beaucoup d’optimisateurs dans votre réseau, vous pouvez exécuter scaler.unscale sur l’un d’eux, et vous devez appeler scaler.step sur chacun d’eux. Cependant, scaler.update ne devrait être utilisé qu’une fois, après l’itération de tous les optimisateurs utilisés dans cette itération :

scaler = torch.cuda.amp.GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer1.zero_grad()
        optimizer2.zero_grad()
        with autocast():
            output1 = model1(input)
            output2 = model2(input)
            loss1 = loss_fn(2 * output1 + 3 * output2, target)
            loss2 = loss_fn(3 * output1 - 5 * output2, target)

       #Bien que le graphe de conservation soit indépendant de l'amp, il est présent dans cet exemple car les deux appels de backward() partagent certaines parties du graphe. 
        scaler.scale(loss1).backward(retain_graph=True)
        scaler.scale(loss2).backward()

        # Si vous souhaitez visualiser ou ajuster les gradients des paramètres qu'ils possèdent, vous pouvez spécifier quels optimisateurs reçoivent un désescalage explicit. .
        scaler.unscale_(optimizer1)

        scaler.step(optimizer1)
        scaler.step(optimizer2)

        scaler.update()

Chaque optimiseur examine ses gradients pour les inf et les NaN et rend une décision individuelle sur la nécessité de sauter l’itération. Certains optimiseurs peuvent sauter l’itération, tandis que d’autres ne le font pas. Le saut d’itération survient une fois par quelques centaines d’itérations ; par conséquent, il ne devrait pas affecter la convergence. Pour les modèles à optimiseurs multiples, vous pouvez signaler le problème si vous constatez une mauvaise convergence après l’ajout de l’échelle des gradients.

Travailler avec plusieurs GPUs

L’un des problèmes les plus importants avec les modèles d’apprentissage profond est qu’ils sont devenus trop importants pour être entraînés sur une seule GPU. Il peut prendre trop longtemps pour entraîner un modèle sur une seule GPU, et il est nécessaire d’utiliser l’entraînement Multi-GPU pour préparer les modèles le plus rapidement possible. Un chercheur bien connu a réussi à raccourcir la période d’entraînement d’ImageNet de deux semaines à 18 minutes ou à entraîner le plus grand et le plus avancé Transformer-XL en deux semaines au lieu de quatre ans.

DataParallel et DistributedDataParallel

Sans compromettre la qualité, PyTorch offre la meilleure combinaison de facilité d’utilisation et de contrôle. nn.DataParallel et nn.parallel.DistributedDataParallel sont deux fonctionnalités de PyTorch pour distribuer la formation sur plusieurs GPUs. Vous pouvez utiliser ces wrappers faciles à utiliser et les modifications pour entraîner le réseau sur plusieurs GPUs.

DataParallel dans un seul processus

Sur une seule machine, DataParallel aide à répandre l’entraînement sur de nombreuses GPUs.
Voyons plus précisément comment DataParallel fonctionne vraiment en pratique.
Lors de l’utilisation de DataParallel pour entraîner un réseau de neurones, les étapes suivantes se déroulent :

Source :

  • Le mini-batch est divisé sur GPU : 0.
  • Séparer et distribuer le mini-batch à toutes les GPUs disponibles.
  • Copier le modèle sur les GPUs.
  • Le passage avant se produit sur toutes les GPUs.
  • Calculer la perte en fonction des sorties du réseau sur GPU : 0, ainsi que de retourner les pertes aux différentes GPUs. Les gradients devraient être calculés sur chaque GPU.
  • Somme des gradients sur GPU : 0 et appliquer l’optimiseur pour mettre à jour le modèle.

Il convient de souligner que les préoccupations discutées ici ne s’appliquent qu’à autocast. Le comportement de GradScaler demeure inchangé. Il n’importe si torch.nn.DataParallel crée des threads pour chaque périphérique pour effectuer la phase de dérivation. L’état autocast est communiqué dans chacun d’eux, et ce qui suit fonctionnera :

model = Model_m()
p_model = nn.DataParallel(model)

# Définit l'autocast dans le thread principal
with autocast():
    # Il y aura un autocast dans p_model. 
    output = p_model(input)
    # loss_fn est également autocast
    loss = loss_fn(output)

DistributedDataParallel, un GPU par processus

La documentation pour torch.nn.parallel.DistributedDataParallel recommande d’utiliser une GPU par processus pour obtenir le meilleur rendement. Dans ce cas, DistributedDataParallel n’initialise pas de threads internes ; donc l’utilisation d’autocast et de GradScaler n’est pas affectée.

DistributedDataParallel, plusieurs GPUs par processus

Dans torch.nn.parallel.DistributedDataParallel, il est possible que le passage avant soit exécuté sur chaque périphérique par un thread secondaire, à l’image de torch.nn.DataParallel. La solution est la même : appliquer autocast en tant que partie de la méthode forward de votre modèle pour s’assurer qu’il est activé dans les threads secondaires.

Conclusion

Dans cet article, nous avons :

  • Introduit Apex.
  • Vu comment fonctionnent les Amps.
  • Vu comment effectuer l’échelle des gradientes, le clip des gradientes, l’accumulation des gradientes et les pénalités de gradient.
  • Vu comment travailler avec plusieurs modèles, pertes et optimisateurs.
  • Vu comment effectuer DataParallel dans un seul processus lorsque l’on travaille avec plusieurs GPU.

Références

https://developer.nvidia.com/blog/apex-pytorch-easy-mixed-precision-training/
https://nvidia.github.io/apex/amp.html
https://discuss.pytorch.org/t/accumulating-gradients/30020
https://towardsdatascience.com/how-to-scale-training-on-multiple-gpus-dae1041f49d2

Source:
https://www.digitalocean.com/community/tutorials/automatic-mixed-precision-using-pytorch