Addestra modelli giganteschi con ridimensionamento quasi lineare utilizzando il parallelismo dei dati partizionati su Amazon SageMaker

Nella ricerca di una precisione superiore, i modelli di deep learning in aree come l'elaborazione del linguaggio naturale e la visione artificiale sono cresciuti in modo significativo negli ultimi anni, spesso contati da decine a centinaia di miliardi di parametri. La formazione di questi giganteschi modelli è impegnativa e richiede complesse strategie di distribuzione. I data scientist e gli ingegneri dell'apprendimento automatico sono costantemente alla ricerca del modo migliore per ottimizzare il proprio calcolo di addestramento, ma sono alle prese con il sovraccarico di comunicazione che può aumentare insieme alle dimensioni complessive del cluster.

Questo è il motivo per cui abbiamo lanciato di recente sparallelismo dei dati duro on Amazon Sage Maker, una nuova tecnica di allenamento distribuito a risparmio di memoria nel Libreria SageMaker modello parallelo (SMP).. Il parallelismo dei dati frammentati è stato creato appositamente per modelli su scala estrema e utilizza Amazon internamente MiCS tecnologia sotto il cofano, uno sforzo scientifico per ridurre al minimo la scala di comunicazione riducendo le spese generali di comunicazione costose radicate nella raccolta di parametri e nella sincronizzazione del gradiente. Con un modello GPT-30 con parametri 2B con lunghezza della sequenza 2048, questa nuova funzionalità ha raggiunto 141 TFLOP, una velocità superiore del 39.7% rispetto a DeepSpeed ​​ZeRO-3. Per un modello 10B GPT-2 con lunghezza di sequenza 512, questa nuova funzionalità ha raggiunto anche 564 campioni al secondo, una velocità superiore del 13.9% rispetto a Full Sharded Data Parallel (FSDP) di PyTorch. Ricorda che in un gigantesco addestramento del modello, ogni percentuale di accelerazione si traduce in dollari risparmiati e produttività guadagnata nel tuo team.

In questo post del blog, daremo prima un'occhiata più da vicino ai principali fattori di differenziazione del parallelismo dei dati partizionati e quando utilizzarlo. Quindi, imparerai come addestrare facilmente un modello GPT-30 con parametro 2B su SageMaker con questa nuova funzionalità. Infine confronteremo le prestazioni con altre opzioni open source, in particolare superando DeepSpeed ​​ZeRO fino al 39.7% su 256 GPU.

Come funziona il parallelismo dei dati partizionati e quando utilizzarlo

Prima di introdurre il parallelismo dei dati partizionati, esaminiamo la sua famiglia di tecniche più ampia. I recenti approcci di addestramento distribuito per modelli di grandi dimensioni sono passati a un paradigma in cui i parametri del modello, i gradienti e gli stati dell'ottimizzatore sono condivisi tra nodi paralleli di dati. A differenza di Pipeline Parallelism, che ha l'innata complessità di scegliere i livelli da partizionare tra i dispositivi, specialmente quando il tuo framework non supporta scomposizione automatizzata del modello, questo paradigma preserva elegantemente la semplicità del parallelismo dei dati, rimuovendo il vincolo del parallelismo dei dati in cui un modello deve adattarsi a una singola GPU.

Nei framework esistenti che rientrano in questo paradigma, in particolare DeepSpeed ​​ZeRO-3 e FSDP di PyTorch a monte di FairScale, gli stati del modello sono frammentati contro tutti i GPU, una strategia che riduce il consumo di memoria su ciascuna GPU a costo di incorrere in un elevato sovraccarico di comunicazione che aumenta con le dimensioni del cluster e quindi riduce significativamente la scalabilità su larga scala. Al contrario, il parallelismo dei dati partizionati nel modello delle partizioni della libreria SMP afferma in a consapevole della scala modo partizionando ogni replica degli stati del modello solo all'interno un sottoinsieme di GPU.

