Entraînement sur des TPU

L’entraînement sur TPU peut différer légèrement de l’entraînement sur plusieurs GPU, même avec Accelerate. Ce guide a pour objectif de vous montrer où vous devez faire attention et pourquoi, ainsi que les bonnes pratiques en général.

Entraînement dans un Notebook

Le principal point de vigilance lors de l’entraînement sur des TPU provient du notebook_launcher(). Comme indiqué dans le tutoriel notebook, vous devez restructurer votre code d’entraînement en une fonction qui peut être transmise à la fonction notebook_launcher() et veiller à ne déclarer aucun tenseur sur le GPU.

Bien que sur un TPU cette dernière partie soit moins importante, un point crucial à comprendre est que lorsque vous lancez du code depuis un notebook, vous le faites via un processus appelé forking. Lors du lancement depuis la ligne de commande, vous effectuez un spawning, où un processus python n'est pas actuellement en cours d'exécution et vous spawnez un nouveau processus. Étant donné que votre notebook Jupyter utilise déjà un processus python, vous devez forker un nouveau processus à partir de celui-ci pour lancer votre code.

C’est à ce niveau que cela devient important, en ce qui concerne la déclaration de votre modèle. Sur les processus TPU forkés, il est recommandé d’instantier votre modèle une seule fois et de le passer à votre fonction d’entraînement. Cela diffère de l’entraînement sur GPU, où vous créez n modèles dont les gradients sont synchronisés et rétropropagés à certains moments. À la place, une seule instance de modèle est partagée entre tous les nœuds et est transmise de l’un à l’autre. Cela est particulièrement important lors de l’entraînement sur des TPU à ressources limitées, tels que ceux fournis dans les noyaux Kaggle ou sur Google Colaboratory.

Voici un exemple de fonction d’entraînement passée à notebook_launcher() en cas d’entraînement sur des CPU ou des GPU :

Cet extrait de code est basé sur celui du notebook simple_nlp_example trouvé ici, avec de légères modifications pour la simplicité

def training_function():
    # Initialize accelerator
    accelerator = Accelerator()
    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
    train_dataloader, eval_dataloader = create_dataloaders(
        train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
    )

    # Instantiate optimizer
    optimizer = AdamW(params=model.parameters(), lr=hyperparameters["learning_rate"])

    # Prepare everything
    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
    # prepare method.
    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader
    )

    num_epochs = hyperparameters["num_epochs"]
    # Now we train the model
    for epoch in range(num_epochs):
        model.train()
        for step, batch in enumerate(train_dataloader):
            outputs = model(**batch)
            loss = outputs.loss
            accelerator.backward(loss)

            optimizer.step()
            optimizer.zero_grad()

from accelerate import notebook_launcher

notebook_launcher(training_function)

Le notebook_launcher utilisera 8 processus par défaut si Accelerate a été configuré pour un TPU

Si vous utilisez cet exemple et que vous déclarez le modèle à l’intérieur de la boucle d’entraînement, vous risquez de rencontrer une erreur de ce type sur un système à ressources limitées :

ProcessExitedException: process 0 terminated with signal SIGSEGV

Cette erreur est extrêmement cryptique, mais l'explication de base est que vous avez épuisé la RAM système. Vous pouvez l'éviter entièrement en reconfigurant la fonction d'entraînement pour qu'elle accepte un seul argument model, et en le déclarant dans une cellule externe :

# In another Jupyter cell
model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)

+ def training_function(model):
      # Initialize accelerator
      accelerator = Accelerator()
-     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
      train_dataloader, eval_dataloader = create_dataloaders(
          train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
      )
  ...

Et enfin, en appelant la fonction d'entraînement avec :

  from accelerate import notebook_launcher
- notebook_launcher(training_function)
+ notebook_launcher(training_function, (model,))

Le contournement ci-dessus n’est nécessaire que lors du lancement d’une instance TPU depuis un Jupyter Notebook sur un serveur à ressources limitées, tel que Google Colaboratory ou Kaggle. Si vous utilisez un script ou que vous lancez l’instance sur un serveur beaucoup plus puissant, il n’est pas nécessaire de déclarer le modèle au préalable.

Précision mixte et variables globales

Comme indiqué dans la section sur la précision mixte du tutoriel de migration, Accelerate prend en charge fp16 et bf16, les deux pouvant être utilisés sur des TPU. Cela dit, il est idéal d’utiliser bf16, car il est extrêmement efficace.

Il existe deux « niveaux » lors de l’utilisation de bf16 et d’Accelerate sur les TPU : au niveau de base et au niveau des opérations.

Au niveau de base, cette fonctionnalité est activée en passant mixed_precision="bf16" à Accelerator, par exemple :

accelerator = Accelerator(mixed_precision="bf16")

Par défaut, cela convertit torch.float et torch.double en bfloat16 sur les TPU. La configuration spécifique appliquée consiste à définir la variable d’environnement XLA_USE_BF16 à 1.

Il existe une configuration supplémentaire que vous pouvez effectuer, à savoir la définition de la variable d’environnement XLA_DOWNCAST_BF16. Si elle est définie sur 1, alors torch.float est bfloat16 et torch.double est float32.

Cela est effectué dans l'objet Accelerator lors du passage de downcast_bf16=True :

accelerator = Accelerator(mixed_precision="bf16", downcast_bf16=True)

Utiliser le downcasting plutôt que le bf16 partout est judicieux lorsque vous calculez des métriques, journalisez des valeurs, etc., car les tenseurs bf16 bruts seraient inutilisables dans ces contextes.

Temps d'entraînement sur TPU

Lorsque vous lancez votre script, vous remarquerez peut-être que l’entraînement semble exceptionnellement lent au début. Cela est dû au fait que les TPU parcourent d’abord quelques lots de données pour déterminer la quantité de mémoire à allouer avant d’utiliser cette allocation de mémoire configurée de manière extrêmement efficace.

Si vous constatez que votre code d’évaluation servant à calculer les métriques de votre modèle prend plus de temps en raison de l’utilisation d’une taille de lot plus grande, il est recommandé de conserver la même taille de lot que pour les données d’entraînement si le processus est trop lent. Sinon, la mémoire sera réallouée à cette nouvelle taille de lot après les premières itérations.

Le fait que la mémoire soit allouée ne signifie pas qu'elle sera utilisée ni que la taille du lot augmentera lors du retour à votre chargeur de données d'entraînement.