Generazione di testo più veloce con TensorFlow e XLA

'Generazione di testo più veloce con TensorFlow e XLA' (Fast text generation with TensorFlow and XLA)

TL;DR: La generazione di testo su 🤗 transformers utilizzando TensorFlow può ora essere compilata con XLA. È fino a 100 volte più veloce di prima, e ancora più veloce di PyTorch — controlla il colab qui sotto!

Generazione di testo

All’aumentare della qualità dei grandi modelli di linguaggio, aumentano anche le nostre aspettative su ciò che questi modelli possono fare. Specialmente dopo il rilascio del GPT-2 di OpenAI, i modelli con capacità di generazione di testo sono stati al centro dell’attenzione. E per motivi legittimi: questi modelli possono essere utilizzati per riassumere, tradurre e hanno persino dimostrato capacità di apprendimento zero-shot su alcune attività linguistiche. Questo post mostrerà come sfruttare al massimo questa tecnologia con TensorFlow.

La libreria 🤗 transformers è iniziata con modelli NLP, quindi è naturale che la generazione di testo sia di massima importanza per noi. Fa parte degli sforzi di democratizzazione di Hugging Face per garantire che sia accessibile, facilmente controllabile ed efficiente. C’è un post precedente sui diversi tipi di generazione di testo. Tuttavia, di seguito c’è un breve riepilogo delle funzionalità principali — sentiti libero di saltarlo se conosci la nostra funzione generate e vuoi passare direttamente alle specificità di TensorFlow.

Cominciamo dalle basi. La generazione di testo può essere deterministica o stocastica, a seconda del flag do_sample. Per impostazione predefinita è impostato su False, causando l’output a essere deterministico, anche noto come Decodifica Greedy. Quando è impostato su True, anche noto come campionamento, l’output sarà stocastico, ma è comunque possibile ottenere risultati riproducibili attraverso l’argomento seed (con lo stesso formato della generazione di numeri casuali di TensorFlow senza stato). Come regola generale, si desidera una generazione deterministica se si desidera ottenere informazioni factuali dal modello e una generazione stocastica se si mira a ottenere output più creativi.

• Proximal Policy Optimization (PPO)
• Benvenuti a fastai nell’Hugging Face Hub
• Abbiamo raccolto $100 milioni per l’apprendimento automatico aperto e collaborativo 🚀

# Richiede transformers >= 4.21.0;
# Gli output del campionamento possono essere diversi a seconda dell'hardware.
from transformers import AutoTokenizer, TFAutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
model.config.pad_token_id = model.config.eos_token_id
inputs = tokenizer(["TensorFlow è"], return_tensors="tf")

generated = model.generate(**inputs, do_sample=True, seed=(42, 0))
print("Output del campionamento: ", tokenizer.decode(generated[0]))
# > Output del campionamento: TensorFlow è una piattaforma di apprendimento eccellente per imparare
# sulla struttura dei dati e la struttura nella scienza dei dati.

A seconda dell’applicazione di destinazione, potrebbe essere desiderabile avere output più lunghi. Puoi controllare la lunghezza dell’output di generazione con max_new_tokens, tenendo presente che generazioni più lunghe richiederanno più risorse.

generated = model.generate(
    **inputs, do_sample=True, seed=(42, 0), max_new_tokens=5
)
print("Limitato a 5 nuovi token:", tokenizer.decode(generated[0]))
# > Limitato a 5 nuovi token: TensorFlow è una piattaforma di apprendimento eccellente per
generated = model.generate(
    **inputs, do_sample=True, seed=(42, 0), max_new_tokens=30
)
print("Limitato a 30 nuovi token:", tokenizer.decode(generated[0]))
# > Limitato a 30 nuovi token: TensorFlow è una piattaforma di apprendimento eccellente per
# imparare sulla struttura dei dati e la struttura nella scienza dei dati................

Il campionamento ha alcuni parametri che puoi regolare per controllare la casualità. Il più importante è temperature, che imposta l’entropia complessiva del tuo output — valori inferiori a 1.0 daranno priorità al campionamento di token con una probabilità più alta, mentre valori superiori a 1.0 faranno il contrario. Impostarlo a 0.0 riduce il comportamento alla Decodifica Greedy, mentre valori molto grandi approssimano il campionamento uniforme.

