使用 TPU 微调并运行 Google Gemma 模型推理，全面提升速度与性能

本文是一份关于使用高级技术微调大型语言模型的完整教程。我们将通过示例使用张量处理单元（TPU）、一种名为 LoRA 的技术以及分布式计算——一切都是为了加速与高效。为演示这些技术，我们将访问并微调 Gemma 模型，这是 Google 推出的新一代轻量级、最先进的大语言模型系列。

在本教程结束时，您将掌握使用 Google Cloud 上可用的 TPU 对任意大型语言模型进行微调与推理的技能。为此，我们将按以下顺序覆盖这些主题：

• 计算类型简介：了解 TPU 以及它们为何重要。

• 在 TPU 上使用 Gemma：学习如何在 Kaggle 环境中配置并使用 TPU。

• 模型设置与推理：使用 Python 中的 Keras 库运行 Gemma。

• 微调：使用名为 LoRA 的技术对 Gemma 模型进行微调。

• 分布式训练：进行分布式微调以提升训练效率。

如果您是 AI 和大语言模型的新手，建议先学习全面的 AI Fundamentals 技能路径，以熟悉本教程中使用的术语。

开始吧！

什么是 Google 的 Gemma 模型？

Google Gemma 标志

Google 的 Gemma 模型属于 Google 于 2024 年推出的开源且轻量级的大语言模型家族。Gemma 基于与 Google Gemini 模型相同的研究与技术构建，并且与 Gemini 类似，Gemma 兼容主流机器学习框架，如 Keras 和 Pytorch。它提供两种规模，包含 20 亿或 70 亿参数。

与主要面向终端用户、通过应用与 API 提供的 Gemini 不同，Gemma 是开源的，便于开发者广泛修改与集成。Gemma 模型也更小，更便携且更具成本效益。

什么是张量处理单元（TPU）？

CPU、GPU、TPU 图片来源 Freepik 与 Flaticon

为便于理解，我们先比较一下机器学习的不同硬件选项。硬件可分为不同类型，每一类都针对特定的处理与计算角色而设计。

• 中央处理器：CPU 运行几乎所有计算设备中的操作系统，擅长顺序处理任务。

• 图形处理器：GPU 可同时处理多项任务，适合图形渲染。

• 张量处理单元：TPU 是 Google 为机器学习开发的专用处理器，专为快速执行训练与运行神经网络所需的矩阵计算而设计。

对本教程而言，最后一类 TPU 尤为重要。TPU 在深度神经网络的训练与推理上通常比 GPU 更快，且功耗更低。缺点是 TPU 生态较不成熟，可用的工具与框架更少。当前可用的框架包括 Google Cloud Platform、Colab 和 Kaggle。

机器学习硬件对比表

使用 TPU 访问 Google Gemma

Gemma 模型在分布式计算环境中具备高度可扩展性与高效性，面对海量数据与复杂模型架构时，这能显著提升其性能与速度。

Keras 库在此场景下表现良好，因为它通过多后端实现，支持 Gemma 模型的分布式训练，包括 TensorFlow 与 PyTorch。

本节将配置 Kaggle 笔记本以支持 TPU，确保平稳运行所需的库与依赖均已安装。随后，我们将使用 keras-nlp 包加载 Gemma 2B 或 7B 模型。最后，我们将部署模型执行任务，利用 TPU 强大的计算能力高效处理与生成响应。

环境准备

为准备环境，我们需要集成 Gemma 模型的 Keras 版本，并将计算加速器切换为 TPU 以获得更佳性能。我们将安装 keras-nlp 等关键包，配置 Keras 使用 TensorFlow 作为后端，并优化 TPU 内存分配，以确保模型训练与执行顺畅。

为 Gemma 模型添加 Keras 实现

为准备环境，我们需要为 Gemma 模型添加 Keras 实现。具体而言，即引入与 Keras 兼容的 Gemma 模型版本，以确保其能平滑集成 TensorFlow API。

添加 Keras 实现，按以下步骤进行：

• 在 Kaggle 中创建新笔记本。

• 在右侧面板转到 "Input" 区域。

• 点击 “+ Add Input” 以添加 Gemma 模型的 Keras 实现。