Diamo un'occhiata più da vicino al partizionamento del modello sensibile alla scala in MiCS, la tecnologia principale alla base dei dati frammentati in parallelo. L'intuizione alla base di questo progetto è che il partizionamento degli stati di addestramento nell'intero gruppo di dati paralleli potrebbe non essere necessario per addestrare un modello con decine di miliardi di parametri. Ad esempio, 8 GPU V100 (32 GB ciascuna) sono sufficienti per contenere la replica degli stati del modello di un modello con parametri 10B che richiede circa 200 GB di memoria durante l'allenamento con l'ottimizzatore Adam utilizzando la precisione mista. Limitando una replica completa degli stati modello nel minore sottoinsieme di GPU, possiamo ridurre efficacemente la scala del sovraccarico di comunicazione rispetto a DeepSpeed ​​e PyTorch FSDP. Il parallelo dei dati frammentati sfrutta anche altre tecniche in MiCS come la comunicazione gerarchica e la sincronizzazione del gradiente a 2 salti. Per ulteriori informazioni, controlla Ridimensionamento quasi lineare dell'addestramento di modelli giganteschi su AWS or MiCS: ridimensionamento quasi lineare per l'addestramento di modelli giganteschi su cloud pubblico.

Ora, come fai a sapere quando scegliere i dati partizionati in parallelo rispetto ad altre tecniche di addestramento distribuito? La regola generale è che se il tuo modello ha meno di 1 miliardo di parametri e può adattarsi alla memoria della GPU, Libreria parallela di dati SageMaker or Compilatore di formazione SageMaker può esserti sufficiente. Se disponi di modelli di linguaggio o di visione artificiale più grandi, il nostro suggerimento è di addestrarlo con la tecnica del parallelismo dei dati frammentati combinata con checkpoint di attivazione ed scarico di attivazione nella libreria parallela del modello SageMaker prima, prima di altre tecniche come parallelismo tensoriale o parallelismo della pipeline.

Utilizzo del parallelismo dei dati partizionati per addestrare GPT-2 su Amazon SageMaker

Impariamo ora come addestrare un modello GPT-2 con dati partizionati in parallelo, con SMP che incapsula la complessità per te. Questo quaderno tutorial completo ti guida attraverso l'intero processo, dall'elaborazione dei dati, alla definizione e all'invio dei lavori di formazione, al monitoraggio dei registri di formazione. Quella che segue è una breve panoramica che evidenzia i passaggi chiave per l'utilizzo di questa funzione.

1. Iniziare

Il parallelismo dei dati frammentati è disponibile in PyTorch v1.12.0+ e funziona sia con FP16 che con BF16. Il modo più semplice per utilizzare la libreria SMP è tramite un AWS Deep Learning Container per PyTorch. Tuttavia, se vuoi portare il tuo container Docker, puoi fare riferimento a Crea il tuo contenitore Docker con la libreria parallela del modello distribuito SageMaker. Per iniziare, segui Modifica uno script di addestramento PyTorch per adattare le API degli SMP nello script di addestramento. In questa sezione vengono richiamati solo alcuni passaggi principali con frammenti di codice dallo script di addestramento pronto per l'uso train_gpt_simple.py. Puoi seguire i commenti nello script e Documento API per saperne di più su dove vengono utilizzate le API SMP.

Innanzitutto, importa e inizializza la libreria chiamando smdistributed.modelparallel.torch.init() all'inizio dello script di formazione:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

In secondo luogo, avvolgere il modello con cui partizionare smdistributed.modelparallel.torch.DistributedModel e usa il reso DistributedModel oggetto andando avanti:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_config(model_config)
model = smp.DistributedModel(model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation)

Avvolgi l'ottimizzatore con smdistributed.modelparallel.torch.DistributedOptimizer per salvare e caricare gli stati dell'ottimizzatore.

from torch import optim

optimizer = optim.Adam(
    param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
)

optimizer = smp.DistributedOptimizer(
        optimizer, 
        static_loss_scale=None, 
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
        )