generated = model.generate(
    **inputs, do_sample=True, seed=(42, 0), temperature=0.7
)
print("Temperatura 0.7: ", tokenizer.decode(generated[0]))
# > Temperatura 0.7: TensorFlow è un ottimo modo per fare cose come queste........
generated = model.generate(
    **inputs, do_sample=True, seed=(42, 0), temperature=1.5
)
print("Temperatura 1.5: ", tokenizer.decode(generated[0]))
# > Temperatura 1.5: TensorFlow sta venendo sviluppato sia per Cython che per Bamboo.
# Su Bamboo...

A differenza del campionamento, la Generazione Greedy Decoding selezionerà sempre il token più probabile ad ogni iterazione di generazione. Tuttavia, spesso produce risultati sub-ottimali. Puoi aumentare la qualità dei risultati attraverso l’argomento num_beams. Quando è maggiore di 1, attiva la Beam Search, che esplora continuamente sequenze ad alta probabilità. Questa esplorazione comporta un costo aggiuntivo in termini di risorse e tempo di calcolo.

generated = model.generate(**inputs, num_beams=2)
print("Output Beam Search:", tokenizer.decode(generated[0]))
# > Output Beam Search: TensorFlow è un framework open-source, open-source,
# per l'applicazione distribuita

Infine, quando si esegue il Campionamento o la Beam Search, è possibile utilizzare num_return_sequences per restituire diverse sequenze. Nel caso del Campionamento, è equivalente ad eseguire più volte la generazione dalla stessa frase di input, mentre nel caso della Beam Search restituisce le sequenze generate con il punteggio più alto in ordine decrescente.

generated = model.generate(**inputs, num_beams=2, num_return_sequences=2)
print(
    "Tutte le ipotesi generate:",
    "\n".join(tokenizer.decode(out) for out in generated)
)
# > Tutte le ipotesi generate: TensorFlow è un framework open-source, open-source,
# per l'applicazione distribuita
# > TensorFlow è un framework open-source, open-source, per l'applicazione
# distribuita che permette

I fondamenti della generazione di testo, come puoi vedere, sono facili da controllare. Tuttavia, ci sono molte opzioni non coperte negli esempi precedenti ed è consigliato leggere la documentazione per casi d’uso avanzati. Purtroppo, quando si esegue generate con TensorFlow, potresti notare che impiega del tempo ad eseguire. Se la tua applicazione di destinazione richiede una bassa latenza o una grande quantità di prompt di input, eseguire la generazione di testo con TensorFlow sembra una procedura costosa. 😬

Non temere, perché il resto di questo post sul blog mira a dimostrare che una sola riga di codice può apportare un miglioramento drastico. Se preferisci passare direttamente all’azione, il colab ha un esempio interattivo con cui puoi sperimentare!

TensorFlow e XLA

XLA, o Accelerated Linear Algebra, è un compilatore sviluppato originariamente per accelerare i modelli TensorFlow. Oggi è anche il compilatore dietro JAX, e può essere utilizzato anche con PyTorch. Sebbene la parola “compilatore” possa sembrare spaventosa per alcuni, XLA è semplice da usare con TensorFlow: è incluso nella libreria tensorflow e può essere attivato con l’argomento jit_compile in qualsiasi funzione di creazione di grafi.

Per coloro che sono familiari con TensorFlow 1 🧓, il concetto di grafo TensorFlow è naturale, in quanto era l’unico modo di operare. Prima si definivano le operazioni in modo dichiarativo per creare un grafo. Successivamente, si potevano far passare gli input attraverso il grafo e osservare gli output. Veloce, efficiente, ma difficile da debuggare. Con TensorFlow 2 è arrivata l’Eager Execution e la possibilità di codificare i modelli in modo imperativo: il team di TensorFlow spiega le differenze in modo più dettagliato nel loro post sul blog.

