Parallelizzazione del contesto

La parallelizzazione del contesto è un tipo di parallelizzazione del modello che partiziona le attivazioni del modello lungo la dimensione della sequenza. A differenza di altre tecniche di parallelizzazione delle sequenze, che si limitano a partizionare LayerNorm e RMSNorm, la parallelizzazione del contesto partiziona gli input di rete e tutte le attivazioni intermedie lungo la dimensione della sequenza.

SMP v2 si integra con Transformer Engine per la parallelizzazione del contesto e può essere utilizzato insieme a PyTorch FSDP e SMP Parallelizzazione tensoriale. Per l’addestramento dei modelli è possibile abilitare tutti e tre i parallelismi contemporaneamente. La parallelizzazione del contesto è utile per addestrare modelli con grandi dimensioni di attivazione e maggiori lunghezze di sequenza. Accelera il calcolo dei punteggi e degli output di attenzione, consentendo a ciascun dispositivo di calcolare solo una parte dei punteggi e degli output lungo la dimensione della sequenza. Sebbene la parallelizzazione tensoriale acceleri anche il calcolo attraverso il partizionamento lungo la dimensione nascosta, il vantaggio della parallelizzazione del contesto è più sostanziale, poiché i requisiti di calcolo aumentano quadraticamente con la dimensione della sequenza.

Modelli di Hugging Face Transformers compatibili con la parallelizzazione del contesto SMP

SMP v2 offre attualmente il supporto per la parallelizzazione del contesto per i seguenti modelli di Hugging Face Transformers.

Configurazione della parallelizzazione del contesto

Imposta un valore intero per il parametro context_parallel_degree che divide equamente il numero di GPU nel cluster. Ad esempio, se disponi di un’istanza con 8 GPU, utilizza 2, 4 o 8 per context_parallel_degree. È consigliabile iniziare con un valore context_parallel_degree basso e aumentarlo gradualmente fino a quando il modello risulta adatto alla memoria della GPU con la lunghezza della sequenza di input richiesta.

I seguenti frammenti di codice mostrano come aggiungere il modulo di inizializzazione SMP torch.sagemaker.init() allo script di addestramento e impostare il dizionario di configurazione di SMP in formato JSON per il programma di avvio del job di addestramento seguendo il processo in due fasi introdotto in Utilizzo della libreria SageMaker Model Parallelism v2. Non è necessario apportare alcuna modifica al modello PyTorch o alla configurazione PyTorch FSDP. Per ulteriori informazioni sul parametro context_parallel_degree, vedi Parametri di configurazione delle funzionalità principali di SMP v2.

Nello script di addestramento

Durante la Fase 1, inizializza lo script con torch.sagemaker.init() per attivare SMP v2 ed eseguire il wrapping del modello con l’API torch.sagemaker.transform.

A partire da SMP v2.6.0, è possibile utilizzare l’argomento cp_comm_type per determinare quale implementazione di contesto utilizzare. La libreria SMP attualmente supporta due implementazioni: p2p e all_gather. L’implementazione p2p utilizza chiamate di invio e ricezione peer-to-peer per l’accumulo di chiave-valore durante l’implementazione dell’attenzione e viene eseguita in modo asincrono, consentendo sovrapposizioni con il calcolo. L’implementazione all_gather, invece, utilizza l’operazione collettiva AllGather e viene eseguita in modo sincrono.

import torch.sagemaker as tsm
tsm.init()

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model, cp_comm_type="p2p")

Configurazione di SMP

Durante la Fase 2, aggiungi al dizionario di configurazione di SMP il parametro riportato di seguito per lo strumento di stima SageMaker PyTorch.

{   
    ..., # other SMP config parameters
    "context_parallel_degree": 2
}