 Ver en ai.google.dev
|
 Ejecutar en Google Colab
|
 Ejecutar en Kaggle
|
 Abrir en Vertex AI
|
 Ver código fuente en GitHub
|
Descripción general
Gemma es una familia de modelos abiertos, ligeros y de vanguardia creados a partir de la investigación y la tecnología que se usaron para crear los modelos de Google Gemini. Gemma se puede ajustar aún más para satisfacer necesidades específicas. Sin embargo, los modelos de lenguaje extenso, como Gemma, pueden ser muy grandes y es posible que algunos no se ajusten a un acelerador único para el perfeccionamiento. En este caso, existen dos enfoques generales para ajustarlos:
- El ajuste fino eficiente en relación con los parámetros (PEFT), que busca reducir el tamaño efectivo del modelo sacrificando cierta fidelidad. LoRA se encuentra en esta categoría y el instructivo Ajusta los modelos de Gemma en Keras con LoRA muestra cómo ajustar el modelo Gemma 2B
gemma_2b_encon LoRA usando KerasNLP en una sola GPU. - Ajusta los parámetros por completo con el paralelismo de modelos. El paralelismo de modelos distribuye los pesos de un solo modelo en varios dispositivos y habilita la escala horizontal. Puedes obtener más información sobre el entrenamiento distribuido en esta guía de Keras.
En este instructivo, se explica el uso de Keras con un backend de JAX para ajustar el modelo de Gemma 7B con LoRA y el entrenamiento distribuido de paralismo de modelos en la unidad de procesamiento tensorial (TPU) de Google. Ten en cuenta que se puede desactivar LoRA en este instructivo para obtener un ajuste de parámetros completo más lento, pero más preciso.
Usa aceleradores
Técnicamente, puedes usar TPU o GPU para este instructivo.
Notas sobre los entornos de TPU
Google tiene 3 productos que proporcionan TPU:
- Colab proporciona TPU v2 de forma gratuita, lo que es suficiente para este instructivo.
- Kaggle ofrece TPU v3 de forma gratuita, y también funcionan para este instructivo.
- Cloud TPU ofrece TPU v3 y generaciones posteriores. Una forma de configurarlo es la siguiente:
- Crea una nueva VM de TPU
- Configura la redirección de puertos SSH para el puerto del servidor de Jupyter que deseas
- Instala Jupyter y, luego, inícialo en la VM de TPU. Luego, conéctate a Colab a través de "Conectar a un entorno de ejecución local".
Notas sobre la configuración de varias GPU
Si bien este instructivo se enfoca en el caso de uso de TPU, puedes adaptarlo fácilmente para tus propias necesidades si tienes una máquina con varias GPU.
Si prefieres trabajar con Colab, también puedes aprovisionar una VM de varias GPU para Colab directamente a través de "Conéctate a una VM de GCE personalizada" en el menú de Colab Connect.
Aquí, nos enfocaremos en usar la TPU gratuita de Kaggle.
Antes de comenzar
Credenciales de Kaggle
Kaggle aloja los modelos de Gemma. Para usar Gemma, solicita acceso en Kaggle:
- Accede o regístrate en kaggle.com.
- Abre la tarjeta de modelo Gemma y selecciona "Solicitar acceso".
- Completa el formulario de consentimiento y acepta los Términos y Condiciones.
Luego, para usar la API de Kaggle, crea un token de API:
- Abre la configuración de Kaggle.
- Selecciona "Create New Token".
- Se descargó un archivo
kaggle.json. Contiene tus credenciales de Kaggle.
Ejecuta la siguiente celda y, cuando se te solicite, ingresa tus credenciales de Kaggle.
# 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…
Una forma alternativa es configurar KAGGLE_USERNAME y KAGGLE_KEY en tu entorno si kagglehub.login() no funciona.
Instalación
Instalar Keras y KerasNLP con el modelo de 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 el backend de Keras JAX
Importa JAX y ejecuta una verificación de coherencia en TPU. Kaggle ofrece dispositivos TPUv3-8 que tienen 8 núcleos de TPU con 16 GB de memoria cada uno.
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"
Carga el modelo
import keras
import keras_nlp
Notas sobre el entrenamiento de precisión mixta en GPUs de NVIDIA
Cuando se entrena en GPUs de NVIDIA, se puede usar la precisión mixta (keras.mixed_precision.set_global_policy('mixed_bfloat16')) para acelerar el entrenamiento con un efecto mínimo en la calidad del entrenamiento. En la mayoría de los casos, se recomienda activar la precisión mixta, ya que ahorra memoria y tiempo. Sin embargo, ten en cuenta que, con tamaños de lotes pequeños, puede aumentar el uso de la memoria en 1.5 veces (las ponderaciones se cargarán dos veces, con precisión media y precisión completa).
Para la inferencia, la precisión media (keras.config.set_floatx("bfloat16")) funcionará y ahorrará memoria mientras la precisión mixta no es aplicable.
# Uncomment the line below if you want to enable mixed precision training on GPUs
# keras.mixed_precision.set_global_policy('mixed_bfloat16')
Para cargar el modelo con las ponderaciones y los tensores distribuidos en las TPU, primero crea un DeviceMesh nuevo. DeviceMesh representa una colección de dispositivos de hardware configurados para el procesamiento distribuido y se introdujo en Keras 3 como parte de la API de distribución unificada.
La API de distribución habilita el paralelismo de datos y modelos, lo que permite escalar de manera eficiente los modelos de aprendizaje profundo en varios aceleradores y hosts. Aprovecha el framework subyacente (por ejemplo, JAX) para distribuir el programa y los tensores de acuerdo con las directivas de fragmentación a través de un procedimiento llamado expansión de programa único, varios datos (SPMD). Obtén más detalles en la nueva guía de la API de distribución de 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 de la API de distribución especifica cómo se deben particionar o replicar las ponderaciones y los tensores, con las claves de cadena, por ejemplo, token_embedding/embeddings a continuación, que se tratan como regex para hacer coincidir las rutas de acceso de tensor. Los tensores coincidentes se fragmentan con las dimensiones del modelo (8 TPU); otros se replicarán por completo.
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 te permite dividir los pesos del modelo o los tensores de activación en todos los dispositivos del DeviceMesh. En este caso, algunos de los pesos del modelo Gemma 7B se dividen en 8 chips de TPU según el layout_map definido anteriormente. Ahora carga el modelo de forma distribuida.
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.
Ahora, verifica que el modelo se haya particionado correctamente. Tomemos decoder_block_1 como ejemplo.
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')
Inferencia antes del perfeccionamiento
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)'
El modelo genera una lista de grandes películas de comedia de los 90 para mirar. Ahora, ajustaremos el modelo Gemma para cambiar el estilo de salida.
Optimiza 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)
Realiza el ajuste fino con la adaptación de bajo rango (LoRA). LoRA es una técnica de ajuste fino que reduce en gran medida la cantidad de parámetros entrenables para las tareas posteriores, ya que inmoviliza los pesos completos del modelo y, luego, inserta una cantidad menor de pesos entrenables nuevos en él. Básicamente, LoRA reparamatriza las matrices de pesos completas más grandes por 2 matrices AxB de rango bajo más pequeñas para entrenar, y esta técnica hace que el entrenamiento sea mucho más rápido y eficiente en la 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>
Tenga en cuenta que habilitar LoRA reduce de manera significativa la cantidad de parámetros entrenables, de 7,000 millones a solo 11 millones.
Inferencia después del perfeccionamiento
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."
Después de ajustarlo, el modelo aprendió el estilo de las opiniones sobre películas y ahora genera resultados en ese estilo en el contexto de las películas de comedia de los 90.
¿Qué sigue?
En este instructivo, aprendiste a usar el backend de JAX de KerasNLP para ajustar un modelo de Gemma en el conjunto de datos de IMDb de forma distribuida en las potentes TPU. A continuación, te damos algunas sugerencias sobre qué más puedes aprender:
- Obtén más información para comenzar a usar Keras Gemma.
- Obtén información para ajustar el modelo Gemma en la GPU.