Hugging Face scrive i suoi modelli TensorFlow tenendo presente l’Eager Execution. La trasparenza è un valore fondamentale e la possibilità di ispezionare le parti interne del modello in qualsiasi momento è molto vantaggiosa per questo scopo. Tuttavia, ciò significa che alcuni utilizzi dei modelli non traggono automaticamente vantaggio dalle prestazioni in modalità grafica (ad esempio quando si chiama model(args)).

Fortunatamente, il team di TensorFlow ha pensato a utenti come noi 🥳! Avvolgere una funzione che contiene codice TensorFlow con tf.function tenterà di convertirlo in un grafo quando si chiama la funzione avvolta. Se si sta addestrando un modello, chiamare model.compile() (senza run_eagerly=True) fa proprio questo avvolgimento, in modo da beneficiare della modalità grafica quando si chiama model.fit(). Poiché tf.function può essere utilizzato in qualsiasi funzione che contiene codice TensorFlow, significa che è possibile utilizzarlo anche nelle funzioni che vanno oltre l’inferenza del modello, creando un singolo grafo ottimizzato.

Ora che sai come creare i grafi TensorFlow, compilando con XLA è semplice: basta aggiungere jit_compile=True come argomento alle funzioni menzionate in precedenza (tf.function e tf.keras.Model.compile). Supponendo che tutto sia andato bene (ne parleremo più avanti) e che tu stia usando una GPU o un TPU, noterai che la prima chiamata richiederà del tempo, ma le successive saranno molto, molto più veloci. Ecco un semplice esempio di una funzione che esegue l’inferenza del modello e qualche post-elaborazione dei suoi output:

# Nota: i tempi di esecuzione dipendono fortemente dall'hardware -- qui è stato utilizzato un 3090.
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
inputs = tokenizer(["TensorFlow è"], return_tensors="tf")

def token_successivo_piu_probabile(inputs):
    model_output = model(inputs)
    return tf.argmax(model_output.logits[:, -1, :], axis=-1)

print("Chiamata alla funzione regolare con codice TensorFlow...")
token_successivo_piu_probabile(inputs)
# > Tempo di esecuzione -- 48.8 ms

In una riga, puoi creare una funzione accelerata XLA dalla funzione precedente.

xla_token_successivo_piu_probabile = tf.function(token_successivo_piu_probabile, jit_compile=True)

print("Chiamata alla funzione XLA... (per la prima volta -- sarà lento)")
xla_token_successivo_piu_probabile(inputs)
# > Tempo di esecuzione -- 3951.0 ms
print("Chiamata alla funzione XLA... (per la seconda volta -- sarà veloce)")
xla_token_successivo_piu_probabile(inputs)
# > Tempo di esecuzione -- 1.6 ms

Generazione di Testo utilizzando TensorFlow con XLA

Come con qualsiasi procedura di ottimizzazione, non esiste un pranzo gratis — XLA non fa eccezione. Dal punto di vista di un utente di generazione di testo, c’è solo un aspetto tecnico che devi tenere a mente. Senza addentrarci troppo nei dettagli, XLA utilizzato in questa modalità esegue la compilazione just-in-time (JIT) di una tf.function quando la chiami, che si basa sul polimorfismo.

Quando compili una funzione in questo modo, XLA tiene traccia della forma e del tipo di ogni tensore, così come dei dati di ogni input di funzione non tensore. La funzione viene compilata in un binario, e ogni volta che viene chiamata con la stessa forma e tipo di tensore (con QUALSIASI dato di tensore) e gli stessi argomenti non tensori, la funzione compilata può essere riutilizzata. Al contrario, se chiami la funzione con una forma o un tipo diverso in un tensore di input, o se usi un argomento non tensore diverso, allora verrà eseguito un nuovo passaggio di compilazione costoso. Riassunto in un semplice esempio:

# Nota: i tempi di esecuzione dipendono fortemente dall'hardware -- qui è stato utilizzato un 3090.
import tensorflow as tf

@tf.function(jit_compile=True)
def massimo_piu_costante(tensore, scalare):
    return tf.math.reduce_max(tensore) + scalare

