Paralelismo especializado

As traduções são geradas por tradução automática. Em caso de conflito entre o conteúdo da tradução e da versão original em inglês, a versão em inglês prevalecerá.

As traduções são geradas por tradução automática. Em caso de conflito entre o conteúdo da tradução e da versão original em inglês, a versão em inglês prevalecerá.

Um modelo Mixture of Experts (MoE) é um modelo tipo transformador que utiliza uma abordagem esparsa, tornando-o mais leve para treinamento em comparação com o treinamento tradicional de modelos densos. Nessa arquitetura de rede neural de MoE, apenas um subconjunto dos componentes do modelo, chamados especialistas, é utilizado para cada entrada. Essa abordagem oferece várias vantagens, incluindo treinamento mais eficiente e inferência mais rápida, mesmo com um tamanho de modelo maior. Ou seja, com o mesmo orçamento computacional para treinar um modelo totalmente denso, você pode ajustar um modelo ou conjunto de dados maior ao usar o MoE.

Um modelo de MoE é composto por vários especialistas, cada um composto por uma rede neural, normalmente uma rede de feed-forward (FFN). Uma rede de portas chamada roteador determina quais tokens são enviados para qual especialista. Esses especialistas são especializados no processamento de aspectos específicos dos dados de entrada, o que permite que o modelo seja treinado mais rapidamente, reduza o custo de computação e, ao mesmo tempo, alcance a mesma qualidade de desempenho do modelo denso equivalente. Para saber mais sobre o Mixture of Experts em geral, consulte o blog Aplicação da combinação de especialistas em arquiteturas de LLM no site para desenvolvedores da NVIDIA.

O paralelismo especializado é um tipo de paralelismo que divide especialistas de um modelo de MoE em dispositivos de GPU.

O SMP v2 se integra ao NVIDIA Megatron para implementar paralelismo especializado para suportar modelos de treinamento de MoE e é executado com base no FSDP. PyTorch APIs Você continua usando seu código de treinamento PyTorch FSDP como está e ativa o paralelismo especializado em SMP para treinar modelos MoE.

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

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

Configurar o paralelismo especializado

Para expert_parallel_degree, você seleciona um valor para o grau de paralelismo especializado. 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 gradativamente 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 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 Use 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 o parâmetro expert_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 com o torch.sagemaker.init() para ativar o SMP v2 e envolver seu modelo com a API torch.sagemaker.transform, ao adicionar o parâmetro config à API para ativar a MoE. O trecho de código a seguir mostra como ativar o a MoE do SMP para a função de modelo genérico AutoModelForCausalLM com uso de uma configuração de modelo tipo transformador de MoE, pelo método from_config de treinamento do zero ou o método de ajuste from_pretrained. Para saber mais sobre a função MoEConfig do SMP, consulte torch.sagemaker.moe.moe_config.MoEConfig.

# Import the torch.sagemaker.transform API and initialize.
import torch.sagemaker as tsm
tsm.init()

# Import transformers AutoModelForCausalLM class.
from transformers import AutoModelForCausalLM

# Import the SMP-implementation of MoE configuration class.
from torch.sagemaker.moe.moe_config import MoEConfig

# Define a transformer model with an MoE model configuration
model = AutoModelForCausalLM.from_config(MoEModelConfig)

# Wrap it by torch.sagemaker.transform with the SMP MoE configuration.
model = tsm.transform(
    model, 
    config=MoEConfig(
        smp_moe=True,
        random_seed=12345,
        moe_load_balancing="sinkhorn",
        global_token_shuffle=False,
        moe_all_to_all_dispatcher=True,
        moe_aux_loss_coeff=0.001,
        moe_z_loss_coeff=0.001
    )
)

Configuração do SMP

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

{   
    ..., # other SMP config parameters
    "expert_parallel_degree": 8
}