Paralelismo de contexto

O paralelismo de contexto é um tipo de paralelismo de modelo que divide as ativações do modelo ao longo da dimensão da sequência. Ao contrário de outras técnicas de paralelismo de sequência, que apenas dividem LayerNorm e RMSNorm, o paralelismo de contexto divide as entradas da rede e todas as ativações intermediárias ao longo da dimensão da sequência.

O SMP v2 integra-se ao Transformer Engine para paralelismo de contexto e pode ser usado com o FSDP do PyTorch e o SMP Paralelismo de tensores. Você pode ativar os três paralelismos simultaneamente para o treinamento do modelo. O paralelismo de contexto é benéfico para modelos de treinamento com grandes tamanhos de ativação e longos comprimentos de sequência. Ele acelera o cálculo das pontuações e saídas de atenção, permitindo que cada dispositivo calcule apenas uma parte das pontuações e saídas ao longo da dimensão da sequência. Embora o paralelismo de tensores também acelere a computação por meio da divisão pela dimensão oculta, a vantagem do paralelismo de contexto é mais substancial, pois os requisitos computacionais aumentam de maneira quadrática com a dimensão da sequência.

Modelos tipo transformador da Hugging Face compatíveis com o paralelismo de contexto no SMP

Atualmente, o SMP v2 é compatível com paralelismo de contexto para os seguintes modelo tipo transformador da Hugging Face:

Configurar paralelismo de contexto

Defina um valor inteiro para o parâmetro context_parallel_degree, que divide uniformemente o número de GPUs em seu cluster. Por exemplo, se você tiver uma instância de oito GPUs, use dois, quatro ou oito para context_parallel_degree. Recomendamos começar com um valor pequeno de context_parallel_degree e aumentá-lo gradativamente até que o modelo caiba na memória da GPU com o comprimento de sequência de entrada necessário.

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

No script de treinamento

Como parte da Etapa 1, inicialize seu script torch.sagemaker.init() para ativar o SMP v2 e envolver seu modelo com a API torch.sagemaker.transform.

A partir do SMP v2.6.0, você pode usar o argumento cp_comm_type para determinar qual implementação de paralelismo de contexto usar. Atualmente, a biblioteca de SMP é compatível com duas implementações: p2p e all_gather. A implementação p2p usa chamadas de envio e recebimento ponto a ponto para acúmulo de valores-chave durante a implementação da atenção e é executada de forma assíncrona, permitindo sobreposições com a computação. Já a implementação all_gather usa a operação coletiva AllGather, e é executada de forma síncrona.

import torch.sagemaker as tsm
tsm.init()

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model, cp_comm_type="p2p")

Configuração do SMP

Como parte da Etapa 2, adicione o seguinte parâmetro ao dicionário de configuração do SMP do estimador do PyTorch do SageMaker:

{   
    ..., # other SMP config parameters
    "context_parallel_degree": 2
}