v5e Cloud TPU VM에서 JetStream PyTorch 추론

JetStream은 XLA 기기(TPU)에서 대규모 언어 모델(LLM) 추론을 위한 처리량 및 메모리 최적화 엔진입니다.

시작하기 전에

Cloud TPU 환경 설정의 단계에 따라 Google Cloud 프로젝트를 만들고, TPU API를 활성화하고, TPU CLI를 설치한 후 TPU 할당량을 요청합니다.

CreateNode API를 사용하여 Cloud TPU 만들기의 단계에 따라 --accelerator-type을 v5litepod-8로 설정하는 TPU VM을 만듭니다.

JetStream 저장소 클론 및 종속 항목 설치

  1. SSH를 사용하여 TPU VM에 연결하기

    • ${TPU_NAME}을 TPU 이름으로 설정합니다.
    • ${PROJECT}를 Google Cloud 프로젝트로 설정
    • TPU를 만들 Google Cloud 영역에 ${ZONE}을 설정합니다.
      gcloud compute config-ssh
  gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT} --zone ${ZONE}

  2. jetstream-pytorch 코드 가져오기 bash git clone https://github.com/google/jetstream-pytorch.git git checkout jetstream-v0.2.4

(선택사항) venv 또는 conda를 사용하여 가상 환경을 만들고 활성화합니다.

sudo apt install python3.10-venv
python -m venv venv
source venv/bin/activate
  1. 설치 스크립트를 실행합니다.
cd jetstream-pytorch
source install_everything.sh

jetstream pytorch 실행

지원되는 모델 나열

jpt list

그러면 지원되는 모델 및 변형 목록이 출력됩니다.

meta-llama/Llama-2-7b-chat-hf
meta-llama/Llama-2-7b-hf
meta-llama/Llama-2-13b-chat-hf
meta-llama/Llama-2-13b-hf
meta-llama/Llama-2-70b-hf
meta-llama/Llama-2-70b-chat-hf
meta-llama/Meta-Llama-3-8B
meta-llama/Meta-Llama-3-8B-Instruct
meta-llama/Meta-Llama-3-70B
meta-llama/Meta-Llama-3-70B-Instruct
google/gemma-2b
google/gemma-2b-it
google/gemma-7b
google/gemma-7b-it
mistralai/Mixtral-8x7B-v0.1
mistralai/Mixtral-8x7B-Instruct-v0.1

하나의 모델로 jetstream-pytorch 서버를 실행하려면 다음 단계를 따르세요. bash jpt serve --model_id meta-llama/Llama-2-7b-chat-hf

이 모델을 처음 실행하면 jpt serve 명령어에서 HuggingFace에서 가중치를 다운로드하려고 시도하며, 이때 HuggingFace로 인증해야 합니다.

인증하려면 huggingface-cli login를 실행하여 액세스 토큰을 설정하거나 --hf_token 플래그를 사용하여 HuggingFace 액세스 토큰을 jpt serve 명령어에 전달합니다.

jpt serve --model_id meta-llama/Llama-2-7b-chat-hf --hf_token=...

HuggingFace 액세스 토큰에 대한 자세한 내용은 액세스 토큰을 참고하세요.

HuggingFace Hub를 사용하여 로그인하려면 다음 명령어를 실행하고 프롬프트를 따르세요.

pip install -U "huggingface_hub[cli]"
huggingface-cli login

가중치가 다운로드되면 더 이상 --hf_token 플래그를 지정할 필요가 없습니다.

int8 정규화로 이 모델을 실행하려면 --quantize_weights=1를 추가합니다. 무게가 로드되면 비행 중에 양자화가 실행됩니다.

HuggingFace에서 다운로드한 가중치는 기본적으로 jpt를 실행하는 디렉터리의 checkpoints 폴더라는 디렉터리에 저장됩니다. --working_dir 플래그를 사용하여 디렉터리를 변경하거나 지정할 수도 있습니다.

자체 체크포인트를 사용하려면 checkpoints/<org>/<model>/hf_original 디렉터리 (또는 --working_dir의 해당 하위 디렉터리)에 배치합니다. 예를 들어 Llama2-7b 체크포인트는 checkpoints/meta-llama/Llama-2-7b-hf/hf_original/*.safetensors에 있습니다. 이러한 파일을 HuggingFace 형식의 수정된 가중치로 대체할 수 있습니다.

벤치마크 실행

install_everything.sh를 실행할 때 다운로드된 deps/JetStream 폴더로 변경합니다.

cd deps/JetStream
wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
export dataset_path=ShareGPT_V3_unfiltered_cleaned_split.json
python benchmarks/benchmark_serving.py --tokenizer $tokenizer_path --num-prompts 2000  --dataset-path  $dataset_path --dataset sharegpt --save-request-outputs

자세한 내용은 deps/JetStream/benchmarks/README.md를 참조하세요.

일반적인 오류

Unexpected keyword argument 'device' 오류가 발생하면 다음을 시도해 보세요.

  • jax 및 jaxlib 종속 항목 제거
  • source install_everything.sh를 사용하여 재설치

Out of memory 오류가 발생하면 다음을 시도해 보세요.

  • 더 작은 배치 크기 사용
  • 양자화 사용

삭제

이 튜토리얼에서 사용된 리소스 비용이 Google Cloud 계정에 청구되지 않도록 하려면 리소스가 포함된 프로젝트를 삭제하거나 프로젝트를 유지하고 개별 리소스를 삭제하세요.

  1. GitHub 저장소 삭제

      # Clean up the JetStream repository
  rm -rf JetStream

  # Clean up the xla repository
  rm -rf xla

  2. Python 가상 환경 삭제

    rm -rf .env

  3. TPU 리소스 삭제

    자세한 내용은 TPU 리소스 삭제를 참고하세요.