将 Keras 实现添加到 Gemma 模型

将加速器更改为 TPU

接下来需要将加速器更改为 TPU，也就是在笔记本设置中切换计算加速器为 TPU，以在训练与推理中发挥其算力。

更改加速器按以下步骤进行。环境重载可能需要几分钟，属正常现象。

• 进入 "Session options" 区域。
• 将加速器从 "None" 更改为 "TPU VM v3-8"。

将 TPU 设置为加速器

安装微调与推理所需的 Python 包

现在安装用于微调与推理的新版 Python 包，包括 TensorFlow 与 keras-nlp。

!pip install -q tensorflow-cpu!pip install -q -U keras-nlp tensorflow-hub!pip install -q -U keras==3.1.1技术支持

设置 Keras 后端

我们将 Keras 后端设置为 JAX，以便无缝访问 TPU。您也可以尝试 TensorFlow 或 PyTorch。

预分配 TPU 内存

最后，我们预分配 TPU 内存以优化性能，并在模型执行时避免运行时问题。这里在 JAX 后端预分配 100% 的 TPU 总内存，以避免内存碎片。

import osos.environ["KERAS_BACKEND"] = "jax"os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"]="1.00"技术支持

加载模型

加载用于模型推理的必要 Python 包。

import kerasimport keras_nlp技术支持

然后加载模型并显示模型摘要。

gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset("gemma_instruct_2b_en")gemma_lm.summary()技术支持

我们仅用一条命令就成功加载了分词器和模型。这非常令人惊叹，因为即便是最小的 Gemma 模型也有 25 亿参数，大小 9.34 GB。

Gemma 模型摘要

您可能需要等待几分钟才能运行笔记本并体验超快性能。这是因为 TPU 资源需求旺盛。若在夜间运行，可能无需等待。

排队等待 TPU 访问

模型推理

运行推理指使用模型基于新的、未见过的输入生成输出。在本例中，我们将使用 .generate() 函数以文本形式提供输入。响应会立即返回。

print(gemma_lm.generate("What is DataCamp?", max_length=30))技术支持

What is DataCamp?DataCamp is a leading online data science education platform that empowers individuals and organizations to build data-driven careers. They技术支持

我们也可以尝试批量推理，即提供多个提示以生成多条响应。输出将是列表形式。

print(gemma_lm.generate(["How far is the Sun from earth", "What is the Sun made of?"], max_length=30))技术支持

['How far is the Sun from earth?\n\nThe Sun is about 149.6 million kilometers (93 million miles) away from', 'What is the Sun made of?\n\nThe Sun is primarily composed of hydrogen and helium. Hydrogen makes up about 73.4% of']技术支持

接下来，我们将尝试创建一个包含指令与响应的模板以引导模型。您可以添加系统提示或任何初始命令以调整生成的响应。

我们将使用 .format() 函数将用户参数填充到模板文本中。这样，我们就得到了一份学习 Python 的步骤清单。

template = "Instruction:\n{instruction}\n\nResponse:\n{response}"prompt = template.format(    instruction="How do I start learning Python?",    response="",)print(gemma_lm.generate(prompt, max_length=250))技术支持

Instruction:How do I start learning Python?Response:**Step 1: Choose a learning path*** **Online Courses:**    * DataCamp    * Codecademy    * Coursera    * edX    * Udemy* **Books:**    * "Automate the Boring Stuff with Python" by Al Sweigart    * "Python Crash Course" by Eric Matthes    * "Head First Python" by Kathy Sierra and Bert Bates* **Video Tutorials.........技术支持

如果您在使用 TPU 加载模型并生成响应时遇到问题，可参考此 Kaggle 笔记本：Accessing Gemma-instruct-2b-en using TPUs。

如果您已开始熟悉 Gemma，想进一步了解如何通过自定义指令增强 Gemma，请查看我们的Fine Tuning Google Gemma: Enhancing LLMs with Customized Instructions 教程。

使用 TPU 微调 Gemma

现在我们来学习如何使用名为 OpenHermes 的公开数据集微调 Gemma 模型。OpenHermes 数据集包含 242,000 行，两列分别为指令与响应，均由 GPT-4 生成。

