Sincronización de gradientes

El módulo distribuido de PyTorch funciona comunicándose de ida y vuelta entre todas las GPU de su sistema. Esta comunicación toma tiempo, y garantizar que todos los procesos conozcan los estados de los demás ocurre en puntos de activación específicos al usar el módulo ddp.

Estos puntos de activación se añaden al modelo de PyTorch, concretamente a sus métodos forward() y backward(). Esto ocurre cuando el modelo se envuelve con DistributedDataParallel:

import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel

model = nn.Linear(10, 10)
ddp_model = DistributedDataParallel(model)

En Accelerate, esta conversión se realiza automáticamente al llamar a prepare() y pasar tu modelo.

+ from accelerate import Accelerator
+ accelerator = Accelerator()
  import torch.nn as nn
- from torch.nn.parallel import DistributedDataParallel

  model = nn.Linear(10,10)
+ model = accelerator.prepare(model)

La ralentización en la acumulación de gradientes

Ahora comprende que PyTorch añade hooks al método forward y al método backward de su modelo de PyTorch cuando se entrena en una configuración distribuida. Pero, ¿cómo puede esto hacer que su código se ejecute más lento?

En DDP (paralelismo de datos distribuido), se espera que el orden específico en el que se realizan y ejecutan los procesos ocurra en puntos específicos y que estos también se produzcan aproximadamente al mismo tiempo antes de continuar.

El ejemplo más directo es cuando se actualizan los parámetros del modelo a través de optimizer.step(). Sin acumulación de gradientes, todas las instancias del modelo deben haber calculado, agrupado y actualizado sus gradientes antes de pasar al siguiente lote de datos. Al realizar la acumulación de gradientes, se acumulan n gradientes de pérdida y se omite optimizer.step() hasta que se hayan alcanzado n lotes. Como todos los procesos de entrenamiento solo necesitan sincronizarse en el momento en que se llama a optimizer.step(), sin ninguna modificación en el paso de entrenamiento, esta comunicación innecesaria entre procesos puede causar un ralentización significativa.

¿Cómo puedes evitar esta sobrecarga?

Solucionar el problema de la ralentización

Dado que se omiten las actualizaciones de los parámetros del modelo al entrenar con estos lotes, no es necesario sincronizar sus gradientes hasta el punto en que se llama realmente a optimizer.step(). PyTorch no puede determinar automáticamente cuándo es necesario hacer esto, pero proporciona una herramienta para ayudar mediante el gestor de contexto no_sync que se añade a su modelo después de convertirlo a DDP.

En este gestor de contexto, PyTorch omitirá la sincronización de los gradientes cuando se llame a .backward(), y la primera llamada a .backward() fuera de este gestor de contexto activará la sincronización. A continuación se muestra un ejemplo:

ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for index, batch in enumerate(dataloader):
    inputs, targets = batch
    # Trigger gradient synchronization on the last batch
    if index != (len(dataloader) - 1):
        with ddp_model.no_sync():
            # Gradients only accumulate
            outputs = ddp_model(inputs)
            loss = loss_func(outputs)
            accelerator.backward(loss)
    else:
        # Gradients finally sync
        outputs = ddp_model(inputs)
        loss = loss_func(outputs)
        accelerator.backward(loss)
        optimizer.step()

En Accelerate, para que esto sea una API que se puede llamar sin importar el dispositivo de entrenamiento (aunque es posible que no haga nada si no estás en un sistema distribuido), ddp_model.no_sync se reemplaza por no_sync() y funciona de la misma manera:

  ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

  for index, batch in enumerate(dataloader):
      inputs, targets = batch
      # Trigger gradient synchronization on the last batch
      if index != (len(dataloader)-1):
-         with ddp_model.no_sync():
+         with accelerator.no_sync(model):
              # Gradients only accumulate
              outputs = ddp_model(inputs)
              loss = loss_func(outputs, targets)
              accelerator.backward(loss)
      else:
          # Gradients finally sync
          outputs = ddp_model(inputs)
          loss = loss_func(outputs)
          accelerator.backward(loss)
          optimizer.step()
          optimizer.zero_grad()

Como puede esperar, la función accumulate() encapsula esta comprobación condicional mediante el seguimiento del número de lote actual, dejándole con la API final de acumulación de gradientes:

ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for batch in dataloader:
    with accelerator.accumulate(model):
        optimizer.zero_grad()
        inputs, targets = batch
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

Como resultado, debe usar accelerator.accumulate o accelerator.no_sync cuando se trata de la elección de la API.

Cuánto es realmente el ralentización y los errores fáciles que puedes cometer

Para configurar un ejemplo realista, considere la siguiente configuración:

• Dos nodos de GPU única T4 y un nodo con dos GPUs
• Cada GPU es una T4 y están alojadas en GCP
• El script utilizado es una modificación del script NLP Example
• El tamaño de lote por GPU es 16, y los gradientes se acumulan cada 4 pasos

Todos los scripts están disponibles en este repositorio.

Si no se presta atención a la sincronización de gradientes y a la comunicación entre GPUs, se puede desperdiciar una gran cantidad de tiempo mientras estas GPUs se comunican entre sí durante periodos innecesarios.

¿En cuánto?

Referencia:

• Línea base: no utiliza las prácticas de sincronización discutidas aquí
• no_sync de forma incorrecta: no_sync solo alrededor de la llamada a backward, no de forward
• no_sync: uso correcto del patrón no_sync
• accumulate: uso correcto de accumulate()

A continuación se muestran los segundos promedio por lote al iterar sobre 29 lotes de datos para cada configuración, tanto en un nodo único como en la configuración de doble nodo:

Como puede ver, si no tiene cuidado con la forma en que configura la sincronización de gradientes, ¡puede experimentar un aumento de más del 2x en el tiempo de entrenamiento!

Si te preocupa asegurarte de que todo se haga correctamente, te recomendamos encarecidamente utilizar la función accumulate() y pasar gradient_accumulation_steps o gradient_accumulation_plugin al objeto Accelerator para que Accelerate lo gestione por ti.

no_sync requiere memoria de GPU adicional al usar FSDP

Tenga en cuenta que no sincronizar los gradientes puede tener efectos adversos durante el entrenamiento con FSDP. Como se ha advertido en torch, el administrador de contexto no_sync para FSDP requerirá memoria adicional.

Por lo tanto, en situaciones intensivas en memoria al usar FSDP, recomendamos establecer sync_each_batch en True en GradientAccumulationPlugin para desactivar no_sync.

Vea el ejemplo siguiente, donde ajustamos Mixtral (47B parámetros) en 8 GPUs A100-80GB. Observamos que, incluso para un modesto gradient_accumulation_steps=2, rápidamente agotamos la memoria (OOM) si no_sync está habilitado. Nuevamente, esto se debe a los sobrecostos adicionales de memoria derivados de no_sync de FSDP. Sin embargo, si no_sync se deshabilita mediante sync_each_batch=True, el consumo de memoria para gradient_accumulation_steps=16 vuelve al de gradient_accumulation_steps=1.

Desactivar no_sync implica que habrá un ralentizamiento debido a las sincronizaciones de datos adicionales, como se explicó en las secciones anteriores de esta guía.