# Lento: la compilazione XLA entrerà in azione, poiché è la prima chiamata
massimo_piu_costante(tf.constant([0, 0, 0]), 1)
# > Tempo di esecuzione -- 520.4 ms

# Veloce: Non è la prima chiamata con questa forma di tensore, tipo di tensore e stesso
# argomento non tensore
massimo_piu_costante(tf.constant([1000, 0, -10]), 1)
# > Tempo di esecuzione -- 0.6 ms

# Lento: Tipo di tensore diverso
massimo_piu_costante(tf.constant([0, 0, 0], dtype=tf.int64), 1)
# > Tempo di esecuzione -- 27.1 ms

# Lento: Forma di tensore diversa
massimo_piu_costante(tf.constant([0, 0, 0, 0]), 1)
# > Tempo di esecuzione -- 25.5 ms

# Lento: Argomento non tensore diverso
massimo_piu_costante(tf.constant([0, 0, 0]), 2)
# > Tempo di esecuzione -- 24.9 ms

Nella pratica, per la generazione di testo, ciò significa semplicemente che l’input dovrebbe essere riempito fino a un multiplo di una certa lunghezza (in modo da avere un numero limitato di forme possibili), e che l’utilizzo di diverse opzioni sarà lento la prima volta che le si utilizzano. Vediamo cosa succede quando chiami in modo ingenuo la generazione con XLA.

# Nota: i tempi di esecuzione dipendono fortemente dall'hardware -- qui è stato utilizzato un 3090.
import time
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

# Nota il nuovo argomento, `padding_side="left"` -- i modelli che generano solo testo, che possono
# essere istanziati con TFAutoModelForCausalLM, devono essere riempiti a sinistra, in quanto
# continuano a generare a partire dall'input iniziale.
tokenizer = AutoTokenizer.from_pretrained(
    "gpt2", padding_side="left", pad_token="</s>"
)
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
model.config.pad_token_id = model.config.eos_token_id
input_1 = ["TensorFlow è"]
input_2 = ["TensorFlow è un"]

# Una riga per creare una funzione XLA di generazione
xla_genera = tf.function(model.generate, jit_compile=True)

# Chiama la generazione XLA senza riempimento
input_tokenizzato_1 = tokenizer(input_1, return_tensors="tf")  # lunghezza = 4
input_tokenizzato_2 = tokenizer(input_2, return_tensors="tf")  # lunghezza = 5
print(f"Forma di `input_tokenizzato_1` = {input_tokenizzato_1.input_ids.shape}")
print(f"Forma di `input_tokenizzato_2` = {input_tokenizzato_2.input_ids.shape}")

print("Chiamata alla generazione XLA con input_tokenizzato_1...")
print("(sarà lenta in quanto è la prima chiamata)")
start = time.time_ns()
xla_genera(**input_tokenizzato_1)
end = time.time_ns()
print(f"Tempo di esecuzione -- {(end - start) / 1e6:.1f} ms\n")
# > Tempo di esecuzione -- 9565.1 ms

print("Chiamata alla generazione XLA con input_tokenizzato_2...")
print("(ha una lunghezza diversa = attiverà nuovamente la tracciatura)")
start = time.time_ns()
xla_genera(**input_tokenizzato_2)
end = time.time_ns()
print(f"Tempo di esecuzione -- {(end - start) / 1e6:.1f} ms\n")
# > Tempo di esecuzione -- 6815.0 ms

Oh no, è terribilmente lento! Una soluzione per tenere sotto controllo le diverse combinazioni di forme è attraverso il padding, come accennato in precedenza. Le classi tokenizer hanno un argomento pad_to_multiple_of che può essere usato per ottenere un equilibrio tra accettare qualsiasi lunghezza di input e limitare il tracciamento.

padding_kwargs = {"pad_to_multiple_of": 8, "padding": True}
tokenized_input_1_with_padding = tokenizer(
    input_1, return_tensors="tf", **padding_kwargs
)  # lunghezza = 8
tokenized_input_2_with_padding = tokenizer(
    input_2, return_tensors="tf", **padding_kwargs
)  # lunghezza = 8
print(
    "Forma di `tokenized_input_1_with_padding` = ",
    f"{tokenized_input_1_with_padding.input_ids.shape}"
)
print(
    "Forma di `tokenized_input_2_with_padding` = ",
    f"{tokenized_input_2_with_padding.input_ids.shape}"
)

