Descarga de activación

Normalmente, la pasada hacia adelante calcula las activaciones en cada capa y las guarda en la memoria de la GPU hasta que finalice la pasada hacia atrás de la capa correspondiente. Si se descargan estos tensores en la memoria de la CPU después de la pasada hacia adelante y se devuelven a la GPU cuando se necesitan, se puede ahorrar una cantidad considerable de memoria de la GPU. PyTorch admite la descarga de activaciones, pero la implementación hace que las GPU estén inactivas mientras las activaciones se recuperan de la CPU durante la pasada hacia atrás. Esto provoca una degradación importante del rendimiento cuando se utiliza la descarga de activaciones.

SMP v2 mejora esta descarga de activaciones. Realiza captura previa de las activaciones antes de que sean necesarias para que la GPU inicie una pasada hacia atrás sobre ellas. La característica de captura previa ayuda a que los avances del entrenamiento se ejecuten de manera más eficiente sin GPU inactivas. El resultado es que ofrece las ventajas de un menor uso de memoria sin que se degrade el rendimiento.

Puede conservar los módulos nativos de PyTorch para descargar activaciones en el script de entrenamiento. A continuación se ofrece un ejemplo de estructura para aplicar la característica de descarga de activaciones de SMP en el script. Tenga en cuenta que la descarga de activaciones solo es aplicable cuando se utiliza junto con Puntos de control de activación. Para obtener más información sobre las herramientas de puntos de comprobación nativas de PyTorch para descarga de activaciones, consulte:

Puede aplicar la característica de descarga de activaciones de SMP en puntos de comprobación de activación de PyTorch. Esto se hace añadiendo los parámetros sm_activation_offloading y activation_loading_horizon al diccionario de configuración de SMP durante Paso 2: inicialización de un trabajo de entrenamiento.

En los siguientes fragmentos de código, se muestra cómo añadir el torch.sagemaker.init() del módulo de inicialización de SMP al script de entrenamiento y cómo configurar el diccionario de configuración de SMP en formato JSON para el lanzador de trabajos de entrenamiento mientras se sigue el proceso de dos pasos descrito en Uso de la biblioteca de paralelismo de modelos de SageMaker v2. No necesita hacer cambios en el modelo de PyTorch o en la configuración del PyTorch FSDP. Para obtener más información acerca de los parámetros sm_activation_offloading y activation_loading_horizon, consulte Parámetros de configuración de las características esenciales de SMP v2.

Configuración de SMP

{
    "activation_loading_horizon": 2,
    "sm_activation_offloading": True
}

En script de entrenamiento

import torch.sagemaker as tsm
tsm.init()

# Native PyTorch module for activation offloading
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing, 
    offload_wrapper,
)

model = FSDP(...)

# Activation offloading requires activation checkpointing.
apply_activation_checkpointing(
    model,
    check_fn=checkpoint_transformer_layers_policy,
)

model = offload_wrapper(model)