Addestrare un modello di linguaggio con 🤗 Transformers utilizzando TensorFlow e TPUs

'Addestramento modello linguaggio 🤗 Transformers con TensorFlow e TPUs'

Introduzione

Il training su TPU è una competenza utile da avere: i pod TPU sono ad alte prestazioni ed estremamente scalabili, rendendo facile allenare modelli di qualsiasi scala, da qualche decina di milioni di parametri fino a dimensioni veramente enormi: il modello PaLM di Google (oltre 500 miliardi di parametri!) è stato allenato interamente su pod TPU.

In precedenza abbiamo scritto un tutorial e un esempio di Colab che mostrano il training su TPU su piccola scala con TensorFlow e introducono i concetti fondamentali che devi comprendere per far funzionare il tuo modello su TPU. Questa volta, faremo un passo avanti e alleneremo un masked language model da zero utilizzando TensorFlow e TPU, includendo ogni passaggio, dal training del tokenizer alla preparazione del dataset fino all’allenamento finale del modello e all’upload. Questo è il tipo di compito per il quale probabilmente vorrai dedicare un nodo TPU (o una VM), anziché utilizzare solo Colab, ed è su questo che ci concentreremo.

Come nel nostro esempio di Colab, sfruttiamo il supporto molto pulito di TensorFlow per TPU tramite XLA e TPUStrategy. Beneficeremo anche del fatto che la maggior parte dei modelli TensorFlow in 🤗 Transformers sono completamente compatibili con XLA. Quindi sorprendentemente, è necessario fare poco lavoro per farli funzionare su TPU.

A differenza del nostro esempio di Colab, però, questo esempio è progettato per essere scalabile e molto più simile a un vero allenamento realistico, anche se di default utilizziamo un modello di dimensioni BERT, il codice potrebbe essere espanso per un modello molto più grande e una slice di TPU pod molto più potente cambiando alcune opzioni di configurazione.

Motivazione

Perché stiamo scrivendo questa guida ora? Dopotutto, 🤗 Transformers ha supportato TensorFlow per diversi anni. Ma far allenare quei modelli su TPUs è stato un punto critico per la comunità. Ciò è dovuto al fatto che:

  • Molti modelli non erano compatibili con XLA
  • I collatori di dati non utilizzavano operazioni native di TF

Crediamo che XLA sia il futuro: è il compilatore principale per JAX, ha un supporto di prima classe in TensorFlow e puoi persino utilizzarlo da PyTorch. Pertanto, abbiamo fatto uno sforzo considerevole per rendere il nostro codice compatibile con XLA e rimuovere qualsiasi altro ostacolo che potesse impedire la compatibilità tra XLA e TPU. Ciò significa che gli utenti dovrebbero essere in grado di allenare la maggior parte dei nostri modelli TensorFlow su TPUs senza problemi.

C’è anche un’altra importante ragione per preoccuparsi del training su TPU in questo momento: i recenti grandi progressi nei LLM e nell’IA generativa hanno creato un enorme interesse pubblico nell’allenamento dei modelli, ed è diventato incredibilmente difficile per la maggior parte delle persone accedere alle GPU all’avanguardia. Sapere come allenarsi su TPU ti offre un altro percorso per accedere all’hardware di calcolo ad alte prestazioni, il che è molto più dignitoso che perdere una guerra di offerte per l’ultimo H100 su eBay e poi piangere disperatamente alla tua scrivania. Meriti di meglio. E parlando per esperienza: una volta che ti abitui all’allenamento su TPU, potresti non voler tornare indietro.

Cosa aspettarsi

Stiamo per allenare un RoBERTa (modello base) da zero sul dataset WikiText (v1). Oltre ad allenare il modello, alleneremo anche il tokenizer, suddividiamo i dati in token e li caricheremo su Google Cloud Storage in formato TFRecord, dove saranno accessibili per il training su TPU. Puoi trovare tutto il codice in questa directory. Se sei un certo tipo di persona, puoi saltare il resto di questo post sul blog e passare direttamente al codice. Se invece rimani, daremo uno sguardo più approfondito ad alcune delle idee chiave nel codice.

Molte delle idee qui menzionate sono state anche citate nel nostro esempio di Colab, ma volevamo mostrare agli utenti un esempio completo end-to-end che mette tutto insieme e lo mostra in azione, anziché coprire i concetti a un livello superficiale. Il diagramma seguente ti offre una panoramica visiva dei passaggi coinvolti nell’allenamento di un modello di linguaggio con 🤗 Transformers utilizzando TensorFlow e TPUs:

