Paralelismo experto

Un modelo de mezcla de expertos (MoE) es un tipo de modelo de transformador que emplea un enfoque disperso, lo que facilita el entrenamiento en comparación con los tradicionales modelos densos de entrenamiento. En esta arquitectura de red neuronal de MoE, solo se utiliza para cada entrada un subconjunto de los componentes del modelo, denominados expertos. Este enfoque ofrece varias ventajas, como un entrenamiento más eficiente y una inferencia más rápida, incluso con un tamaño de modelo mayor. Es decir, con el mismo presupuesto de computación para entrenar un modelo denso completo, puede adaptarse a un modelo o conjunto de datos mayor al utilizar MoE.

Un modelo de MoE consta de varios expertos, cada uno de los cuales se compone de una red neuronal, normalmente una red prealimentada (FFN). Una red de puertas llamada router determina qué tokens se envían a cada experto. Estos expertos se especializan en procesar aspectos específicos de los datos de entrada, lo que permite que el modelo se entrene más rápido, se reduzcan los costos de cómputo y, al mismo tiempo, se logre la misma calidad de rendimiento que su modelo denso homólogo. Para obtener más información sobre mezcla de expertos en general, consulte el blog Applying Mixture of Experts in LLM Architectures en el sitio web para desarrolladores de NVIDIA.

El paralelismo experto es un tipo de paralelismo que gestiona la división de expertos de un modelo de MoE entre dispositivos de GPU.

El SMP v2 se integra con NVIDIA Megatron para implementar un paralelismo experto que permita el entrenamiento de modelos MoE, y se ejecuta sobre el FSDP. PyTorch APIs Sigues usando tu código de entrenamiento del PyTorch FSDP tal cual y activas el paralelismo experto del SMP para entrenar los modelos MoE.

Modelos de transformador Hugging Face compatibles con el paralelismo experto de SMP

Actualmente, SMP v2 ofrece compatibilidad de paralelismo experto para los siguientes modelos de transformador Hugging Face.

• Mixtral

Configuración de paralelismo experto

Para expert_parallel_degree, seleccione un valor de grado de paralelismo experto. El valor debe dividir en partes iguales el número de miembros del clúster. GPUs Por ejemplo, para fragmentar el modelo mientras usas una instancia con 8 GPUs, elige 2, 4 u 8. Le recomendamos que comience con un número pequeño y vaya aumentado poco a poco hasta que el modelo se ajuste a la memoria de la GPU.

En los siguientes fragmentos de código, se muestra cómo añadir el torch.sagemaker.init() del módulo de inicialización de SMP al script de entrenamiento y cómo configurar el diccionario de configuración de SMP en formato JSON para el lanzador de trabajos de entrenamiento mientras se sigue el proceso de dos pasos descrito en Utilice la biblioteca de paralelismo de SageMaker modelos v2. No es necesario realizar ningún cambio en el PyTorch modelo ni en la configuración del PyTorch FSDP. Para obtener más información sobre el parámetro expert_parallel_degree, consulte Parámetros de configuración de las características esenciales de SMP v2.

nota

Puede utilizar el paralelismo experto con Paralelismo híbrido de datos particionados. Tenga en cuenta que el paralelismo experto no es compatible actualmente con el paralelismo de tensores.

nota

Esta función de entrenamiento experto en paralelismo está disponible en la siguiente combinación de bibliotecas de y la biblioteca: SageMaker PyTorch

• SMP v2.3.0 y versiones posteriores

• El SDK de SageMaker Python v2.214.4 y versiones posteriores

• PyTorch v2.2.0 y versiones posteriores

En el script de entrenamiento

Como parte del paso 1, inicialice el script con torch.sagemaker.init() para activar SMP v2 y encapsular el modelo con la torch.sagemaker.transform API, añadiendo añada el parámetro config a la API para activar MoE. En el siguiente fragmento de código se muestra cómo activar MoE de SMP para la clase de modelo genérico AutoModelForCausalLM extrayendo la configuración de un modelo de transformador de MoE con el método from_config para entrenamiento desde cero o el método from_pretrained para realizar el afinamiento. Para obtener más información sobre la clase MoEConfig de 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
    )
)

Configuración de SMP

Como parte del paso 2, añada el siguiente parámetro al diccionario de configuración SMP del estimador. SageMaker PyTorch

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