正如我们将看到的，即便使用 TPU，微调该数据集也将耗时很长。为提效，我们将采用一种名为 LoRA 的技术。

LoRA（低秩适配）旨在让大型语言模型的微调过程更高效、更易用。它解决了传统微调方法常常计算开销大、资源密集的问题。

在技术层面，LoRA 通过冻结模型的预训练权重，并引入低秩矩阵来修改模型架构中的特定部分。通过聚焦这些低秩矩阵，LoRA 将所需的计算资源最小化，使在较弱硬件上微调大模型成为可能。

环境准备

与之前的模型环境设置类似，我们将设置 Keras 后端并进行内存预分配。我们还可以使用 jax. devices() 函数检查 TPU 设备的可用性。

import osimport jaxos.environ["KERAS_BACKEND"] = "jax"os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.9"jax.devices()技术支持

计算设备数量

加载模型与数据集

加载模型并显示模型摘要。我们有 25 亿个可训练参数。

import kerasimport keras_nlpgemma_lm = keras_nlp.models.GemmaCausalLM.from_preset("gemma_2b_en")gemma_lm.summary()技术支持

Gemma 模型摘要

接着，在笔记本中加载数据集。该过程与添加模型类似。按以下步骤进行：

• 前往 "Add input" 区域。
• 搜索 OpenHermes 数据集。
• 添加数据集。为明确起见，我们选择由 Volodymyr Pivoshenko 托管的版本。

添加 OpenHermes 数据集

现在使用 pandas 读取数据集并显示前 5 行。可以看到共有两列：一列为指令，另一列为输出。

import pandas as pddf = pd.read_csv('/kaggle/input/openhermes/openhermes.csv')df.head()技术支持

OpenHermes 数据集

现在模型和数据集都已加载，我们需要将数据集传递给模型。为此，必须把数据集转换为由字符串组成的列表，格式为指令与响应。

事实证明，训练整个数据集即便在 TPU 上也将耗时近 23 小时。因此，我们仅选择列表中的前 1,000 个样本以缩短训练时间。

template = "Instruction:\n{instruction}\n\nResponse:\n{output}"data = [template.format(**row) for index, row in df.iterrows()]技术支持

data = data[:1000]print(data[0])技术支持

OpenHermes 数据集首条样本

微调前的推理

首先建立一个基线，以便与使用 LoRA 的效果进行对比。为此，我们将使用 template.format() 生成响应。

仔细阅读可发现，下方输出显示基线响应细节不多。Gemma 甚至在第一段末尾开始重复自身。这并非理想结果，但很好地说明了我们需要微调模型以提供更优答案。

prompt = template.format(    instruction="Plan a 5-day Bahamas trip.",    output="",)print(gemma_lm.generate(prompt, max_length=256))技术支持

Instruction:Plan a 5-day Bahamas trip.Response:Day 1:Fly to Nassau, Bahamas.Visit the Atlantis Resort.Day 2:Visit the Exuma Cays.Day 3:Visit the Lucayan National Park.Day 4:Visit the Great Abaco Islands.Day 5:Fly home.Instruction:Plan a 5-day Bahamas trip.Response:Day 1:Fly to Nassau, Bahamas.Visit the Atlantis Resort.Day 2:Visit the Exuma Cays.Day 3:Visit the Lucayan National Park.Day 4:Visit the Great Abaco Islands.Day 5:Fly home.Instruction:Plan a 5-day Bahamas trip.Response:Day 1:Fly to Nassau, Bahamas.Visit the Atlantis Resort.Day 2:Visit the Exuma Cays.Day 3:Visit the Lucayan National Park.Day 4:Visit the Great Abaco Islands.Day 5:Fly home.Instruction:Plan a 5-day技术支持

为微调编译模型

现在使用 LoRA 改进模型，它能减少下游任务的可训练参数数量。

在示例中，我们将 LoRA 等级设为 4，这是计算上最节省的较小等级。如果您想提升模型性能，可以选择 LoRA 等级 8 或 16。参阅 Fine-Tuning LLaMA 2 指南以进一步了解 LoRA 与量化。

随后指定优化器、损失函数与准确率指标。

