Paralelismo de tensores

O paralelismo de tensores é um tipo de paralelismo de modelo no qual pesos, gradientes e estados do otimizador específicos do modelo são divididos entre dispositivos. Em contraste com o paralelismo de tubulação, que mantém os pesos individuais intactos, mas divide o conjunto de pesos, gradientes ou otimizador entre dispositivos, o paralelismo tensorial fragmenta os pesos individuais. Isso normalmente envolve computação distribuída de operações, módulos ou camadas específicas do modelo.

O paralelismo do tensor é necessário nos casos em que um único parâmetro consome a maior parte da memória da GPU (como grandes tabelas de incorporação com um grande tamanho de vocabulário ou uma grande camada softmax com um grande número de classes). Nesse caso, tratar esse grande tensor ou operação como uma unidade atômica é ineficiente e impede o equilíbrio da carga de memória.

O SMP v2 se integra ao Transformer Engine para a implementação do paralelismo de tensores e é executado com base nas APIs do FSDP. PyTorch Você pode ativar o paralelismo de tensores PyTorch FSDP e SMP simultaneamente e determinar o melhor paralelismo do modelo para obter o melhor desempenho.

Na prática, o paralelismo de tensores é especialmente útil nos cenários a seguir.

• Ao treinar com longos comprimentos de contexto, isso leva a uma alta memória de ativação apenas com o FSDP.

• Ao treinar com clusters realmente grandes nos quais o tamanho do lote global excede os limites desejados.

Modelos Hugging Face Transformer compatíveis com o paralelismo do tensor SMP

Atualmente, o SMP v2 oferece suporte a paralelismo de tensores para os seguintes modelos de transformadores Hugging Face.

• GPT-Neox

• Lhama 2

Para obter a configuração de referência para aplicar o paralelismo de tensores nesses modelos, consulte. Dicas de configuração

Configurar o paralelismo do tensor

Paratensor_parallel_degree, você seleciona um valor para o grau de paralelismo do tensor. O valor deve dividir uniformemente o número de GPUs em seu cluster. Por exemplo, para fragmentar seu modelo ao usar uma instância com 8 GPUs, escolha 2, 4 ou 8. Recomendamos que você comece com um número pequeno e aumente gradualmente até que o modelo caiba na memória da GPU.

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 tensor_parallel_degree e random_seed, consulte Parâmetros de configuração do recurso principal do SMP v2.

Configuração SMP

{
    "tensor_parallel_degree": 8,
    "random_seed": 0 
}

Em seu roteiro de treinamento

Inicialize com torch.sagemaker.init() para ativar o SMP v2 e agrupar seu modelo com a API. torch.sagemaker.transform

import torch.sagemaker as tsm
tsm.init()

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model)

Salvando e carregando os pontos de verificação do Hugging Face Transformer

Depois que a biblioteca SMP transforma um modelo, ela altera o dicionário de estado (state_dict) do modelo. Isso significa que o modelo se torna incompatível com as funcionalidades originais de checkpoint do Hugging Face Transformer. Para lidar com isso, a biblioteca SMP fornece APIs para salvar pontos de verificação de um modelo transformado na representação do Hugging Face Transformer e a API torch.sagemaker.transform para carregar um ponto de verificação do modelo Hugging Face Transformer para ajuste fino.

Para obter mais informações sobre como salvar pontos de verificação ao usar o recurso de paralelismo de tensores do SMP v2, consulte. Salve e carregue pontos de verificação ao usar o SMP

Para obter mais informações sobre o ajuste fino de um modelo aplicando o recurso de paralelismo de tensores do SMP v2, consulte. Ajuste fino