JAX 및 Flax를 사용하여 RecurrentGemma로 추론

ai.google.dev에서 보기 Google Colab에서 실행 Vertex AI에서 열기 GitHub에서 소스 보기

이 튜토리얼에서는 JAX (고성능 수치 컴퓨팅 라이브러리), Flax (JAX 기반 신경망 라이브러리), Orbax (JAX 기반 신경망 라이브러리), Orbax (체크포인트 생성과 같은 유틸리티를 위한 JAX 기반 신경망 라이브러리)로 작성된 Google DeepMind의 recurrentgemma 라이브러리를 사용하여 RecurrentGemma 2B Instruct 모델로 기본 샘플링/추론을 수행하는 방법을 보여줍니다.SentencePiece Flax는 이 노트북에서 직접 사용되지는 않지만 Gemma와 RecurrentGemma (Griffin 모델)를 만드는 데 Flax를 사용했습니다.

이 노트북은 T4 GPU를 사용하는 Google Colab에서 실행할 수 있습니다 (수정 > 노트북 설정으로 이동한 후 하드웨어 가속기에서 T4 GPU 선택).

설정

다음 섹션에서는 모델 액세스, API 키 가져오기, 노트북 런타임 구성 등 RecurrentGemma 모델을 사용하기 위해 노트북을 준비하는 단계를 설명합니다.

Gemma에 Kaggle 액세스 권한 설정하기

이 튜토리얼을 완료하려면 먼저 Gemma 설정과 비슷한 설정 안내를 따라야 하지만 몇 가지 예외가 있습니다.

  • kaggle.com에서 Gemma 대신 RecurrentGemma에 액세스하세요.
  • 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의 종속 항목 리졸버'에 관한 오류가 발생하면 일반적으로 무시해도 됩니다.

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

RecurrentGemma 모델 로드 및 준비

  1. 세 가지 인수를 사용하는 kagglehub.model_download를 사용하여 RecurrentGemma 모델을 로드합니다.
  • handle: Kaggle의 모델 핸들
  • path: (선택사항 문자열) 로컬 경로
  • force_download: (선택적 불리언) 모델을 강제로 다시 다운로드합니다.
RECURRENTGEMMA_VARIANT = '2b-it' # @param ['2b', '2b-it'] {type:"string"}

import kagglehub

RECURRENTGEMMA_PATH = kagglehub.model_download(f'google/recurrentgemma/flax/{RECURRENTGEMMA_VARIANT}')
<ph type="x-smartling-placeholder">
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...
</ph> 
print('RECURRENTGEMMA_PATH:', RECURRENTGEMMA_PATH)

RECURRENTGEMMA_PATH: /root/.cache/kagglehub/models/google/recurrentgemma/flax/2b-it/1
  1. 모델 가중치와 tokenizer의 위치를 확인한 다음 경로 변수를 설정합니다. tokenizer 디렉터리는 모델을 다운로드한 기본 디렉터리에 있고 모델 가중치는 하위 디렉터리에 있습니다. 예를 들면 다음과 같습니다.
  • tokenizer.model 파일은 /LOCAL/PATH/TO/recurrentgemma/flax/2b-it/1에 있습니다.
  • 모델 체크포인트는 /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 모델 체크포인트를 로드합니다. "single_device"로 설정된 sharding 인수는 단일 기기에 모든 모델 매개변수를 로드합니다.
import recurrentgemma
from recurrentgemma import jax as recurrentgemma

params = recurrentgemma.load_parameters(checkpoint_path=CKPT_PATH, sharding="single_device")
  1. sentencepiece.SentencePieceProcessor를 사용하여 구성된 RecurrentGemma 모델 tokenizer를 로드합니다.
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 모델 체크포인트/가중치 및 tokenizer 위에 recurrentgemma.jax.Sampler를 사용하여 sampler를 만듭니다.
sampler = recurrentgemma.Sampler(
    model=model,
    vocab=vocab,
    params=params,
)
  1. prompt로 프롬프트를 작성하고 추론을 수행합니다. total_generation_steps (응답을 생성할 때 수행되는 단계 수)를 조정할 수 있습니다. 이 예에서는 50를 사용하여 호스트 메모리를 보존합니다.
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.

자세히 알아보기