Agrega Accelerate a tu código

Cada marco de entrenamiento distribuido tiene su propia forma de funcionar, lo que puede requerir escribir mucho código personalizado para adaptarlo a tu código de entrenamiento de PyTorch y a tu entorno de entrenamiento. Accelerate ofrece una forma amigable de interactuar con estos marcos de entrenamiento distribuido sin tener que aprender los detalles específicos de cada uno. Accelerate se encarga de esos detalles por ti, para que puedas centrarte en el código de entrenamiento y escalarlo a cualquier entorno de entrenamiento distribuido.

En este tutorial aprenderás a adaptar tu código existente de PyTorch con Accelerate y a ponerte en camino hacia el entrenamiento en sistemas distribuidos con facilidad. Comenzarás con un bucle de entrenamiento básico de PyTorch (se asume que todos los objetos de entrenamiento, como model y optimizer, ya están configurados) e irás integrando Accelerate en él de forma progresiva.

device = "cuda"
model.to(device)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()
    scheduler.step()

Acelerador

La clase Accelerator es la principal para adaptar tu código para que funcione con Accelerate. Conoce la configuración distribuida que estás utilizando, como el número de procesos distintos y el tipo de hardware. Esta clase también proporciona acceso a muchos de los métodos necesarios para habilitar que tu código de PyTorch funcione en cualquier entorno de entrenamiento distribuido y para gestionar y ejecutar procesos entre dispositivos.

Por eso siempre debes empezar importando y creando una instancia de Accelerator en tu script.

from accelerate import Accelerator

accelerator = Accelerator()

El Accelerator también sabe a qué dispositivo mover tus objetos de PyTorch, por lo que se recomienda dejar que Accelerate lo haga por ti.

- device = "cuda"
+ device = accelerator.device
  model.to(device)

Preparar objetos de PyTorch

A continuación, debe preparar sus objetos de PyTorch (modelo, optimizador, programador, etc.) para el entrenamiento distribuido. El método prepare() se encarga de colocar su modelo en el contenedor adecuado (como una sola GPU o varias GPUs) para su configuración de entrenamiento, adaptar el optimizador y el programador para usar AcceleratedOptimizer y AcceleratedScheduler de Accelerate, y crear un nuevo dataloader que pueda dividirse entre procesos.

Accelerate solo prepara objetos que heredan de sus respectivas clases de PyTorch, como torch.optim.Optimizer.

Los objetos de PyTorch se devuelven en el mismo orden en que se envían.

model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

Bucle de entrenamiento

Por último, elimina las llamadas a to(device) para las entradas y los objetivos en el bucle de entrenamiento, ya que las clases de DataLoader de Accelerate las colocan automáticamente en el dispositivo correcto. También debes reemplazar el paso habitual de backward() con el método backward() de Accelerate, que escala los gradientes por ti y utiliza el método backward() adecuado según tu configuración distribuida (por ejemplo, DeepSpeed o Megatron).

-   inputs = inputs.to(device)
-   targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
-   loss.backward()
+   accelerator.backward(loss)

Junte todo y su nuevo bucle de entrenamiento de Accelerate debería verse así:

from accelerate import Accelerator
accelerator = Accelerator()

device = accelerator.device
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

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

Características de entrenamiento

Accelerate ofrece funciones adicionales, como acumulación de gradientes, recorte de gradientes, entrenamiento con precisión mixta y más, que puedes añadir a tu script para mejorar tu ejecución de entrenamiento. Exploremos estas tres funciones.

Acumulación de gradientes

La acumulación de gradientes le permite entrenar con tamaños de lote más grandes acumulando los gradientes durante varios lotes antes de actualizar los pesos. Esto puede ser útil para sortear las limitaciones de memoria. Para habilitar esta función en Accelerate, especifique el parámetro gradient_accumulation_steps en la clase Accelerator y agregue el administrador de contexto accumulate() a su script.

