Ativação e descarregamento

Importante

No SMP v2.2.0, a funcionalidade de descarregamento de ativação da biblioteca SMP não funciona. Em vez disso, use o descarregamento de PyTorch ativação nativo.

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

O SMP v2 melhora esse descarregamento de ativação. Ele pré-busca as ativações com antecedência, antes que elas sejam necessárias para que a GPU comece a repassar essas ativações. O recurso de pré-busca ajuda os progressos do treinamento a serem executados com mais eficiência sem GPUs ociosas. 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 ativação SMP 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:

• checkpoint_wrapper.py no PyTorch GitHubrepositório

• Ponto de verificação de ativação no PyTorch blog Scaling Multimodal Foundation Models in with Distributed. TorchMultimodal PyTorch

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

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

Configuração SMP

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

No roteiro de treinamento

nota

Ao ativar o recurso de descarregamento de ativação do SMP, certifique-se de também usar a PyTorch offload_wrapper função e aplicá-la ao módulo raiz. O recurso de descarregamento de ativação do SMP usa o módulo raiz para determinar quando o encaminhamento é feito para iniciar a pré-busca.

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)