 Visualizza su ai.google.dev
|
 Esegui in Google Colab
|
 Corri su Kaggle
|
 Apri in Vertex AI
|
 Visualizza il codice sorgente su GitHub
|
Panoramica
Gemma è una famiglia di modelli aperti leggeri e all'avanguardia creati sulla base della ricerca e della tecnologia utilizzata per creare i modelli Google Gemini. Gemma può essere ulteriormente ottimizzata in base a esigenze specifiche. Tuttavia, i modelli linguistici di grandi dimensioni, come Gemma, possono essere molto grandi e alcuni potrebbero non essere adatti a un singolo acceleratore per la messa a punto. In questo caso, esistono due approcci generali per perfezionarli:
- L'ottimizzazione fine efficiente dei parametri (PEFT), che mira a ridurre le dimensioni effettive del modello sacrificando un po' di fedeltà. LoRA rientra in questa categoria e il tutorial Eseguire l'ottimizzazione fine dei modelli Gemma in Keras utilizzando LoRA mostra come ottimizzare il modello Gemma 2B
gemma_2b_encon LoRA utilizzando KerasNLP su una singola GPU. - Ottimizzazione completa dei parametri con parallelismo dei modelli. Il parallelismo dei modelli distribuisce i pesi di un singolo modello su più dispositivi e consente il ridimensionamento orizzontale. Puoi scoprire di più sull'addestramento distribuito in questa guida di Keras.
Questo tutorial illustra come utilizzare Keras con un backend JAX per perfezionare il modello Gemma 7B con LoRa e l'addestramento distribuito con parallelismo del modello sulla Tensor Processing Unit (TPU) di Google. Tieni presente che in questo tutorial LoRA può essere disattivato per un'ottimizzazione completa dei parametri più lenta, ma più accurata.
Utilizzare gli acceleratori
Tecnicamente, per questo tutorial puoi utilizzare TPU o GPU.
Note sugli ambienti TPU
Google offre tre prodotti che forniscono TPU:
- Colab fornisce TPU v2 senza costi, un valore sufficiente per questo tutorial.
- Kaggle offre TPU v3 senza costi e funzionano anche per questo tutorial.
- Cloud TPU offre TPU v3 e generazioni più recenti. Un modo per configurarlo è:
- Crea una nuova VM TPU
- Configura l'inoltro alla porta SSH per la porta del server Jupyter che vuoi utilizzare
- Installa Jupyter e avvialo sulla VM TPU, quindi connettiti a Colab tramite "Connetti a un runtime locale"
Note sulla configurazione multi-GPU
Sebbene questo tutorial si concentri sul caso d'uso delle TPU, puoi adattarlo facilmente alle tue esigenze se hai una macchina multi-GPU.
Se preferisci lavorare tramite Colab, è anche possibile eseguire il provisioning di una VM multi-GPU per Colab direttamente tramite "Connettiti a una VM GCE personalizzata" nel menu Connetti di Colab.
Qui ci concentreremo sull'utilizzo della TPU senza costi di Kaggle.
Prima di iniziare
Credenziali di Kaggle
I modelli Gemma sono ospitati da Kaggle. Per utilizzare Gemma, richiedi l'accesso su Kaggle:
- Accedi o registrati su kaggle.com
- Apri la scheda del modello Gemma e seleziona "Richiedi accesso".
- Compila il modulo di consenso e accetta i Termini e condizioni
Per utilizzare l'API Kaggle, crea un token API:
- Apri le impostazioni di Kaggle.
- Seleziona "Create New Token" (Crea nuovo token).
- Viene scaricato un file
kaggle.json. Contiene le tue credenziali Kaggle
Esegui la seguente cella e inserisci le tue credenziali Kaggle quando richiesto.
# If you are using Kaggle, you don't need to login again.
!pip install ipywidgets
import kagglehub
kagglehub.login()
VBox(children=(HTML(value='<center> <img\nsrc=https://www.kaggle.com/static/images/site-logo.png\nalt=\'Kaggle…
Un modo alternativo è impostare KAGGLE_USERNAME e KAGGLE_KEY nel tuo ambiente se kagglehub.login() non funziona.
Installazione
Installare Keras e KerasNLP con il modello Gemma.
pip install -q -U keras-nlp# Work around an import error with tensorflow-hub. The library is not used.pip install -q -U tensorflow-hub# Install tensorflow-cpu so tensorflow does not attempt to access the TPU.pip install -q -U tensorflow-cpu tensorflow-text# Install keras 3 last. See https://keras.io/getting_started for details.pip install -q -U keras
Configura il backend JAX di Keras
Importa JAX ed esegui un controllo di convalida su TPU. Kaggle offre dispositivi TPUv3-8 con 8 core TPU e 16 GB di memoria ciascuno.
import jax
jax.devices()
[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]
import os
# The Keras 3 distribution API is only implemented for the JAX backend for now
os.environ["KERAS_BACKEND"] = "jax"
# Pre-allocate 90% of TPU memory to minimize memory fragmentation and allocation
# overhead
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.9"
Carica modello
import keras
import keras_nlp
Note sull'addestramento con precisione mista su GPU NVIDIA
Durante l'addestramento su GPU NVIDIA, è possibile utilizzare la precisione mista (keras.mixed_precision.set_global_policy('mixed_bfloat16')) per accelerare l'addestramento con un effetto minimo sulla qualità dell'addestramento. Nella maggior parte dei casi, è consigliabile attivare la precisione mista, in quanto consente di risparmiare memoria e tempo. Tuttavia, tieni presente che con batch di piccole dimensioni, l'utilizzo della memoria può aumentare di 1,5 volte (i pesi verranno caricati due volte, con una precisione dimezzata e la precisione massima).
Per l'inferenza, la precisione a metà (keras.config.set_floatx("bfloat16")) funzionerà e risparmierà memoria, mentre la precisione mista non è applicabile.
# Uncomment the line below if you want to enable mixed precision training on GPUs
# keras.mixed_precision.set_global_policy('mixed_bfloat16')
Per caricare il modello con i pesi e i tensori distribuiti sulle TPU, crea prima un nuovo DeviceMesh. DeviceMesh rappresenta una raccolta di dispositivi hardware configurati per il calcolo distribuito ed è stata introdotta in Keras 3 nell'ambito dell'API di distribuzione unificata.
L'API di distribuzione consente il parallelismo dei dati e dei modelli, consentendo una scalabilità efficiente dei modelli di deep learning su più acceleratori e host. Sfrutta il framework di base (ad es. JAX) per distribuire il programma e i tensori in base alle direttive di sharding tramite una procedura chiamata espansione SPMD (single program, multiple data). Scopri di più nella nuova guida all'API di distribuzione di Keras 3.
# Create a device mesh with (1, 8) shape so that the weights are sharded across
# all 8 TPUs.
device_mesh = keras.distribution.DeviceMesh(
(1, 8),
["batch", "model"],
devices=keras.distribution.list_devices())
LayoutMap dell'API di distribuzione specifica come devono essere suddivisi o replicati i pesi e i tensori, utilizzando le chiavi di stringa, ad esempio token_embedding/embeddings di seguito, che vengono trattate come regex per abbinare i percorsi dei tensori. I tensori corrispondenti vengono suddivisi in parti con le dimensioni del modello (8 TPU); gli altri verranno replicati completamente.
model_dim = "model"
layout_map = keras.distribution.LayoutMap(device_mesh)
# Weights that match 'token_embedding/embeddings' will be sharded on 8 TPUs
layout_map["token_embedding/embeddings"] = (model_dim, None)
# Regex to match against the query, key and value matrices in the decoder
# attention layers
layout_map["decoder_block.*attention.*(query|key|value).*kernel"] = (
model_dim, None, None)
layout_map["decoder_block.*attention_output.*kernel"] = (
model_dim, None, None)
layout_map["decoder_block.*ffw_gating.*kernel"] = (None, model_dim)
layout_map["decoder_block.*ffw_linear.*kernel"] = (model_dim, None)
ModelParallel ti consente di suddividere i pesi del modello o i tensori di attivazione su tutti i dispositivi su DeviceMesh. In questo caso, alcuni pesi del modello Gemma 7B vengono suddivisi con sharding su 8 chip TPU in base al layout_map definito sopra. Ora carica il modello in modo distribuito.
model_parallel = keras.distribution.ModelParallel(
layout_map=layout_map, batch_dim_name="batch")
keras.distribution.set_distribution(model_parallel)
gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset("gemma_7b_en")
Attaching 'config.json' from model 'keras/gemma/keras/gemma_7b_en/1' to your Kaggle notebook... Attaching 'config.json' from model 'keras/gemma/keras/gemma_7b_en/1' to your Kaggle notebook... Attaching 'model.weights.h5' from model 'keras/gemma/keras/gemma_7b_en/1' to your Kaggle notebook... Attaching 'tokenizer.json' from model 'keras/gemma/keras/gemma_7b_en/1' to your Kaggle notebook... Attaching 'assets/tokenizer/vocabulary.spm' from model 'keras/gemma/keras/gemma_7b_en/1' to your Kaggle notebook... normalizer.cc(51) LOG(INFO) precompiled_charsmap is empty. use identity normalization.
Ora verifica che il modello sia stato partizionato correttamente. Prendiamo decoder_block_1 come esempio.
decoder_block_1 = gemma_lm.backbone.get_layer('decoder_block_1')
print(type(decoder_block_1))
for variable in decoder_block_1.weights:
print(f'{variable.path:<58} {str(variable.shape):<16} {str(variable.value.sharding.spec)}')
<class 'keras_nlp.src.models.gemma.gemma_decoder_block.GemmaDecoderBlock'>
decoder_block_1/pre_attention_norm/scale (3072,) PartitionSpec(None,)
decoder_block_1/attention/query/kernel (16, 3072, 256) PartitionSpec(None, 'model', None)
decoder_block_1/attention/key/kernel (16, 3072, 256) PartitionSpec(None, 'model', None)
decoder_block_1/attention/value/kernel (16, 3072, 256) PartitionSpec(None, 'model', None)
decoder_block_1/attention/attention_output/kernel (16, 256, 3072) PartitionSpec(None, None, 'model')
decoder_block_1/pre_ffw_norm/scale (3072,) PartitionSpec(None,)
decoder_block_1/ffw_gating/kernel (3072, 24576) PartitionSpec('model', None)
decoder_block_1/ffw_gating_2/kernel (3072, 24576) PartitionSpec('model', None)
decoder_block_1/ffw_linear/kernel (24576, 3072) PartitionSpec(None, 'model')
Interruzione prima del perfezionamento
gemma_lm.generate("Best comedy movies in the 90s ", max_length=64)
'Best comedy movies in the 90s 1. The Naked Gun 2½: The Smell of Fear (1991) 2. Wayne’s World (1992) 3. The Naked Gun 33⅓: The Final Insult (1994)'
Il modello genera un elenco di film commedia degli anni '90 da guardare. Ora perfezioniamo il modello Gemma per cambiare lo stile dell'output.
Ottimizza con IMDB
import tensorflow_datasets as tfds
imdb_train = tfds.load(
"imdb_reviews",
split="train",
as_supervised=True,
batch_size=2,
)
# Drop labels.
imdb_train = imdb_train.map(lambda x, y: x)
imdb_train.unbatch().take(1).get_single_element().numpy()
Downloading and preparing dataset 80.23 MiB (download: 80.23 MiB, generated: Unknown size, total: 80.23 MiB) to /root/tensorflow_datasets/imdb_reviews/plain_text/1.0.0... Dl Completed...: 0 url [00:00, ? url/s] Dl Size...: 0 MiB [00:00, ? MiB/s] Generating splits...: 0%| | 0/3 [00:00<?, ? splits/s] Generating train examples...: 0%| | 0/25000 [00:00<?, ? examples/s] Shuffling /root/tensorflow_datasets/imdb_reviews/plain_text/1.0.0.incompleteAJDUZT/imdb_reviews-train.tfrecord… Generating test examples...: 0%| | 0/25000 [00:00<?, ? examples/s] Shuffling /root/tensorflow_datasets/imdb_reviews/plain_text/1.0.0.incompleteAJDUZT/imdb_reviews-test.tfrecord*… Generating unsupervised examples...: 0%| | 0/50000 [00:00<?, ? examples/s] Shuffling /root/tensorflow_datasets/imdb_reviews/plain_text/1.0.0.incompleteAJDUZT/imdb_reviews-unsupervised.t… Dataset imdb_reviews downloaded and prepared to /root/tensorflow_datasets/imdb_reviews/plain_text/1.0.0. Subsequent calls will reuse this data. b"This was an absolutely terrible movie. Don't be lured in by Christopher Walken or Michael Ironside. Both are great actors, but this must simply be their worst role in history. Even their great acting could not redeem this movie's ridiculous storyline. This movie is an early nineties US propaganda piece. The most pathetic scenes were those when the Columbian rebels were making their cases for revolutions. Maria Conchita Alonso appeared phony, and her pseudo-love affair with Walken was nothing but a pathetic emotional plug in a movie that was devoid of any real meaning. I am disappointed that there are movies like this, ruining actor's like Christopher Walken's good name. I could barely sit through it."
# Use a subset of the dataset for faster training.
imdb_train = imdb_train.take(2000)
Esegui l'ottimizzazione utilizzando l'adattamento a basso ranking (LoRA). LoRA è una tecnica di ottimizzazione che riduce notevolmente il numero di parametri addestrabili per le attività downstream bloccando i pesi completi del modello e inserendo nel modello un numero inferiore di nuovi pesi addestrabili. Fondamentalmente, LoRA riparametrizza le matrici a tutto peso più grandi con 2 matrici di basso livello più piccole AxB da addestrare e questa tecnica rende l'addestramento molto più veloce ed efficiente in termini di memoria.
# Enable LoRA for the model and set the LoRA rank to 4.
gemma_lm.backbone.enable_lora(rank=4)
# Fine-tune on the IMDb movie reviews dataset.
# Limit the input sequence length to 128 to control memory usage.
gemma_lm.preprocessor.sequence_length = 128
# Use AdamW (a common optimizer for transformer models).
optimizer = keras.optimizers.AdamW(
learning_rate=5e-5,
weight_decay=0.01,
)
# Exclude layernorm and bias terms from decay.
optimizer.exclude_from_weight_decay(var_names=["bias", "scale"])
gemma_lm.compile(
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=optimizer,
weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],
)
gemma_lm.summary()
gemma_lm.fit(imdb_train, epochs=1)
/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py:756: UserWarning: Some donated buffers were not usable: ShapedArray(float32[256000,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,384,256]), ShapedArray(float32[16,256,384]), ShapedArray(float32[384,24576]), ShapedArray(float32[384,24576]), ShapedArray(float32[24576,384]).
See an explanation at https://jax.readthedocs.io/en/latest/faq.html#buffer_donation.
warnings.warn("Some donated buffers were not usable:"
2000/2000 ━━━━━━━━━━━━━━━━━━━━ 358s 163ms/step - loss: 2.7145 - sparse_categorical_accuracy: 0.4329
<keras.src.callbacks.history.History at 0x7e9cac7f41c0>
Tieni presente che l'attivazione di LoRA riduce notevolmente il numero di parametri addestrabili, da 7 miliardi a soli 11 milioni.
Inferenza dopo l'ottimizzazione
gemma_lm.generate("Best comedy movies in the 90s ", max_length=64)
"Best comedy movies in the 90s \n\nThis is the movie that made me want to be a director. It's a great movie, and it's still funny today. The acting is superb, the writing is excellent, the music is perfect for the movie, and the story is great."
Dopo il perfezionamento, il modello ha appreso lo stile delle recensioni cinematografiche e ora sta generando output in questo stile nel contesto dei film comici degli anni '90.
Passaggi successivi
In questo tutorial hai imparato a utilizzare il backend JAX di KerasNLP per ottimizzare un modello Gemma sul set di dati IMDb in modo distribuito sulle potenti TPU. Ecco alcuni suggerimenti su cosa altro imparare:
- Scopri come iniziare a utilizzare Keras Gemma.
- Scopri come ottimizzare il modello Gemma sulla GPU.