Ottenere i dati e allenare un tokenizer

Come accennato, abbiamo utilizzato il dataset WikiText (v1). Puoi visitare la pagina del dataset su Hugging Face Hub per esplorare il dataset.

Dato che l’insieme di dati è già disponibile su Hub in un formato compatibile, possiamo caricarlo e interagire con esso facilmente utilizzando 🤗 datasets. Tuttavia, per questo esempio, poiché stiamo anche addestrando un tokenizer da zero, ecco cosa abbiamo fatto:

  • Caricato la divisione train del WikiText utilizzando 🤗 datasets.
  • Sfruttato i 🤗 tokenizers per addestrare un modello Unigram.
  • Caricato il tokenizer addestrato su Hub.

Puoi trovare il codice di addestramento del tokenizer qui e il tokenizer qui. Questo script ti permette anche di eseguirlo con qualsiasi insieme di dati compatibile da Hub.

💡 È facile utilizzare 🤗 datasets per ospitare i tuoi insiemi di dati testuali. Consulta questa guida per saperne di più.

Tokenizzazione dei dati e creazione di TFRecords

Una volta addestrato il tokenizer, possiamo utilizzarlo su tutte le divisioni dell’insieme di dati ( train , validation , e test in questo caso) e creare frammenti di TFRecord da essi. Avere le divisioni dell’insieme di dati distribuite su più frammenti di TFRecord aiuta con l’elaborazione massivamente parallela rispetto ad avere ciascuna divisione in singoli file TFRecord.

Tokenizziamo i campioni singolarmente. Prendiamo poi un batch di campioni, li concateniamo insieme e li dividiamo in diversi chunk di dimensione fissa (128 nel nostro caso). Seguiamo questa strategia anziché tokenizzare un batch di campioni con una lunghezza fissa per evitare la riduzione aggressiva del contenuto testuale (a causa della troncatura).

Poi prendiamo questi campioni tokenizzati a batch e serializziamo quei batch come multipli frammenti di TFRecord, dove la lunghezza totale dell’insieme di dati e la dimensione individuale del frammento determinano il numero di frammenti. Infine, questi frammenti vengono caricati in un bucket di Google Cloud Storage (GCS).

Se stai usando un nodo TPU per l’addestramento, allora i dati devono essere trasmessi da un bucket di GCS poiché la memoria host del nodo è molto piccola. Ma per le VM TPU, possiamo utilizzare insiemi di dati in locale o persino allegare archiviazione persistente a quelle VM. Poiché i nodi TPU sono ancora abbastanza utilizzati, abbiamo basato il nostro esempio sull’utilizzo di un bucket di GCS per lo storage dei dati.

Puoi vedere tutto questo nel codice in questo script. Per comodità, abbiamo anche ospitato i frammenti resultant di TFRecord in questo repository su Hub.

Addestramento di un modello su dati in GCS

Se sei familiare con l’utilizzo di 🤗 Transformers, allora conosci già il codice di modellazione:

from transformers import AutoConfig, AutoTokenizer, TFAutoModelForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("tf-tpu/unigram-tokenizer-wikitext")

config = AutoConfig.from_pretrained("roberta-base")
config.vocab_size = tokenizer.vocab_size
model = TFAutoModelForMaskedLM.from_config(config) 

Ma poiché siamo nel territorio TPU, dobbiamo eseguire questa inizializzazione sotto un contesto di strategia in modo che possa essere distribuita tra i lavoratori TPU con addestramento parallelo dei dati:

import tensorflow as tf

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(...)
strategy = tf.distribute.TPUStrategy(tpu)

with strategy.scope():
    tokenizer = AutoTokenizer.from_pretrained("tf-tpu/unigram-tokenizer-wikitext")
    config = AutoConfig.from_pretrained("roberta-base")
    config.vocab_size = tokenizer.vocab_size
    model = TFAutoModelForMaskedLM.from_config(config) 

Allo stesso modo, l’ottimizzatore deve essere inizializzato anche sotto lo stesso contesto di strategia con cui il modello verrà successivamente compilato. Non vogliamo affrontare il codice completo di addestramento in questo post, quindi ti invitiamo a leggerlo qui. Invece, discutiamo un altro punto chiave di — un collatore di dati nativo di TensorFlow — DataCollatorForLanguageModeling.

