 前往 ai.google.dev 查看
|
 在 Kaggle 中執行 |
 在 Vertex AI 中開啟
|
 在 GitHub 上查看原始碼 |
總覽
Gemma 是一系列最先進的開放式模型，以研究及技術建立 Google Gemini 模型。Gemma 可以進一步微調，以滿足特定需求。但 Gemma 等大型語言模型的大小非常大，其中部分模型可能不適合用來進行微調。在這種情況下，一般的微調方法有兩種：
- 「有效微調 (PEFT)」：藉由犧牲部分擬真度，來縮減有效模型大小。LoRA 屬於這個類別，使用 LoRA 的 Keras 微調 Gemma 模型教學課程示範如何在單一 GPU 上使用 KerasNLP，透過 LoRA 微調 Gemma 2B 模型
gemma_2b_en。 - 使用模型平行處理功能微調完整參數。模型平行處理功能可將單一模型的權重分配給多部裝置，並啟用水平資源調度。如要進一步瞭解分散式訓練，請參閱 Keras 指南。
本教學課程將逐步引導您使用 Keras 搭配 JAX 後端，以 LoRA 和 Google Tensor Processing Unit (TPU) 上的模型參數分散式訓練，微調 Gemma 7B 模型。請注意，在本教學課程中，您可以關閉 LoRA，以加快但更準確的完整參數調整作業。
使用加速器
技術上來說，您可以選用 TPU 或 GPU 進行這個教學課程。
TPU 環境注意事項
Google 有 3 項提供 TPU 的產品：
多 GPU 設定注意事項
雖然本教學課程著重於 TPU 用途，但如果您有多重 GPU 機器，也可以根據個人需求輕鬆調整設定。
如果你想使用 Colab，也可以透過 Colab Connect 選單中的「連線至自訂的 GCE VM」直接為 Colab 佈建多 GPU VM。
這裡會著重於使用 Kaggle 提供的免費 TPU。
事前準備
Kaggle 憑證
Gemma 模型由 Kaggle 託管。如要使用 Gemma，請在 Kaggle 上要求存取權：
- 登入或註冊 kaggle.com
- 開啟 Gemma 模型資訊卡，然後選取「要求存取權」
- 填妥同意書並接受條款及細則
接著，如要使用 Kaggle API，請建立 API 權杖：
- 開啟 Kaggle 設定
- 選取「Create New Token」(建立新權杖)
- 系統會下載
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…
另一種方法是如果 kagglehub.login() 無法運作，就在您的環境中設定 KAGGLE_USERNAME 和 KAGGLE_KEY。
安裝
使用 Gemma 模型安裝 Keras 和 KerasNLP。
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# 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 GB 記憶體。
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 GPU 的混合精確度訓練注意事項
在 NVIDIA GPU 上進行訓練時，可使用混合精確度 (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 代表設定用於分散式運算的一組硬體裝置，並做為統合發行 API 中在 Keras 3 中導入的。
發布 API 讓資料和模型平行處理，可在多個加速器和主機上有效率地擴充深度學習模型。這個平台運用基礎架構 (例如 JAX)，根據資料分割指令發布程式和張量，整個程序稱為「單一程式，擴充多個資料 (SPMD)」。詳情請參閱新的 Keras 3 發布 API 指南。
# 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())
Distribution API 中的 LayoutMap 會指定使用字串鍵分割或複製權重和張量的方式。例如，下方的 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 模型權重是根據上述定義的 layout_map，在 8 個 TPU 晶片之間進行資料分割。現在，以分散的方式載入模型。
model_parallel = keras.distribution.ModelParallel(
device_mesh, 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 會將較大的完整矩陣重新參數化，依照 2 個較小的低排名矩陣 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 後，可訓練的參數數量大幅減少，從 70 億次提高為 1, 100 萬組。
調整後推論
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 年代喜劇片產生該風格的輸出內容。
後續步驟
在這個教學課程中，您學到瞭如何使用 KerasNLP JAX 後端，在強大的 TPU 上以分散式方式微調 IMDb 資料集上的 Gemma 模型。以下幾點建議還有其他學習問題：
- 瞭解如何開始使用 Keras Gemma。
- 瞭解如何在 GPU 上調整 Gemma 模型。