# 使用 Keras 进行 Gemma 模型的分布式微调和推理 Gemma 是 Google 基于 Gemini 研究和技术构建的轻量级开放模型家族，虽然小巧但性能不俗。对于很多实际场景，我们往往需要对这类大型语言模型进行微调，让它们更贴合自己的业务需求。但问题来了——像 Gemma 7B 这样规模的模型，单个翻跟斗根本塞不下，更别提训练了。 那么，面对这种 "模型太大、设备太少" 的尴尬，有两条路可以走： 1. **参数高效微调（PEFT）**，比如 LoRA。这类方法通过缩小有效模型尺寸来节省资源，代价是部分保真度。之前在 Keras 中如何使用 LoRA 微调 Gemma 2B 模型，已经有教程演示过——本质上就是牺牲一点精度，换来能在单张 GPU 上跑起来的机会。 2. **全参数微调 + 模型并行**。把一个大模型的权重拆开，分散到多个设备上协同训练。这就引出了分布式训练的话题，Keras 官方指南里有详细的介绍。 今天这篇文章，我们就来手把手走一遍：如何利用 Keras 搭配 JAX 后端，通过 LoRA + 模型并行分布式训练，在 Google 的 TPU 上搞定 Gemma 7B 模型的微调。另外，如果你真想玩全参数微调，LoRA 那一层也可以关掉——不过时间嘛，自然会长一些。 ## 使用翻跟斗 从硬件上来说，TPU 和 GPU 都能跑这套流程，不过这里有个细节得提醒一下。 ### 关于 TPU 环境的说明 Google 目前提供了三种 TPU 产品，但不是每种都适合这个任务： - Colab 提供的 TPU v2 —— 内存太小，不够用，直接跳过。 - Kaggle 免费提供的 TPU v3 —— 正好够用，本教程主要就靠它。 - Cloud TPU —— 提供 v3 及以上版本。设置起来稍麻烦一些，大致流程是：先创建一个新的 TPU VM，配置好 SSH 端口转发到 Jupyter 服务器的端口，然后在 TPU VM 上装 Jupyter 并启动，最后通过 Colab 的 "连接到本地运行时" 方式连过去。 ### 关于多 GPU 设置的说明 如果你手头有多 GPU 机器，完全可以按自己的需求调整这套流程，本教程虽然以 TPU 为主线，但思路是通用的。 如果想在 Colab 里用多 GPU，也可以通过 "连接到自定义 GCE VM" 的方式来配置。当然，我们这次还是把重心放在 Kaggle 免费 TPU 上——毕竟门槛最低嘛。 ## 开始之前 ### Gemma 设置 正式开始前，得先把 Gemma 的访问权限搞定。Gemma 模型由 Kaggle 托管，所以操作步骤如下： - 在 [kaggle.com](http://kaggle.com) 登录或注册。 - 打开 Gemma 模型卡片，点击 "请求访问权限"。 - 填写同意表格，接受条款和条件。 完成这些，才能正式拿到模型的访问权。 ## 安装 环境准备好了，接下来装包。需要安装 Keras 和 KerasNLP。 ``` # 安装最新的 Keras 3。更多信息查看 https://keras.io/getting_started/ !pip install -q tensorflow-cpu !pip install -q -U keras-nlp tensorflow-hub !pip install -q -U keras>=3 !pip install -U tensorflow-text ``` ### 设置 Keras JAX 后端 导入 JAX，先确认一下 TPU 是否正常工作。Kaggle 提供的 TPUv3-8 设备有 8 个核心，每个核心 16GB 内存。 ``` 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)] ``` 8 个核心全部就绪。接下来设置环境变量： ``` import os # Keras 3 分布式 API 目前只对 JAX 后端实现 os.environ["KERAS_BACKEND"] = "jax" # 预分配 90% 的 TPU 内存，减少碎片和开销 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')`。这个操作对训练质量影响很小，在大部分情况下都建议开启，既省内存又省时间。不过要注意：如果 batch size 比较小，混合精度反而可能让内存占用增加 1.5 倍（因为权重会被同时加载为半精度和全精度）。 推理时，半精度就够了，直接用 `keras.config.set_floatx("bfloat16")` 即可省内存，混合精度反而不适用。 ``` # 如果你想在 GPU 上启用混合精度训练，取消这行的注释 # keras.mixed_precision.set_global_policy('mixed_bfloat16') ``` 要在 TPU 上以分布式方式加载模型及其权重，首先得创建一个 `DeviceMesh`——简单说就是一张 "设备地图"，告诉 Keras 哪些硬件参与分布式计算。这是 Keras 3 统一分布式 API 的一部分。 这套 API 支持数据和模型并行，能高效地在多个翻跟斗和主机间扩展模型。它底层利用 JAX 的 SPMD（单程序多数据）机制来分布计算。想深入了解的，可以去看 Keras 3 分布式 API 指南。 ``` # 创建一个形状为 (1, 8) 的 DeviceMesh，把所有 8 个 TPU 核心纳入进来 device_mesh = keras.distribution.DeviceMesh( (1, 8), ["batch", "model"], devices=keras.distribution.list_devices()) ``` 接着用 `LayoutMap` 来指定权重和张量的分片策略——它就像一个规则表，用字符串键（支持正则）来匹配张量路径，然后决定哪些张量要分片，哪些要复制。 ``` model_dim = "model" layout_map = keras.distribution.LayoutMap(device_mesh) # token_embedding/embeddings 会在 8 个 TPU 上分片 layout_map["token_embedding/embeddings"] = (None, model_dim) # 匹配解码器中注意力层的 query、key、value 矩阵 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) ``` `ModelParallel` 负责在 `DeviceMesh` 的所有设备之间分片模型权重或激活张量。我们按照上面定义的 `layout_map`，把 Gemma 7B 的部分权重拆到 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") ``` 现在验证一下模型是否正确分区。拿 `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)}') ``` 输出： ``` 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') ``` 从 `PartitionSpec` 来看，注意力部分的 query、key、value 以及 FFW 层的权重都按 'model' 维度进行了分片，说明模型并行策略已经生效。 ## 微调前推理 先看看微调前，原模型对 "90 年代最佳喜剧电影" 这个问题的反应如何： ``` 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 年代喜剧片。接下来我们微调它，让输出风格往另一个方向靠。 ## 使用 IMDB 数据微调 ``` import tensorflow_datasets as tfds imdb_train = tfds.load( "imdb_reviews", split="train", as_supervised=True, batch_size=2, ) # 丢掉标签，只要文本 imdb_train = imdb_train.map(lambda x, y: x) imdb_train.unbatch().take(1).get_single_element().numpy() ``` ``` # 取 2000 条数据作为子集，加快训练速度 imdb_train = imdb_train.take(2000) ``` 这次微调我们用 LoRA（低秩适应）。LoRA 的原理其实很简单：冻结模型的全部原始权重，只在模型里插入少量新的可训练参数——准确来说是通过两个低秩矩阵 A 和 B 来重新参数化。这样做的好处是训练速度更快、内存更省。 ``` # 启用 LoRA，设置秩为 4 gemma_lm.backbone.enable_lora(rank=4) ``` ``` # 在 IMDB 电影评论数据集上微调 # 把输入序列长度限制在 128，控制内存 gemma_lm.preprocessor.sequence_length = 128 # 使用 AdamW——Transformer 模型的常见优化器 optimizer = keras.optimizers.AdamW( learning_rate=5e-5, weight_decay=0.01, ) # layernorm 和偏置项不参与权重衰减 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) ``` 这里有个数字值得注意：启用 LoRA 之后，可训练参数从 70 亿降到了仅仅 1100 万——降幅相当惊人。 ## 微调后推理 微调完毕，再问一次同样的问题： ``` gemma_lm.generate("Best comedy movies in the 90s ", max_length=64) ``` 输出： ``` 'Best comedy movies in the 90s 1990-1999. 10. Austin Powers - International Man of Mystery (1997) 9. The Wedding Singer (1998) 8. The Cable Guy (1996) 7' ``` 可以看到，模型已经吸收了电影评论数据集的风格，输出不再是简单地列片名，而是带有评分和年份格式的列表——虽然内容上还是 90 年代喜剧片，但叙事风格完全不同了。这就说明，微调确实生效了。

抖音极速版赚钱快速提现方法☜☜☜☜☜点击保存