Metti la logica avanti e indietro in una funzione passo e decorala con smdistributed.modelparallel.torch.step.  Qualsiasi calcolo definito all'interno di smp.step-decorated la funzione viene eseguita in modo distribuito.

@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    model.backward(loss)

    return loss

@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    
    return loss

2. Preparare il set di dati

Usiamo il testo web aperto è il set di dati che utilizziamo in questo esempio. Il taccuino utilizza lo script data_prep_512.py per scaricare e preelaborare il set di dati. Puoi anche allenarti con altri set di dati modificando data_pipeline.py. Quando si ha a che fare con set di dati e modelli di grandi dimensioni, è possibile accelerare il processo di addestramento utilizzando i dati archiviati Amazon FSx per Lustre, che fornisce un file system ad alte prestazioni integrato nativamente con Servizio di archiviazione semplice Amazon (S3). Si prega di consultare le istruzioni da Configura il canale di input dei dati per utilizzare Amazon FSx for Lustre per indicazioni sull'impostazione di un file system FSx Lustre come canale di input dei dati.

3. Inizia i lavori di formazione

Questo passaggio presuppone che tu l'abbia già fatto modificato il tuo script di allenamento e preparato il dataset come indicato nelle sezioni precedenti. Per abilitare il parallelismo dei dati partizionati, impostare semplicemente il sharded_data_parallel_degree nel Stimatore PyTorch. In questo tutorial, abbiamo impostato sharded_data_parallel_degree=128 ed instace_count=32 per i nodi p4d.24xlarge, che indica che gli stati del modello verranno suddivisi in 128 GPU su un totale di 256 GPU. Sulla base di questo valore selezionato, SMP imposterà automaticamente il grado di parallelismo dei dati su 2 (perché 256/128=2), il che significa che avremo due repliche per il parallelismo dei dati. Una regola generale per scegliere un valore ideale per sharded_data_parallel_degree consiste nell'aggiungere un altro nodo al gruppo di condivisione per ogni 3B di parametri del modello. In questo tutorial, la dimensione del nostro modello è 30B, quindi dovremmo usare almeno 10 nodi per lo sharding. E poiché 16 nodi (128 GPU) è la più piccola potenza di 2 sopra la soglia, abbiamo impostato sharded_data_parallel_degree=128.

Per il checkpoint, forniamo anche una serie di utilità di checkpoint in sharded_data_parallel_checkpoint.py , compresa un'utilità per ricostruire il pieno state_dict per casi d'uso avanzati. Infine, possiamo avviare un lavoro di formazione distribuito chiamando fit() su Estimator.

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    instance_type="ml.p4d.24xlarge",
    source_dir=os.getcwd(),
    volume_size=500,
    instance_count=32,
    distribution={
        "mpi": {
            "enabled": True,
            "processes_per_host": processes_per_host,
            "custom_mpi_options": mpioptions,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "ddp": True,
                    "skip_tracing": True,
                    "delayed_parameter_initialization": True,
                    "offload_activations": True,
                    "activation_loading_horizon": 4,
                    # To enable sharded data parallelism.
                    # Here we shard model states across 128 GPUs. 
                    "sharded_data_parallel_degree": 128, 
                    "fp16": False,
                    "bf16": True,
                    # This is to disable pipeline parallelism.
                    "partitions": 1,
                },
            }
        },
    },
    framework_version="1.12",
    py_version="py38",
    hyperparameters=hyperparameters,
    checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,
    checkpoint_local_path=hyperparameters["checkpoint-dir"] if use_fsx else None,
    ...
)

smp_estimator.fit(inputs=data_channels)

4. Monitorare i lavori di formazione

Puoi accedere ai registri di formazione e monitorare l'utilizzo della GPU e della memoria Amazon Cloud Watch. Assicurati di guardare i log di "algo-1" perché questo è il nodo principale il cui flusso di output ha i registri del lavoro di addestramento di tutte le istanze.

Prestazioni di benchmarking

