用途
RecurrentGemma 是 Google DeepMind 推出的开放权重语言模型系列，基于新颖的 Griffin 架构。该架构通过将全局注意力替换为局部注意力和线性递归的混合，在生成长序列时实现快速推理。它适用于需要长序列生成且注重推理效率的场景，例如文本生成、翻译等任务。
主要能力
- 提供开放权重语言模型，基于 Griffin 架构。
- 包含模型实现以及采样和微调的示例代码。
- 官方推荐使用高度优化的 Flax 实现，同时提供未优化的 PyTorch 参考实现。
- 支持在 CPU、GPU 或 TPU 上运行，Flax 实现针对 TPU 进行了优化，并包含用于递归层线性扫描的低层 Pallas 内核。
使用方式
安装依赖
项目使用 Poetry 管理依赖。完整安装步骤：
- 检出代码。
- 运行
poetry install -E full创建包含所有依赖的虚拟环境。 - 运行
poetry shell激活虚拟环境。
如果仅需部分依赖，可使用以下替代命令：
- 若使用 pip 而非 Poetry，先创建虚拟环境（
python -m venv recurrentgemma-demo和source recurrentgemma-demo/bin/activate），然后运行pip install .[full]。 - 仅安装 JAX 路径依赖：
poetry install -E jax或pip install .[jax]。 - 仅安装 PyTorch 路径依赖：
poetry install -E torch或pip install .[torch]。 - 安装运行单元测试所需依赖：
poetry install -E test或pip install .[test]。
下载模型
模型检查点可通过 Kaggle 获取：http://kaggle.com/models/google/recurrentgemma。选择 Flax 或 PyTorch 模型变体，点击下载按钮下载模型存档，然后解压到本地目录。存档中包含模型权重和分词器。
运行单元测试
在源码树根目录安装 [test] 可选依赖（如 pip install .[test]），然后运行测试。
运行采样脚本
传递权重目录和分词器的路径即可运行示例采样脚本。
Colab 教程
colabs/sampling_tutorial_jax.ipynb：使用 JAX 的采样示例。colabs/sampling_tutorial_pytorch.ipynb：使用 PyTorch 的采样示例。colabs/fine_tuning_tutorial_jax.ipynb：使用 JAX 微调 RecurrentGemma 的基础教程，例如英法翻译任务。
运行这些笔记本需要 Kaggle 账户，并先在 RecurrentGemma 页面阅读并接受 Gemma 许可条款。之后运行笔记本时会自动下载权重和分词器。目前不同笔记本支持的硬件暂未提供详细信息。
适用场景
- 需要高效生成长序列的语言模型应用。
- 学术研究或实验，探索新型 Griffin 架构。
- 基于开放权重模型进行微调，适配特定任务。
限制
- 仓库中的 PyTorch 实现未经过优化，仅供参考。
- 使用模型需遵守 Gemma 许可条款。
- 目前不同 Colab 笔记本支持的硬件暂未提供详细信息。
- 该项目并非 Google 官方产品。
许可
代码采用 Apache License 2.0 许可。版权所有 2024 DeepMind Technologies Limited。
