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á.

Paralelismo especializado

Um modelo Mixture of Experts (MoE) é um tipo de modelo de transformador que emprega uma abordagem esparsa, tornando-o mais leve para treinamento em comparação com o treinamento de modelos densos tradicionais. Nessa arquitetura de rede neural 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. Em outras palavras, 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 MoE consiste em vários especialistas, cada um consistindo em uma rede neural, normalmente uma rede de feedback (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, permitindo 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 a mistura de especialistas em geral, consulte o blog Aplicando a mistura de especialistas em arquiteturas LLM no site para desenvolvedores da NVIDIA.

O paralelismo especializado é um tipo de paralelismo que divide especialistas de um modelo 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 nas APIs do FSDP. PyTorch Você continua usando seu código de treinamento PyTorch FSDP como está e ativa o paralelismo especializado em SMP para treinar modelos MoE.

Modelos Hugging Face Transformer compatíveis com o paralelismo especializado em SMP

O paralelismo especializado do SMP v2 suporta o seguinte modelo Hugging Face Transformer.

Configure o paralelismo especializado

Paraexpert_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 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 o parâmetro expert_parallel_degree, consulte Parâmetros de configuração do recurso principal do SMP v2.

Em seu roteiro de treinamento

Como parte da Etapa 1, inicialize seu script torch.sagemaker.init() para ativar o SMP v2 e encapsular seu modelo com a torch.sagemaker.transform API, adicionando o config parâmetro à API para ativar o MoE. O trecho de código a seguir mostra como ativar o SMP MoE para a classe de modelo genérico usando uma configuração de AutoModelForCausalLM modelo de transformador MoE usando o from_config método de treinamento do zero ou o método de ajuste fino. from_pretrained Para saber mais sobre a MoEConfig classe SMP, consultetorch.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 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
}