Expertenparallelität

Ein Mixture-of-Experts-Modell (MoE) ist eine Art Transformer-Modell, das einen Ansatz mit geringer Dichte verwendet, sodass es leichter zu trainieren ist als das Training mit herkömmlichen Modellen mit hoher Dichte. In der neuronalen Netzwerkarchitektur von MoE wird für jede Eingabe nur eine Teilmenge der Modellkomponenten, die als Experten bezeichnet werden, verwendet. Dieser Ansatz bietet mehrere Vorteile, darunter effizienteres Training und schnellere Inferenz, selbst bei einer größeren Modellgröße. Mit anderen Worten: Mit dem gleichen Budget für das Training eines Modells mit hoher Dichte können Sie ein größeres Modell oder einen größeren Datensatz einfügen, wenn Sie MoE verwenden.

Ein MoE-Modell besteht aus mehreren Experten, die jeweils aus einem neuronalen Netzwerk bestehen, in der Regel einem Feed-Forward-Netzwerk (FFN). Ein Gate-Netzwerk, das als Router bezeichnet wird, bestimmt, welche Token an welchen Experten gesendet werden. Diese Experten sind auf die Verarbeitung bestimmter Aspekte der Eingabedaten spezialisiert, sodass das Modell schneller trainiert werden kann, die Datenverarbeitungskosten reduziert werden und gleichzeitig dieselbe Leistungsqualität wie bei dem Modell mit hoher Dichte erreicht wird. Weitere Informationen zu Mixture of Experts im Allgemeinen finden Sie im Blog Anwendung von Mixture of Experts in LLM-Architekturen auf der NVIDIA-Entwickler-Website.

Expertenparallelität ist eine Art von Parallelität, bei der Experten eines MoE-Modells auf verschiedene GPU-Geräte aufgeteilt werden.

SMP v2 ist in NVIDIA Megatron integriert, um die Expertenparallelität zur Unterstützung des Trainings von MoE-Modellen zu implementieren, und wird auf den PyTorch-FSDP-APIs ausgeführt. Sie verwenden Ihren PyTorch-FSDP-Trainingscode unverändert und aktivieren die SMP-Expertenparallelität für das Training von MoE-Modellen.

Transformer-Modelle von Hugging Face, die mit der SMP-Expertenparallelität kompatibel sind

SMP v2 bietet derzeit Unterstützung der Expertenparallelität für die folgenden Transformer-Modelle von Hugging Face.

Expertenparallelität konfigurieren

Für expert_parallel_degree wählen Sie einen Wert für den Grad der Expertenparallelität aus. Der Wert muss die Anzahl der GPUs in Ihrem Cluster gleichmäßig verteilen. Wählen Sie z. B. 2, 4 oder 8 aus, um Ihr Modell zu fragmentieren, wenn Sie eine Instance mit 8 GPUs verwenden. Wir empfehlen, mit einer kleinen Zahl zu beginnen und diese schrittweise zu erhöhen, bis das Modell in den GPU-Speicher passt.

Die folgenden Codefragmente zeigen, wie Sie das SMP-Initialisierungsmodul torch.sagemaker.init() zu Ihrem Trainingsskript hinzufügen und das SMP-Konfigurationswörterbuch im JSON-Format für den Trainingsjob-Launcher einrichten. Dabei folgen Sie dem in Verwendung der SageMaker-Modellparallelitätsbibliothek v2 beschriebenen zweistufigen Prozess. Sie müssen keine Änderungen an Ihrem PyTorch-Modell oder der PyTorch-FSDP-Konfiguration vornehmen. Weitere Informationen zum Parameter expert_parallel_degree erhalten Sie unter Konfigurationsparameter für die Kernfunktionen von SMP v2.

In Ihrem Trainingsskript

Initialisieren Sie im Rahmen von Schritt 1 Ihr Skript mit torch.sagemaker.init(), um SMP v2 zu aktivieren, und umschließen Sie Ihr Modell mit der torch.sagemaker.transform-API, wobei Sie der API den config-Parameter hinzufügen, um MoE zu aktivieren. Der folgende Codeausschnitt zeigt, wie Sie SMP MoE für die generische Modellklasse AutoModelForCausalLM aktivieren, indem Sie mit der from_config-Methode für das Training von Grund auf oder mit der from_pretrained-Methode für die Optimierung eine MoE-Transformer-Modellkonfiguration abrufen. Weitere Informationen zur SMP-MoEConfig-Klasse finden Sie unter 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
    )
)

SMP-Konfiguration

Fügen Sie im Rahmen von Schritt 2 den folgenden Parameter zum SMP-Konfigurationswörterbuch für den SageMaker-PyTorch-Schätzer hinzu.

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