Paralelismo especializado

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 o paralelismo especializado para aceitar os modelos de treinamento de MoE e é executado com base nas APIs do FSDP do PyTorch. Você continua usando seu código de treinamento do FSDP do PyTorch como está e ativa o paralelismo especializado do SMP para treinar modelos de 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 o modelo ao usar uma instância com oito GPUs, escolha dois, quatro ou oito. 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 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 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 do SMP do estimador do PyTorch do SageMaker:

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