DataCollatorForLanguageModeling è responsabile della mascheratura dei token selezionati casualmente dalla sequenza di input e della preparazione delle etichette. Per impostazione predefinita, restituiamo i risultati da questi collatori come array NumPy. Tuttavia, molti collatori supportano anche il ritorno di questi valori come tensori TensorFlow se specificato return_tensor="tf". Questo è stato cruciale affinché il nostro flusso di dati fosse compatibile con l’addestramento TPU.

Fortunatamente, TensorFlow fornisce un supporto senza soluzione di continuità per la lettura di file da un bucket GCS:

training_records = tf.io.gfile.glob(os.path.join(args.train_dataset, "*.tfrecord"))

Se args.dataset contiene l’identificatore gs://, TensorFlow capirà che deve cercare in un bucket GCS. Caricare localmente è semplice come rimuovere l’identificatore gs://. Per il resto del codice legato al flusso di dati, puoi fare riferimento a questa sezione nello script di addestramento.

Una volta che i dataset sono stati preparati, il modello e l’ottimizzatore sono stati inizializzati e il modello è stato compilato, possiamo fare quello che la comunità preferisce – model.fit(). Per l’addestramento, non abbiamo effettuato una sintonizzazione iperparametrica approfondita. L’abbiamo semplicemente addestrato più a lungo con un tasso di apprendimento di 1e-4. Abbiamo anche sfruttato il PushToHubCallback per il salvataggio dei checkpoint del modello e la sincronizzazione con il Hub. Puoi trovare i dettagli degli iperparametri e un modello addestrato qui: https://huggingface.co/tf-tpu/roberta-base-epochs-500-no-wd.

Una volta addestrato il modello, eseguire inferenze con esso è semplice come:

from transformers import pipeline

model_id = "tf-tpu/roberta-base-epochs-500-no-wd"
unmasker = pipeline("fill-mask", model=model_id, framework="tf")
unmasker("L'obiettivo della mia vita è [MASK].")

[{'score': 0.1003185287117958,
  'token': 52,
  'token_str': 'essere',
  'sequence': "L'obiettivo della mia vita è essere."},
 {'score': 0.032648514956235886,
  'token': 5,
  'token_str': '',
  'sequence': "L'obiettivo della mia vita è ."},
 {'score': 0.02152673341333866,
  'token': 138,
  'token_str': 'lavorare',
  'sequence': "L'obiettivo della mia vita è lavorare."},
 {'score': 0.019547373056411743,
  'token': 984,
  'token_str': 'recitare',
  'sequence': "L'obiettivo della mia vita è recitare."},
 {'score': 0.01939118467271328,
  'token': 73,
  'token_str': 'avere',
  'sequence': "L'obiettivo della mia vita è avere."}]

Conclusione

Se c’è una cosa che vogliamo sottolineare con questo esempio, è che l’addestramento su TPU è potente, scalabile e facile. Infatti, se stai già utilizzando modelli Transformers con TF/Keras e dati in streaming da tf.data, potresti rimanere sorpreso da quanto poco lavoro sia necessario per spostare l’intero flusso di addestramento su TPU. Hanno una reputazione di hardware complesso e di fascia alta, ma sono abbastanza accessibili e istanziare una grande fetta di pod è sicuramente più facile che mantenere sincronizzati server GPU multipli!

Diversificare l’hardware su cui vengono addestrati i modelli all’avanguardia sarà fondamentale negli anni 2020, soprattutto se continua la carenza di GPU. Speriamo che questa guida ti darà gli strumenti necessari per alimentare esecuzioni di addestramento all’avanguardia, indipendentemente dalle circostanze che affronti.

Come disse una volta il grande poeta GPT-4:

Se riesci a mantenere la calma quando tutto intorno a te sta perdendo la testa a causa delle carenze di GPU, e a fidarti del tuo codice, mentre gli altri dubitano, per addestrare su TPUs, senza esitazioni;

Se riesci a imparare dagli errori e a procedere, e a ottimizzare il tuo obiettivo per raggiungere il cielo, La tua sarà la strada per la padronanza dell’AI, E prevarrai, amico mio, col passare del tempo.

Sì, è un palese plagio di Rudyard Kipling e non ha idea di come pronunciare “drought”, ma speriamo che ti senta comunque ispirato.