 Посмотреть на ai.google.dev |  Запустить в Google Colab |  Бегать в Kaggle |  Открыть в Vertex AI |  Посмотреть исходный код на GitHub |
Обзор
Gemma — это семейство легких современных открытых моделей, созданных на основе исследований и технологий, используемых при создании моделей Google Gemini. Gemma можно дополнительно настроить в соответствии с конкретными потребностями. Но большие языковые модели, такие как Gemma, могут быть очень большими по размеру, и некоторые из них могут не поместиться в ускоритель сингла для точной настройки. В этом случае есть два общих подхода к их точной настройке:
- Эффективная точная настройка параметров (PEFT), целью которой является уменьшение эффективного размера модели за счет некоторой точности. LoRA попадает в эту категорию, и в учебнике «Точная настройка моделей Gemma в Keras с использованием LoRA» показано, как точно настроить модель Gemma 2B
gemma_2b_enс помощью LoRA с использованием KerasNLP на одном графическом процессоре. - Полная точная настройка параметров с параллелизмом моделей. Параллелизм моделей распределяет веса одной модели по нескольким устройствам и обеспечивает горизонтальное масштабирование. Подробнее о распределенном обучении вы можете узнать в этом руководстве по Keras .
В этом руководстве вы узнаете, как использовать Keras с серверной частью JAX для точной настройки модели Gemma 7B с помощью LoRA и распределенного обучения параллизму моделей на тензорном процессоре Google (TPU). Обратите внимание, что в этом руководстве LoRA можно отключить для более медленной, но более точной настройки полных параметров.
Использование ускорителей
Технически для этого урока вы можете использовать TPU или GPU.
Примечания по средам TPU
У Google есть 3 продукта, которые предоставляют TPU:
- Colab бесплатно предоставляет TPU v2, чего достаточно для этого руководства.
- Kaggle предлагает TPU v3 бесплатно, и они также подходят для этого руководства.
- Cloud TPU предлагает TPU v3 и более новые поколения. Один из способов его настройки:
- Создайте новую виртуальную машину TPU.
- Настройте переадресацию портов SSH для предполагаемого порта сервера Jupyter.
- Установите Jupyter и запустите его на виртуальной машине TPU, затем подключитесь к Colab через «Подключение к локальной среде выполнения».
Примечания по настройке нескольких графических процессоров
Хотя это руководство посвящено варианту использования TPU, вы можете легко адаптировать его для своих нужд, если у вас машина с несколькими графическими процессорами.
Если вы предпочитаете работать через Colab, также можно предоставить виртуальную машину с несколькими графическими процессорами для Colab напрямую через «Подключение к пользовательской виртуальной машине GCE» в меню Colab Connect.
Здесь мы сосредоточимся на использовании бесплатного TPU от Kaggle .
Прежде чем начать
Учетные данные Kaggle
Модели Джеммы размещаются на Kaggle. Чтобы использовать Gemma, запросите доступ на Kaggle:
- Войдите или зарегистрируйтесь на сайте kaggle.com.
- Откройте карточку модели Джеммы и выберите «Запросить доступ».
- Заполните форму согласия и примите условия использования
Затем, чтобы использовать API Kaggle, создайте токен API:
- Откройте настройки Kaggle
- Выберите «Создать новый токен».
- Будет загружен файл
kaggle.json. Он содержит ваши учетные данные Kaggle.
Запустите следующую ячейку и введите свои учетные данные 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…
Альтернативный способ — установить KAGGLE_USERNAME и KAGGLE_KEY в вашей среде, если kagglehub.login() у вас не работает.
Установка
Установите Keras и KerasNLP с моделью 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
Настройте бэкэнд Keras JAX
Импортируйте JAX и запустите проверку работоспособности TPU. Kaggle предлагает устройства TPUv3-8 с 8 ядрами TPU по 16 ГБ памяти каждое.
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"
Загрузить модель
import keras
import keras_nlp
Примечания по обучению смешанной точности на графических процессорах NVIDIA
При обучении на графических процессорах NVIDIA смешанная точность ( keras.mixed_precision.set_global_policy('mixed_bfloat16') ) может использоваться для ускорения обучения с минимальным влиянием на качество обучения. В большинстве случаев рекомендуется включать смешанную точность, поскольку это экономит и память, и время. Однако имейте в виду, что при небольших размерах пакетов это может увеличить использование памяти в 1,5 раза (веса будут загружены дважды, с половинной и полной точностью).
Для вывода, половинная точность ( keras.config.set_floatx("bfloat16") ) будет работать и экономить память, тогда как смешанная точность неприменима.
# Uncomment the line below if you want to enable mixed precision training on GPUs
# keras.mixed_precision.set_global_policy('mixed_bfloat16')
Чтобы загрузить модель с весами и тензорами, распределенными по TPU, сначала создайте новый DeviceMesh . DeviceMesh представляет собой набор аппаратных устройств, настроенных для распределенных вычислений, и был представлен в Keras 3 как часть API унифицированного распространения.
API распространения обеспечивает параллелизм данных и моделей, что позволяет эффективно масштабировать модели глубокого обучения на нескольких ускорителях и хостах. Он использует базовую структуру (например, JAX) для распределения программы и тензоров в соответствии с директивами сегментирования с помощью процедуры, называемой расширением одной программы, нескольких данных (SPMD). Более подробную информацию можно найти в новом руководстве по API распространения 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 из API распространения определяет, как веса и тензоры должны быть сегментированы или реплицированы, используя строковые ключи, например token_embedding/embeddings ниже, которые обрабатываются как регулярное выражение для соответствия тензорным путям. Соответствующие тензоры сегментируются по размерам модели (8 TPU); другие будут полностью воспроизведены.
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 позволяет сегментировать веса моделей или тензоры активации на всех устройствах в DeviceMesh . В этом случае некоторые из весов модели Gemma 7B распределяются по 8 чипам TPU в соответствии с layout_map определенным выше. Теперь загрузим модель распределенным способом.
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.
Теперь убедитесь, что модель правильно разделена. Давайте возьмем decoder_block_1 в качестве примера.
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')
Выводы перед точной настройкой
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)'
Модель генерирует список отличных комедий 90-х годов, которые стоит посмотреть. Теперь мы настраиваем модель Gemma, чтобы изменить стиль вывода.
Точная настройка с помощью 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)
Выполните точную настройку с помощью адаптации низкого ранга (LoRA). LoRA — это метод тонкой настройки, который значительно сокращает количество обучаемых параметров для последующих задач за счет замораживания полных весов модели и добавления в модель меньшего количества новых обучаемых весов. По сути, LoRA перепараметризует большие полновесовые матрицы с помощью двух меньших матриц низкого ранга AxB для обучения, и этот метод делает обучение намного быстрее и эффективнее с точки зрения памяти.
# 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>
Обратите внимание, что включение LoRA значительно сокращает количество обучаемых параметров: с 7 миллиардов до всего 11 миллионов.
Вывод после точной настройки
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."
После тонкой настройки модель освоила стиль обзоров фильмов и теперь генерирует результаты в этом стиле в контексте комедийных фильмов 90-х годов.
Что дальше
В этом руководстве вы узнали, как использовать JAX-сервер KerasNLP для точной настройки модели Gemma в наборе данных IMDb распределенным образом на мощных TPU. Вот несколько советов о том, чему еще стоит научиться:
- Узнайте, как начать работу с Keras Gemma .
- Узнайте, как точно настроить модель Gemma на графическом процессоре .