gemma_lm.backbone.enable_lora(rank=4)技术支持

gemma_lm.preprocessor.sequence_length = 512optimizer = keras.optimizers.AdamW(    learning_rate=5e-5,    weight_decay=0.01,)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()技术支持

我们的可训练参数已从 25 亿大幅减少至 130 万。且适配层大小仅为 5.2 MB。

Gemma 模型：总参数 vs 可训练参数

训练模型

现在将数据拟合到模型上。这里选择 1 个 epoch，batch size 为 1。若想提升模型性能，可在完整数据集上训练至少 5 个 epoch。

gemma_lm.fit(data, epochs=1, batch_size=1)技术支持

微调后的推理

现在使用与基线实验相同的提示，看看模型是否改进。这一次，模型不再是一两句并出现重复，而是给出了详尽的巴哈马 5 日行程规划。

prompt = template.format(    instruction="Plan a 5-day Bahamas trip.",    output="",)print(gemma_lm.generate(prompt, max_length=256))技术支持

Instruction:Plan a 5-day Bahamas trip.Response:Day 1: Arrive in Nassau, BahamasUpon arrival in Nassau, you will be greeted by a local guide who will take you on a tour of the city. Visit the famous Straw Market, where you can shop for souvenirs and local crafts. Enjoy lunch at a local restaurant before heading to the Atlantis Resort for a day of fun and relaxation. Spend the afternoon exploring the resort's world-class amenities, including the Aquaventure Water Park, Dolphin Cay, and the massive marine habitat, The Dig.Day 2: Explore Paradise IslandOn day 2, you will explore Paradise Island, home to the world-famous Atlantis Resort. Visit the famous Atlantis Casino and enjoy a round of blackjack or roulette. Take a stroll along the beach and enjoy the crystal-clear waters of the Atlantic Ocean. Visit the Dolphin Cay, where you can swim with dolphins and other marine animals.Day 3: Snorkel at Cabbage BeachOn day 3, you will head to Cabbage Beach, a secluded beach located on Paradise Island. Snorkel in the crystal-clear waters and explore the coral reefs. Enjoy a delicious lunch at a local restaurant before heading back to Nassau.技术支持

您可以保存模型权重以便推理，并将模型部署到生产环境。

gemma_lm.save_weights('gemma_2b_openhermes.weights.h5')技术支持

我们微调后的模型总大小为 10.04GB。

Gemma 微调并保存的模型文件

如果您在微调时遇到问题，可参考此 Kaggle 笔记本：Finetuning Gemma using TPUs

使用 TPU 对 Gemma 进行分布式微调与推理

最后一节我们将深入了解分布式微调。分布式微调通过模型并行实现，它将单个模型的权重分布到多个设备上，从而实现横向扩展并加速训练。

在这一部分，我们将显著缩短大型模型的训练时间。具体来说，我们将使用以 JAX 为后端的 Keras，在 TPU 上结合 LoRA 与分布式训练微调 Gemma。

环境准备

此处新建一个 Kaggle 笔记本并安装必要的库。随后设置 Keras 后端并预分配内存。

!pip install -q tensorflow-cpu!pip install -q -U keras-nlp tensorflow-hub!pip install -q -U keras==3.1.1技术支持

import osos.environ["KERAS_BACKEND"] = "jax"os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.9"技术支持

为 8 个 TPU 设置分布

首先创建一个 DeviceMesh，帮助我们将模型权重与张量分布在多个 TPU 设备上。它支持数据并行与模型并行，使 LLM 能在多个加速器上高效扩展。

创建形状为 (1, 8) 的 DeviceMesh，使模型权重在 8 个 TPU 上分片。

import kerasimport keras_nlpdevice_mesh = keras.distribution.DeviceMesh(    (1, 8), ["batch", "model"], devices=keras.distribution.list_devices())技术支持

接下来，使用正则表达式创建一个布局映射，指定权重与张量应如何分片或复制。

请注意以下要点：

• 匹配 token_embedding/embeddings 的权重将被共享。

• 使用正则匹配解码器 attention、attention_output、ffw_gating 与 ffw_linear 层中的 query、key、value 矩阵。

• 匹配到的张量将按 DeviceMesh 分片，其余将完全复制。

