Paralelismo experto

Un modelo de mezcla de expertos (MoE) es un tipo de modelo de transformador que emplea un enfoque disperso, lo que lo hace más ligero a la hora de entrenar en comparación con el entrenamiento de los modelos densos tradicionales. En esta arquitectura de red neuronal del MoE, solo se utiliza un subconjunto de los componentes del modelo, denominados expertos, para cada entrada. 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 más grande. En otras palabras, con el mismo presupuesto de cómputo para entrenar un modelo totalmente denso, puede adaptarse a un modelo o conjunto de datos más grande al utilizar el MoE.

Un modelo MoE consta de varios expertos, cada uno de los cuales consiste en una red neuronal, normalmente una red de retroalimentación (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, reduzca los costos de cómputo y, al mismo tiempo, logre la misma calidad de rendimiento que su modelo denso homólogo. Para obtener más información sobre Mixture of Experts en general, consulta 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 permite dividir a los expertos de un modelo MoE entre dispositivos de GPU.

El SMP v2 se integra con NVIDIA Megatron para implementar un paralelismo experto que permita el entrenamiento de modelos de MoE, y se ejecuta sobre las API del FSDP. PyTorch Sigue usando el código de entrenamiento del PyTorch FSDP tal como está y activa el paralelismo experto del SMP para entrenar los modelos MoE.

Modelos Hugging Face Transformer compatibles con el paralelismo SMP Expert

El paralelismo experto SMP v2 es compatible con el siguiente modelo de Hugging Face Transformer.

• Mixtral

Configure el paralelismo experto

Paraexpert_parallel_degree, selecciona un valor para el grado de paralelismo experto. El valor debe dividir en partes iguales el número de GPU del clúster. Por ejemplo, para fragmentar el modelo mientras usas una instancia con 8 GPU, elige 2, 4 u 8. Te recomendamos empezar con un número pequeño y aumentarlo gradualmente hasta que el modelo quepa en la memoria de la GPU.

En los siguientes fragmentos de código se muestra cómo añadir el módulo de inicialización de SMP torch.sagemaker.init() a su guion de formación y cómo configurar el diccionario de configuración de SMP en formato JSON para el lanzador de tareas de entrenamiento, siguiendo el proceso de dos pasos descrito en la presente descripción. Comience con la biblioteca de paralelismo de SageMaker modelos v2 No es necesario realizar ningún cambio en el modelo ni en la configuración del FSDP. PyTorch PyTorch Para obtener más información sobre el parámetro expert_parallel_degree, consulte Parámetros de configuración de las funciones principales del SMP v2.

nota

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

nota

Esta función de formación experta 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 tu guion de entrenamiento

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

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

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