跳轉到

使用 JAX 與 Flax 以 RecurrentGemma 執行推論

Models RecurrentGemma Uses JAX Colab

本教學示範如何使用 Google DeepMind 的 recurrentgemma 函式庫（以 JAX（高效能數值運算函式庫）、Flax（基於 JAX 的類神經網路函式庫）、Orbax（基於 JAX 的訓練工具函式庫，例如檢查點）與 SentencePiece（分詞/還原詞元函式庫）撰寫）對 RecurrentGemma 2B Instruct 模型執行基本的取樣/推論。雖然本筆記本未直接使用 Flax，但 Gemma 與 RecurrentGemma（Griffin 模型）正是用 Flax 建立的。

這本筆記本可在 Google Colab 上使用 T4 GPU 執行（前往編輯 > 筆記本設定 > 在硬體加速器下方選取 T4 GPU）。

設定

以下各節說明準備筆記本以使用 RecurrentGemma 模型的步驟，包括模型存取、取得 API 金鑰，以及設定筆記本執行環境

設定 Kaggle 存取以使用 Gemma

要完成本教學，您首先需要依照與 Gemma 設定「類似」的設定說明進行，但有以下例外：

  • 在 kaggle.com 取得 RecurrentGemma（而非 Gemma）的存取權。
  • 選取資源足以執行 RecurrentGemma 模型的 Colab 執行環境。
  • 產生並設定 Kaggle 使用者名稱與 API 金鑰。

完成 RecurrentGemma 設定後，請繼續下一節，在 Colab 環境中設定環境變數。

設定環境變數

設定 KAGGLE_USERNAME 與 KAGGLE_KEY 的環境變數。當出現「是否授權存取？」的提示時，請同意提供密鑰存取權。

import os
from google.colab import userdata # `userdata` is a Colab API.

os.environ["KAGGLE_USERNAME"] = userdata.get('KAGGLE_USERNAME')
os.environ["KAGGLE_KEY"] = userdata.get('KAGGLE_KEY')

安裝 recurrentgemma 函式庫

本筆記本聚焦於使用免費的 Colab GPU。若要啟用硬體加速，請點選編輯 > 筆記本設定 > 選取 T4 GPU > 儲存。

接著，您需要從 github.com/google-deepmind/recurrentgemma 安裝 Google DeepMind 的 recurrentgemma 函式庫。如果看到關於「pip 依賴解析器」的錯誤，通常可以忽略。

Note

安裝 recurrentgemma 時，也會一併安裝 flax、核心 jax、optax（基於 JAX 的梯度處理與最佳化函式庫）、orbax 與 sentencepiece。

pip install git+https://github.com/google-deepmind/recurrentgemma.git

載入並準備 RecurrentGemma 模型

  1. 使用 kagglehub.model_download 載入 RecurrentGemma 模型，它接受三個參數：

  2. handle：來自 Kaggle 的模型代碼

  3. path：（選用字串）本機路徑
  4. force_download：（選用布林值）強制重新下載模型

Note

請注意，recurrentgemma-2b-it 模型大小約為 3.85Gb。

RECURRENTGEMMA_VARIANT = '2b-it' # @param ['2b', '2b-it'] {type:"string"}

import kagglehub

RECURRENTGEMMA_PATH = kagglehub.model_download(f'google/recurrentgemma/flax/{RECURRENTGEMMA_VARIANT}')

Downloading from https://www.kaggle.com/api/v1/models/google/recurrentgemma/flax/2b-it/1/download...
100%|██████████| 3.85G/3.85G [00:52<00:00, 78.2MB/s]
Extracting model files...

print('RECURRENTGEMMA_PATH:', RECURRENTGEMMA_PATH)

RECURRENTGEMMA_PATH: /root/.cache/kagglehub/models/google/recurrentgemma/flax/2b-it/1

Note

上方的輸出路徑是模型權重與分詞器在本機的儲存位置，稍後會用到。

  1. 檢查模型權重與分詞器的位置，然後設定路徑變數。分詞器目錄位於您下載模型的主目錄中，而模型權重則位於子目錄中。例如：

  2. tokenizer.model 檔案位於 /LOCAL/PATH/TO/recurrentgemma/flax/2b-it/1。

  3. 模型檢查點位於 /LOCAL/PATH/TO/recurrentgemma/flax/2b-it/1/2b-it。
CKPT_PATH = os.path.join(RECURRENTGEMMA_PATH, RECURRENTGEMMA_VARIANT)
TOKENIZER_PATH = os.path.join(RECURRENTGEMMA_PATH, 'tokenizer.model')
print('CKPT_PATH:', CKPT_PATH)
print('TOKENIZER_PATH:', TOKENIZER_PATH)

CKPT_PATH: /root/.cache/kagglehub/models/google/recurrentgemma/flax/2b-it/1/2b-it
TOKENIZER_PATH: /root/.cache/kagglehub/models/google/recurrentgemma/flax/2b-it/1/tokenizer.model

執行取樣/推論

  1. 使用 recurrentgemma.jax.load_parameters 方法載入 RecurrentGemma 模型檢查點。sharding 參數設為 "single_device" 會將所有模型參數載入單一裝置。
import recurrentgemma
from recurrentgemma import jax as recurrentgemma

params = recurrentgemma.load_parameters(checkpoint_path=CKPT_PATH, sharding="single_device")
  1. 載入 RecurrentGemma 模型分詞器，它是以 sentencepiece.SentencePieceProcessor 建構的：
import sentencepiece as spm

vocab = spm.SentencePieceProcessor()
vocab.Load(TOKENIZER_PATH)

True
  1. 若要從 RecurrentGemma 模型檢查點自動載入正確的設定，請使用 recurrentgemma.GriffinConfig.from_flax_params_or_variables。接著，使用 recurrentgemma.jax.Griffin 具現化 Griffin 模型。
model_config = recurrentgemma.GriffinConfig.from_flax_params_or_variables(
    flax_params_or_variables=params)

model = recurrentgemma.Griffin(model_config)
  1. 在 RecurrentGemma 模型檢查點/權重與分詞器之上，使用 recurrentgemma.jax.Sampler 建立 sampler：
sampler = recurrentgemma.Sampler(
    model=model,
    vocab=vocab,
    params=params,
)
  1. 在 prompt 中寫入提示並執行推論。您可以調整 total_generation_steps（生成回應時執行的步數——本範例使用 50 以節省主機記憶體）。

Note

如果記憶體不足，請點選執行階段 > 中斷連線並刪除執行階段，然後點選執行階段 > 全部執行。

prompt = [
    "\n# 5+9=?",
]

reply = sampler(input_strings=prompt,
                total_generation_steps=50,
                )

for input_string, out_string in zip(prompt, reply.text):
    print(f"Prompt:\n{input_string}\nOutput:\n{out_string}")

/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/mlir.py:920: UserWarning: Some donated buffers were not usable: ShapedArray(int32[1,8]).
See an explanation at https://jax.readthedocs.io/en/latest/faq.html#buffer_donation.
  warnings.warn("Some donated buffers were not usable:"
Prompt:

# 5+9=?
Output:


# Answer: 14

# Explanation: 5 + 9 = 14.

了解更多