Descarga de activación

Normalmente, la pasada hacia delante calcula las activaciones en cada capa y las guarda en la GPU memoria hasta que finaliza la pasada hacia atrás de la capa correspondiente. Al descargar estos tensores a la CPU memoria después de la pasada hacia adelante y volver a recuperarlos GPU cuando se necesiten, se puede ahorrar una cantidad considerable de memoria. GPU PyTorch admite la descarga de activaciones, pero la implementación hace que permanezca inactiva mientras se recuperan las activaciones durante la transferencia GPUs hacia atrás. CPU Esto provoca una importante degradación del rendimiento cuando se utiliza la descarga de activaciones.

SMPLa versión 2 mejora esta descarga de activación. Realiza una búsqueda previa de las activaciones antes de que sean necesarias para poder empezar y pasar esas activaciones GPU hacia atrás. La función de captura previa ayuda a que los avances del entrenamiento se ejecuten de manera más eficiente sin estar inactivos. GPUs El resultado es que ofrece ventajas al reducir el uso de memoria sin degradar el rendimiento.

Puedes conservar los PyTorch módulos nativos para descargar las activaciones en tu guion de entrenamiento. El siguiente es un ejemplo de estructura para aplicar la función de SMP activación y descarga en su script. Tenga en cuenta que la activación y descarga solo se aplica cuando se utiliza junto con. Puntos de control de activación Para obtener más información sobre las herramientas de PyTorch puntos de control nativas para la activación y la descarga, consulte:

Puede aplicar la función de SMP activación y descarga en los puntos de control de activación. PyTorch Esto se hace añadiendo los activation_loading_horizon parámetros sm_activation_offloading y al diccionario de SMP configuración durante. Paso 2: Inicie un trabajo de formación

En los siguientes fragmentos de código se muestra cómo añadir el módulo de SMP inicialización torch.sagemaker.init() a su script de formación y cómo configurar el diccionario de SMP configuración en el JSON formato adecuado para el lanzador de tareas de formación, siguiendo el proceso de dos pasos introducido en el presente documento. Utilice la biblioteca de paralelismo de SageMaker modelos v2 No es necesario realizar ningún cambio en el modelo o la configuración. PyTorch PyTorch FSDP Para obtener más información acerca de los parámetros sm_activation_offloading y activation_loading_horizon, consulte SMPv2: parámetros de configuración de las funciones principales.

SMPconfiguración

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

En el guion 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)