Punto de control de gradientes

La pasada hacia adelante (forward pass) suele almacenar en caché todas las activaciones intermedias para reutilizarlas en la pasada hacia atrás (backward pass). Sin embargo, las activaciones escalan con el tamaño del lote (batch size) y la longitud de la secuencia. El checkpointing de gradientes solo guarda ciertas activaciones y descarta el resto. Esto obliga a la pasada hacia atrás a recalcular algunas de las activaciones sobre la marcha a medida que se necesitan.

Normal training:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save ALL activations)
  Backward:  ←uses cached activations everywhere

Gradient checkpointing:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save only at checkpoints, discard the rest)
  Backward:  ←reaches L2, recomputes L2→L3 from scratch, uses it, discards it

El entrenamiento suele ser más lento porque la pasada hacia atrás recalcula las activaciones descartadas, pero el checkpointing reduce la memoria de activaciones.

Establezca gradient_checkpointing=True para habilitarlo.

Úsalo con acumulación de gradientes para reducir aún más el uso de memoria.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
)

Punto de control parcial

El checkpointing completo de gradientes recalcula cada capa que admite checkpointing. Si tu ejecución tiene algo de memoria libre, realiza checkpointing de menos capas para intercambiar parte de los ahorros de memoria por velocidad.

Pasa every_n_layers a gradient_checkpointing_enable() para elegir el intervalo de checkpointing.

every_n_layers=2

Forward:   input  -> [L1] -> [L2] -> [L3] -> [L4] -> [L5] -> [L6]
                      CP     keep     CP     keep     CP     keep

Backward:  output <- [L6] <- [L5] <- [L4] <- [L3] <- [L2] <- [L1]
                     keep   rerun    keep   rerun    keep   rerun

Con every_n_layers=2, la primera capa y cada segunda capa a partir de ella se marcan para checkpoint. Las capas marcadas descartan sus activaciones durante la pasada hacia adelante y las recalculan durante la pasada hacia atrás, mientras que las demás capas mantienen sus activaciones en memoria.

model.gradient_checkpointing_enable(every_n_layers=2)

El valor predeterminado, every_n_layers=1, realiza un checkpoint de cada capa. Los valores más grandes realizan un checkpoint de la primera capa y luego de cada n capas posteriores, dejando las activaciones de las otras capas en memoria. Por ejemplo, every_n_layers=2 realiza un checkpoint de las capas 1, 3, 5, y así sucesivamente. Solo se cuentan los módulos que heredan de GradientCheckpointingLayer. Otros módulos que admiten el checkpoint de gradientes permanecen habilitados.

Para usar el checkpointing parcial de gradientes con Trainer, establezca every_n_layers en gradient_checkpointing_kwargs.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"every_n_layers": 4},
)

Descarga de las activaciones guardadas

El checkpointing de gradientes mantiene una activación por capa con checkpoint en la GPU. Con longitudes de secuencia largas, esto puede consumir una cantidad considerable de memoria (aproximadamente layers x sequence x hidden x bytes_per_element). Establezca offload para mantener esas activaciones en memoria del host fijada, a costa de una copia de dispositivo a host durante la pasada hacia adelante y una copia de host a dispositivo durante la pasada hacia atrás.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"offload": True},
)

Ambas copias se ejecutan en el stream de cómputo, por lo que esto intercambia un paso más lento por memoria. Úsalo cuando una ejecución no quepa de otra manera, no para acelerarla.

Próximos pasos

• Lea el documento Uso de la memoria de la GPU para comprender qué factores determinan el uso de la memoria en la GPU durante el entrenamiento.
• Consulte la guía de Entrenamiento con precisión mixta para aprender a usar tipos de datos de menor precisión y reducir aún más la memoria y acelerar el entrenamiento.
• Consulte la guía de Kernels para aprender a acelerar el entrenamiento con kernels fusionados personalizados.