Ativação e descarregamento

Normalmente, a passagem para frente calcula as ativações em cada camada e as mantém na GPU memória até que a passagem para trás da camada correspondente termine. Descarregar esses tensores para a CPU memória após passar para frente e recuperá-los GPU quando forem necessários pode economizar um uso substancial de memória. GPU PyTorch suporta o descarregamento de ativações, mas a implementação faz com que fique ociosa enquanto as ativações são recuperadas durante a passagem GPUs para trás. CPU Isso causa uma grande degradação do desempenho ao usar o descarregamento de ativação.

SMPv2 melhora esse descarregamento de ativação. Ele pré-busca as ativações com antecedência, antes que as ativações sejam necessárias para começar a repassar GPU essas ativações. O recurso de pré-busca ajuda os progressos do treinamento a serem executados com mais eficiência sem inatividade. GPUs Isso resulta na oferta de benefícios de menor uso de memória sem degradação do desempenho.

Você pode manter os PyTorch módulos nativos para descarregar as ativações em seu script de treinamento. Veja a seguir um exemplo de estrutura de aplicação do recurso de descarregamento de SMP ativação em seu script. Observe que o descarregamento de ativação é aplicável somente quando usado em conjunto com. Ponto de verificação de ativação Para saber mais sobre as ferramentas nativas de PyTorch ponto de verificação para descarregamento de ativação, consulte:

Você pode aplicar o recurso de descarregamento de SMP ativação no ponto de verificação de PyTorch ativação. Isso é feito adicionando os activation_loading_horizon parâmetros sm_activation_offloading e ao dicionário SMP de configuração duranteEtapa 2: iniciar um trabalho de treinamento.

Os trechos de código a seguir mostram como adicionar o módulo de SMP inicialização torch.sagemaker.init() ao seu script de treinamento e configurar o dicionário de SMP configuração no JSON formato para o lançador de tarefas de treinamento, seguindo o processo de duas etapas apresentado em. Use a biblioteca de paralelismo de SageMaker modelos v2 Você não precisa fazer nenhuma alteração em seu PyTorch modelo ou PyTorch FSDPconfiguração. Para obter mais informações sobre os parâmetros sm_activation_offloading e activation_loading_horizon, consulte SMPparâmetros de configuração do recurso principal v2.

SMPconfiguração

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

No roteiro de treinamento

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)