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

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

개요

Gemma는 Google DeepMind Gemini의 연구 및 기술을 기반으로 하는 최첨단 개방형 대규모 언어 모델 제품군입니다. 이 튜토리얼에서는 JAX(고성능 수치 컴퓨팅 라이브러리), Flax(JAX 기반 신경망 라이브러리), Orbax(checkpointing과 같은 학습 유틸리티를 위한 JAX 기반 라이브러리){/10gemmaSentencePiece 이 노트북에서 Flax를 직접 사용하지는 않지만 Flax를 사용하여 Gemma를 만들었습니다.

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

설정

1. Gemma용 Kaggle 액세스 권한 설정하기

이 튜토리얼을 완료하려면 먼저 Gemma 설정의 설정 안내에 따라 다음 작업을 수행해야 합니다.

  • kaggle.com에서 Gemma에 액세스하세요.
  • Gemma 모델을 실행하기에 충분한 리소스가 포함된 Colab 런타임을 선택합니다.
  • Kaggle 사용자 이름과 API 키를 생성하고 구성합니다.

Gemma 설정을 완료한 후에는 다음 섹션으로 이동하여 Colab 환경의 환경 변수를 설정합니다.

2. 환경 변수 설정하기

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')

3. gemma 라이브러리 설치

이 노트북에서는 무료 Colab GPU 사용에 중점을 둡니다. 하드웨어 가속을 사용 설정하려면 수정 > 노트북 설정을 클릭하고 T4 GPU > 저장을 선택합니다.

다음으로 github.com/google-deepmind/gemma에서 Google DeepMind gemma 라이브러리를 설치해야 합니다. 'pip의 종속 항목 리졸버'에 관한 오류가 발생하는 경우 일반적으로 무시해도 됩니다.

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

Gemma 모델 로드 및 준비

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

import kagglehub

GEMMA_PATH = kagglehub.model_download(f'google/gemma/flax/{GEMMA_VARIANT}')

Downloading from https://www.kaggle.com/api/v1/models/google/gemma/flax/2b-it/2/download...
100%|██████████| 3.67G/3.67G [00:35<00:00, 110MB/s]
Extracting model files...

print('GEMMA_PATH:', GEMMA_PATH)

GEMMA_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2
  1. 모델 가중치와 tokenizer의 위치를 확인한 후 경로 변수를 설정하세요. tokenizer 디렉터리는 모델을 다운로드한 기본 디렉터리에 있고 모델 가중치는 하위 디렉터리에 있습니다. 예를 들면 다음과 같습니다.
  • tokenizer.model 파일은 /LOCAL/PATH/TO/gemma/flax/2b-it/2에 있습니다.
  • 모델 체크포인트는 /LOCAL/PATH/TO/gemma/flax/2b-it/2/2b-it에 위치합니다.
CKPT_PATH = os.path.join(GEMMA_PATH, GEMMA_VARIANT)
TOKENIZER_PATH = os.path.join(GEMMA_PATH, 'tokenizer.model')
print('CKPT_PATH:', CKPT_PATH)
print('TOKENIZER_PATH:', TOKENIZER_PATH)

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

샘플링/추론 수행

  1. gemma.params.load_and_format_params 메서드를 사용하여 Gemma 모델 체크포인트를 로드하고 형식을 지정합니다.
from gemma import params as params_lib

params = params_lib.load_and_format_params(CKPT_PATH)
  1. sentencepiece.SentencePieceProcessor를 사용하여 구성된 Gemma tokenizer를 로드합니다.
import sentencepiece as spm

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

True
  1. Gemma 모델 체크포인트에서 올바른 구성을 자동으로 로드하려면 gemma.transformer.TransformerConfig를 사용하세요. cache_size 인수는 Gemma Transformer 캐시의 시간 단계 수입니다. 그런 다음 flax.linen.Module에서 상속되는 gemma.transformer.Transformer를 사용하여 Gemma 모델을 transformer로 인스턴스화합니다.
from gemma import transformer as transformer_lib

transformer_config = transformer_lib.TransformerConfig.from_params(
    params=params,
    cache_size=1024
)

transformer = transformer_lib.Transformer(transformer_config)
  1. Gemma 모델 체크포인트/가중치 및 tokenizer 위에 있는 gemma.sampler.Sampler를 사용하여 sampler를 만듭니다.
from gemma import sampler as sampler_lib

sampler = sampler_lib.Sampler(
    transformer=transformer,
    vocab=vocab,
    params=params['transformer'],
)
  1. input_batch로 프롬프트를 작성하고 추론을 실행합니다. total_generation_steps (응답을 생성할 때 수행되는 단계 수 - 이 예에서는 호스트 메모리를 보존하기 위해 100 사용)를 조정할 수 있습니다.
prompt = [
    "\n# What is the meaning of life?",
]

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

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

Prompt:

# What is the meaning of life?
Output:


The question of what the meaning of life is one that has occupied the minds of philosophers, theologians, and individuals for centuries. There is no single, universally accepted answer, but there are many different perspectives on this complex and multifaceted question.

**Some common perspectives on the meaning of life include:**

* **Biological perspective:** From a biological standpoint, the meaning of life is to survive and reproduce.
* **Existential perspective:** Existentialists believe that life is not inherently meaningful and that
  1. (선택사항) 노트북을 완료한 후 다른 프롬프트를 사용해 보려면 이 셀을 실행하여 메모리를 확보하세요. 그런 다음 3단계에서 sampler를 다시 인스턴스화하고 4단계에서 프롬프트를 맞춤설정하고 실행할 수 있습니다.
del sampler

자세히 알아보기