Descarga de activación

importante

En la versión 2.2.0 de SMP, la funcionalidad de activación y descarga de la biblioteca SMP no funciona. En su lugar, utilice la descarga de activaciones nativa. PyTorch

Normalmente, la pasada hacia delante calcula las activaciones en cada capa y las guarda en la memoria de la GPU hasta que finaliza la pasada hacia atrás de la capa correspondiente. Si se descargan estos tensores en la memoria de la CPU después de la transferencia 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 permanezcan inactivas mientras las activaciones se recuperan de la CPU durante la transferencia hacia atrás. Esto provoca una importante degradación del rendimiento cuando se utiliza la descarga de activaciones.

SMP v2 mejora esta descarga de activación. Obtiene previamente las activaciones antes de que sean necesarias para que la GPU comience a transferirlas hacia atrás. La función 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 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 activación y descarga de SMP en su script. Tenga en cuenta que la activación y la descarga solo se aplican cuando se utilizan 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 descarga, consulte:

• checkpoint_wrapper.py en el repositorio PyTorch GitHub

• Activation Checkpoints en el PyTorch blog Scaling Multimodal Foundation Models in TorchMultimodal with PyTorch Distributed.

Puede aplicar la función SMP de 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 configuración del SMP 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 inicialización de SMP torch.sagemaker.init() a su guion de formación y cómo configurar el diccionario de configuración de SMP en formato JSON para el lanzador de tareas de formación, siguiendo el proceso de dos pasos introducido en la presente versión. Comience con la biblioteca de paralelismo de SageMaker modelos v2 No es necesario realizar ningún cambio en el modelo ni en la configuración del FSDP. PyTorch PyTorch 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 funciones principales del SMP v2.

Configuración SMP

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

En el guion de entrenamiento

nota

Al activar la función de activación y descarga de SMP, asegúrese de utilizarla también y aplicarla al módulo raíz. PyTorch offload_wrapper La función de descarga y activación de SMP utiliza el módulo raíz para determinar cuándo se ha realizado la transferencia directa e iniciar la captura previa.

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)