print("Chiamata alla generazione XLA con tokenized_input_1_with_padding...")
print("(lento, la prima volta che viene eseguito con questa lunghezza)")
start = time.time_ns()
xla_generate(**tokenized_input_1_with_padding)
end = time.time_ns()
print(f"Tempo di esecuzione -- {(end - start) / 1e6:.1f} ms\n")
# > Tempo di esecuzione -- 6815.4 ms

print("Chiamata alla generazione XLA con tokenized_input_2_with_padding...")
print("(sarà veloce!)")
start = time.time_ns()
xla_generate(**tokenized_input_2_with_padding)
end = time.time_ns()
print(f"Tempo di esecuzione -- {(end - start) / 1e6:.1f} ms\n")
# > Tempo di esecuzione -- 19.3 ms

Questo è molto meglio, le chiamate successive alla generazione eseguite in questo modo saranno ordini di grandezza più veloci rispetto a prima. Tieni presente che provare nuove opzioni di generazione, in qualsiasi momento, attiverà il tracciamento.

print("Chiamata alla generazione XLA con lo stesso input, ma con nuove opzioni...")
print("(lento di nuovo)")
start = time.time_ns()
xla_generate(**tokenized_input_1_with_padding, num_beams=2)
end = time.time_ns()
print(f"Tempo di esecuzione -- {(end - start) / 1e6:.1f} ms\n")
# > Tempo di esecuzione -- 9644.2 ms

Dal punto di vista dello sviluppatore, fare affidamento su XLA implica essere consapevoli di alcune sfumature aggiuntive. XLA brilla quando la dimensione delle strutture dati è nota in anticipo, come ad esempio nell’allenamento del modello. D’altra parte, quando le loro dimensioni sono impossibili da determinare o vengono utilizzate determinate parti dinamiche, XLA non riesce a compilare. Le implementazioni moderne della generazione di testo sono auto-regressive, il cui comportamento naturale è espandere i tensori e interrompere bruscamente alcune operazioni man mano che procede – in altre parole, non sono di default XLA-friendly. Abbiamo riscritto l’intero codice di generazione di testo di TensorFlow per vettorizzare le operazioni e utilizzare strutture di dimensioni fisse con padding. I nostri modelli NLP sono stati anche modificati per utilizzare correttamente i loro embedding posizionali in presenza di strutture con padding. Il risultato dovrebbe essere invisibile agli utenti della generazione di testo di TensorFlow, ad eccezione della disponibilità della compilazione XLA.

Benchmarks e Conclusioni

Nel paragrafo precedente hai visto che puoi convertire le funzioni TensorFlow in un grafo e accelerarle con la compilazione XLA. Le forme attuali della generazione di testo sono semplicemente funzioni auto-regressive che alternano tra un passaggio in avanti del modello e una post-elaborazione, producendo un token per iterazione. Attraverso la compilazione XLA, l’intero processo viene ottimizzato, con conseguente esecuzione più veloce. Ma quanto più veloce? La demo di Gradio qui sotto contiene alcuni benchmark che confrontano la generazione di testo di Hugging Face su modelli GPU multipli per i due principali framework di ML, TensorFlow e PyTorch.

Se esplori i risultati, due conclusioni diventano rapidamente visibili:

1. Come questo post del blog ha dimostrato finora, la generazione di testo di TensorFlow è molto più veloce quando si utilizza XLA. Stiamo parlando di velocizzazioni superiori a 100 volte in alcuni casi, il che dimostra veramente la potenza di un grafo compilato 🚀
2. La generazione di testo di TensorFlow con XLA è l’opzione più veloce nella stragrande maggioranza dei casi, in alcuni di essi addirittura fino a 9 volte più veloce, sfatando il mito che PyTorch sia il framework di riferimento per compiti NLP seri 💪

Prova il colab e goditi il potere della generazione di testo potenziata da XLA!