model_dim = "model"layout_map = keras.distribution.LayoutMap(device_mesh)layout_map["token_embedding/embeddings"] = (None, model_dim)layout_map["decoder_block.*attention.*(query|key|value).*kernel"] = (    None,    model_dim,    None,)layout_map["decoder_block.*attention_output.*kernel"] = (None, None, model_dim)layout_map["decoder_block.*ffw_gating.*kernel"] = (model_dim, None)layout_map["decoder_block.*ffw_linear.*kernel"] = (None, model_dim)技术支持

现在使用 device_mesh 与 layout_map 启用模型并行。

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")gemma_lm.summary()技术支持

这一次，因为使用的是 Gemma 7B 模型，我们看到有 85 亿个可训练参数。

Gemma 7B 模型摘要

为验证模型是否正确分片，我们将打印 decoder_block_1 层权重的 path、shape 与 spec。

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)}')技术支持

解码器块

加载数据集

加载 OpenHermes 数据集。依旧将数据框转换为使用指令与响应模板的字符串列表，并为加快训练仅选择前 1,000 个样本。

import pandas as pddf = pd.read_csv('/kaggle/input/openhermes/openhermes.csv')# Format and convert the dataframe into a list of strings.template = "Instruction:\n{instruction}\n\nResponse:\n{output}"data = [template.format(**row) for index, row in df.iterrows()]# Select a subset of the dataset. data = data[:1000]技术支持

微调前的推理

我们将通过提供格式化提示来生成基线响应。有趣的是，这次的响应甚至比小模型更差，但我们继续。

prompt = template.format(    instruction="Plan a 5-day Bahamas trip.",    output="",)print(gemma_lm.generate(prompt, max_length=256))技术支持

 微调前的推理结果

为微调编译模型

我们将使用与之前相同的模型配置、LoRA 等级、优化器、损失函数与准确率指标来编译模型。现在有 1100 万个可训练参数，大小为 42.22 MB。

gemma_lm.backbone.enable_lora(rank=4)gemma_lm.preprocessor.sequence_length = 512optimizer = keras.optimizers.AdamW(    learning_rate=5e-5,    weight_decay=0.01,)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 7B LoRA 模型摘要

训练模型

微调 1 个 epoch 用时 305 秒，考虑到模型规模，这是很不错的表现。

gemma_lm.fit(data, epochs=1, batch_size=1)技术支持

如果您觉得微调有难度，可以学习使用 OpenAI 的 API，并参考这份逐步微调 GPT-4 的指南。

微调后的推理

生成响应并与基线结果对比。

prompt = template.format(    instruction="Plan a 5-day Bahamas trip.",    output="",)print(gemma_lm.generate(prompt, max_length=256))技术支持

微调后的模型表现理想。它生成了涵盖五天巴哈马行程安排的清单式响应。

注意：您可以关闭 LoRA，配合模型并行进行全参数微调，速度更慢但精度更高。

微调后的推理结果

gemma_lm.save_weights('gemma_7b_openhermes.weights.h5')技术支持

如果您在分布式微调方面遇到困难，也可参考这份 Kaggle 笔记本作为补充资源：Accessing Gemma-instruct-2b-en using TPUs。

结语

在本教程中，我们了解了 TPU，以及如何使用它们加速大语言模型的响应生成。我们还学习了如何在 TPU 与分布式训练的加持下，使用 OpenHermes 数据集微调 Gemma 模型。

在 TPU 上微调 LLM，并使用模型并行进行分布式学习，是缩短训练时间的重要手段。通过分布式学习，您甚至可以微调 Gemma 7B 这样的大模型。借助专为高性能机器学习任务设计的 TPU，您可在训练与推理阶段都获得显著加速。

如果您喜欢本教程并想进一步探索大语言模型的世界，请学习Master Large Language Models (LLMs) Concepts 课程，了解 LLM 的应用、训练方法、伦理考量与最新研究。如果您对与 Gemma 竞争的其他大语言模型感兴趣，请查看我们的Getting Started With Mixtral 8X22B 教程，进一步了解 Mistral AI 的新模型及其稀疏专家混合（SMoE）架构。