Abbiamo confrontato il parallelismo dei dati partizionati nella libreria SMP su entrambi i nodi 16 e 32 p4d.24xlarge per la lunghezza della sequenza 512 e 2048, rispettivamente. Il modello GPT30 con parametro 2B è configurato per utilizzare una larghezza nascosta di 7168, 48 strati e 64 teste. È possibile adottare la stessa identica configurazione in cui la lunghezza della sequenza è 2048 impostando model_config = "gpt2-30b" nel taccuino tutorial. Con questa impostazione, SMP ha raggiunto 73.52 campioni al secondo, una velocità superiore del 39.7% rispetto a DeepSpeed ​​ZeRO-3. Se la dimensione del tuo token è di 500 miliardi, questa accelerazione significa quasi 367 ore di risparmio sui nodi p4d.24xlarge, l'equivalente di un budget di oltre $ 12,000 risparmiato per formazione! La tabella seguente riassume i nostri risultati di benchmark.

1/ Per ciascuna configurazione del modello, abbiamo testato diverse caratteristiche, fasi e configurazioni in DeepSpeed ​​ZeRO e abbiamo scelto quella che fornisce il miglior throughput come linea di base di DeepSpeed. Il benchmark è stato eseguito Cloud di calcolo elastico di Amazon (Amazon EC2). 2/ Questi risultati si basano su collettivi di comunicazione migliorati ottimizzati per AWS che saranno presto resi disponibili. 3/ Il tempo di allenamento viene proiettato dalla velocità in base al numero di gettoni elaborati.

In sintesi, abbiamo osservato un throughput costantemente più elevato con il parallelismo dei dati partizionati in SMP rispetto a DeepSpeed ​​su una vasta gamma di modelli e configurazioni. Questa nuova funzionalità ha anche dimostrato una migliore efficienza della memoria rispetto a DeepSpeed, consentendo a SMP di adattarsi a un batch di dimensioni maggiori e di ridurre il livello di accumulo del gradiente richiesto per adattarsi a una particolare dimensione di batch globale.

Conclusione

In questo post, abbiamo introdotto una nuova tecnica di addestramento distribuito, il parallelismo dei dati partizionati, e il modo in cui accelera l'addestramento di modelli giganteschi con un ridimensionamento quasi lineare su Amazon SageMaker. Abbiamo anche spiegato come addestrare un modello GPT-2 con la nuova tecnica in seguito esempio completo. Puoi seguire il Esempi di Amazon SageMaker repository GitHub per tenere traccia di tutti gli esempi di modelli paralleli SageMaker o partecipare al nostro prossimo laboratori di formazione distribuiti. Per ulteriori informazioni sul parallelismo dei dati partizionati, vedere il documentazione.

Circa gli autori

Emilia Webber si è unito ad AWS subito dopo il lancio di SageMaker e da allora ha cercato di parlarne al mondo! Oltre a creare nuove esperienze ML per i clienti, Emily ama meditare e studiare il buddismo tibetano.

Può Karakus è un Senior Applied Scientist presso AWS, che ottimizza il deep learning distribuito su larga scala su AWS. I suoi interessi di ricerca riguardano il deep learning, l'ottimizzazione distribuita, i sistemi distribuiti e la teoria dell'informazione. Al di fuori del lavoro, gli piace andare in bicicletta, viaggiare, leggere e imparare.

Rahul Huilgol è un Senior Software Engineer presso AWS. Lavora su sistemi di deep learning distribuiti, per rendere facile e performante il training di grandi modelli di deep learning nel cloud. Nel tempo libero ama la fotografia, il ciclismo e il giardinaggio.

Suhit Kodgule è un ingegnere di sviluppo software con il gruppo di intelligenza artificiale di AWS che lavora su framework di deep learning. Nel tempo libero ama fare escursioni, viaggiare e cucinare.

Erin Ho è un Product Manager per AWS Deep Learning. Lavora su prodotti che rendono più facile per i clienti addestrare modelli di deep learning su AWS. Per divertirsi fuori dal lavoro, le piace fare escursioni e sciare.