+ accelerator = Accelerator(gradient_accumulation_steps=2)
  model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)

  for input, label in training_dataloader:
+     with accelerator.accumulate(model):
          predictions = model(input)
          loss = loss_function(predictions, label)
          accelerator.backward(loss)
          optimizer.step()
          scheduler.step()
          optimizer.zero_grad()

Recorte de gradientes

El recorte de gradientes es una técnica para evitar los “gradientes explosivos”, y Accelerate ofrece:

• clipgrad_value() para recortar los gradientes a un valor mínimo y máximo
• clipgrad_norm() para normalizar los gradientes a un valor determinado

Precisión mixta

La precisión mixta acelera el entrenamiento al usar un tipo de datos de menor precisión, como fp16 (precisión media), para calcular los gradientes. Para obtener el mejor rendimiento con Accelerate, la pérdida debe calcularse dentro de su modelo (como en los modelos de Transformers), ya que los cálculos fuera del modelo se realizan con precisión completa.

Establezca el tipo de precisión mixta que se usará en el Accelerator y, a continuación, utilice el gestor de contexto autocast() para convertir automáticamente los valores al tipo de datos especificado.

Accelerate habilita la precisión mixta automática, por lo que autocast() solo es necesario si hay otras operaciones de precisión mixta además de las realizadas en la pérdida por backward(), que ya gestiona la escalado.

+ accelerator = Accelerator(mixed_precision="fp16")
+ with accelerator.autocast():
      loss = complex_loss_function(outputs, target)

Guardar y cargar

Accelerate también puede guardar y cargar un modelo una vez que el entrenamiento está completo, o también puede guardar el estado del modelo y el optimizador, lo cual puede ser útil para reanudar el entrenamiento.

Modelo

Una vez que todos los procesos estén completos, desenrolle el modelo con el método unwrap_model() antes de guardarlo, ya que el método prepare() envolvió su modelo en la interfaz adecuada para el entrenamiento distribuido. Si no desenrolla el modelo, al guardar el diccionario de estado del modelo también se guardarán las capas adicionales potenciales del modelo más grande y no podrá volver a cargar los pesos en su modelo base.

Debes usar el método save_model() para desenrollar y guardar el diccionario de estado del modelo. Este método también puede guardar un modelo en puntos de control fragmentados o en el formato safetensors.

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory)

unwrapped_model = accelerator.unwrap_model(model)
path_to_checkpoint = os.path.join(save_directory,"pytorch_model.bin")
unwrapped_model.load_state_dict(torch.load(path_to_checkpoint))

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory, max_shard_size="1GB", safe_serialization=True)

load_checkpoint_in_model(unwrapped_model, save_directory, device_map={"":device})

Estado

Durante el entrenamiento, es posible que desee guardar el estado actual del modelo, el optimizador, los generadores aleatorios y, potencialmente, los programadores de tasa de aprendizaje, para que puedan restaurarse en el mismo script. Debe agregar los métodos save_state() y load_state() a su script para guardar y cargar estados.

Para personalizar aún más dónde y cómo se guardan los estados a través de save_state(), utilice la clase ProjectConfiguration. Por ejemplo, si automatic_checkpoint_naming está habilitado, cada checkpoint guardado se almacena en Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}.

Los demás elementos con estado que se vayan a almacenar deben registrarse con el método register_for_checkpointing() para que puedan guardarse y cargarse. Cada objeto que se pase a este método para su almacenamiento debe tener una función load_state_dict y una función state_dict.

Si tienes instalado torchdata>=0.8.0, además puedes pasar use_stateful_dataloader=True a tu DataLoaderConfiguration. Esto extiende las clases de DataLoader de Accelerate con una función load_state_dict y una state_dict, y hace que Accelerator.save_state y Accelerator.load_state también rastreen hasta dónde se ha leído en el conjunto de datos de entrenamiento al persistir el modelo.