Griffin 架构开源实现：RNN 版 LLM

Hacker News · github.com#模型架构#RNN#开源

Editor brief · 编辑速览

社区发布 Google Griffin 架构实现，这是一种基于 RNN 的 LLM 架构，在特定场景相比 Transformer 更高效。

文章思维导图

Knowledge map

拖拽缩放

Full translation

完整中文译文

RecurrentGemma 是 Google DeepMind 推出的开源权重语言模型家族，基于新颖的 Griffin 架构。这一架构通过用本地注意力和线性递推的混合替代全局注意力，在生成长序列时实现快速推理。

本仓库包含模型实现以及采样和微调示例。我们建议大多数用户采用 Flax 实现，它经过高度优化。我们也提供了一个未经优化的 PyTorch 实现供参考。

RecurrentGemma 技术报告提供了关于 RecurrentGemma 训练和评估的具体细节。

Griffin 论文描述了底层模型架构。

RecurrentGemma 使用 Poetry 进行依赖管理。

完整项目的依赖安装：

如果只需安装部分依赖，使用下面的库特定命令：

如果你更喜欢用 pip 而不是 Poetry，则创建虚拟环境（运行 python -m venv recurrentgemma-demo 和 . recurrentgemma-demo/bin/activate），然后按上面的 pip 命令安装。

模型检查点可通过 Kaggle 获取，网址：http://kaggle.com/models/google/recurrentgemma。选择 Flax 或 PyTorch 模型变体，点击 ⤓ 按钮下载模型归档，然后将内容解压到本地目录。

两种情况下，归档都包含模型权重和分词器。

运行测试需要先从源代码树根目录安装可选的 [test] 依赖（例如使用 pip install .[test]），然后：

pytest .

运行示例采样脚本，传入权重目录和分词器的路径：

python examples/sampling_jax.py \
  --path_checkpoint=/path/to/archive/contents/2b/ \
  --path_tokenizer=/path/to/archive/contents/tokenizer.model

colabs/sampling_tutorial_jax.ipynb 包含一个 JAX 采样示例的 Colab 笔记本。

colabs/sampling_tutorial_pytorch.ipynb 包含一个 PyTorch 采样示例的 Colab 笔记本。

colabs/fine_tuning_tutorial_jax.ipynb 包含一个使用 JAX 进行 RecurrentGemma 微调的基础教程 Colab，例如英法翻译任务。

要运行这些笔记本，你需要拥有 Kaggle 账户，并首先从 RecurrentGemma 页面读取并接受 Gemma 许可条款。之后可以运行笔记本，它会自动从那里下载权重和分词器。

目前不同笔记本支持的硬件如下：

RecurrentGemma 代码可在 CPU、GPU 或 TPU 上运行。该代码已针对使用 Flax 实现在 TPU 上运行进行了优化，其中包含一个低级 Pallas 内核来在循环层中执行线性扫描。

我们欢迎错误报告和问题反馈。详见 CONTRIBUTING.md。

本代码采用 Apache License, Version 2.0（"许可证"）授权；你不可在不遵守该许可证的情况下使用本文件。你可在以下网址获取许可证副本：http://www.apache.org/licenses/LICENSE-2.0

除非适用法律要求或以书面形式同意，本软件按"现状"基础提供，不附带任何明示或暗示的担保或条件。详见许可证中对权利和限制的具体规定。

本不是官方 Google 产品。