Cloud TPU v5e 추론 소개

개요 및 이점

Cloud TPU v5e는 변환기 기반, 텍스트-이미지 및 CNN 기반 학습, 미세 조정, 서빙(추론)에 최적화된 Google에서 개발된 AI 가속기입니다. TPU v5e 슬라이스에는 최대 256개의 칩을 포함할 수 있습니다.

서빙이란 추론에 사용할 수 있는 학습된 머신러닝 모델을 프로덕션 환경에 배포하는 프로세스입니다. 서빙에서는 지연 시간 SLO가 우선시됩니다.

이 문서에서는 단일 호스트 TPU에서의 모델 서빙에 대해 설명합니다. 칩이 8개 이하인 TPU 슬라이스는 하나의 TPU VM 또는 호스트를 가지며 단일 호스트 TPU라고 합니다.

시작하기

v5e TPU의 할당량이 필요합니다. 주문형 TPU에는 tpu-v5s-litepod-serving 할당량이 필요합니다. 예약된 TPU에는 tpu-v5s-litepod-serving-reserved 할당량이 필요합니다. 자세한 내용은 Cloud 영업팀에 문의하세요.

Cloud TPU를 사용하려면 Google Cloud 계정과 프로젝트가 필요합니다. 자세한 내용은 Cloud TPU 환경 설정을 참조하세요.

큐에 추가된 리소스를 사용하여 v5e TPU를 프로비저닝합니다. 서빙에 사용 가능한 v5e 구성에 대한 자세한 내용은 서빙을 위한 Cloud TPU v5e 유형을 참조하세요.

Cloud TPU 모델 추론 및 서빙

추론을 위한 모델 서빙 방법은 모델을 작성한 ML 프레임워크에 따라 달라집니다. TPU v5e는 JAX, TensorFlow, PyTorch로 작성된 서빙 모델을 지원합니다.

JAX 모델 추론 및 서빙

TPU VM에서 모델을 서빙하려면 다음을 수행해야 합니다.

TensorFlow SavedModel 형식으로 모델을 직렬화
추론 변환기를 사용하여 저장된 모델을 서빙할 수 있도록 준비
TensorFlow Serving을 사용하여 모델 서빙

SavedModel 형식

SavedModel에는 학습된 매개변수 및 연산을 포함한 완전한 TensorFlow 프로그램이 포함되어 있습니다. 원본 모델 빌드 코드를 실행할 필요가 없습니다.

모델이 JAX로 작성된 경우 jax2tf를 사용하여 모델을 SavedModel 형식으로 직렬화해야 합니다.

추론 변환기

Cloud TPU 추론 변환기는 TPU 추론을 위해 SavedModel 형식으로 내보낸 모델을 준비하고 최적화합니다. 로컬 셸 또는 TPU VM에서 추론 변환기를 실행할 수 있습니다. 변환기를 실행하는 데 필요한 모든 명령줄 도구가 있으므로 TPU VM 셸을 사용하는 것이 좋습니다. 추론 변환기에 대한 자세한 내용은 추론 변환기 사용자 가이드를 참조하세요.

추론 변환기 요구사항

TensorFlow 또는 JAX에서 SavedModel 형식으로 모델을 내보내야 합니다.
TPU 함수의 함수 별칭을 정의해야 합니다. 자세한 내용은 추론 변환기 사용자 가이드를 참조하세요. 이 가이드의 예시에서는 tpu_func를 TPU 함수 별칭으로 사용합니다.
TensorFlow 라이브러리(Cloud TPU 추론 변환기의 종속 항목)가 AVX 명령을 사용하도록 컴파일되므로 머신 CPU가 Advanced Vector eXtensions(AVX) 명령을 지원하는지 확인합니다. 대부분의 CPU는 AVX를 지원합니다.
참고: lscpu | grep avx를 실행하여 AVX 명령 집합이 지원되는지 확인할 수 있습니다.

JAX 모델 추론 및 서빙

이 섹션에서는 jax2tf 및 TensorFlow Serving을 사용하여 JAX 모델을 서빙하는 방법을 설명합니다.

jax2tf를 사용하여 모델을 SavedModel 형식으로 직렬화
추론 변환기를 사용하여 저장된 모델을 서빙할 수 있도록 준비
TensorFlow Serving을 사용하여 모델 서빙

`jax2tf`를 사용하여 JAX 모델을 SavedModel 형식으로 직렬화

다음 Python 함수는 모델 코드 내에서 jax2tf를 사용하는 방법을 보여줍니다.

# Inference function
def model_jax(params, inputs):
  return params[0] + params[1] * inputs

# Wrap the parameter constants as tf.Variables; this will signal to the model
# saving code to save those constants as variables, separate from the
# computation graph.
params_vars = tf.nest.map_structure(tf.Variable, params)

# Build the prediction function by closing over the `params_vars`. If you
# instead were to close over `params` your SavedModel would have no variables
# and the parameters will be included in the function graph.
prediction_tf = lambda inputs: jax2tf.convert(model_jax)(params_vars, inputs)

my_model = tf.Module()
# Tell the model saver what the variables are.
my_model._variables = tf.nest.flatten(params_vars)
my_model.f = tf.function(prediction_tf, jit_compile=True, autograph=False)
tf.saved_model.save(my_model)

jax2tf에 대한 자세한 내용은 JAX와 Cloud TPU 상호 운용성을 참조하세요.

추론 변환기를 사용하여 저장된 모델을 서빙할 수 있도록 준비

추론 변환기 사용 안내는 추론 변환기 가이드를 참조하세요.

TensorFlow Serving 사용

TensorFlow Serving 사용 안내는 TensorFlow 서빙을 참조하세요.

JAX 모델 서빙 예시

기본 요건

Docker 사용자 인증 정보를 설정하고 추론 변환기 및 Cloud TPU Serving Docker 이미지를 가져옵니다.

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
   us-docker.pkg.dev
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

SSH를 사용하여 TPU VM에 연결하고 추론 데모 코드를 설치합니다.
```
gsutil -m cp -r \
"gs://cloud-tpu-inference-public/demo" \
.
```

JAX 데모 종속 항목을 설치합니다.

pip install -r ./demo/jax/requirements.txt

추론을 위한 JAX BERT 모델 서빙

Hugging Face에서 사전 학습된 BERT 모델을 다운로드할 수 있습니다.

Flax BERT 모델에서 TPU 호환 TensorFlow 저장된 모델을 내보냅니다.
```
cd demo/jax/bert
python3 export_bert_model.py
```

Cloud TPU 모델 서버 컨테이너를 시작합니다.

docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=/tmp/jax/bert_tpu,target=/models/bert \
  -e MODEL_NAME=bert \
  us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

컨테이너가 시작되고 약 30초 후 모델 서버 컨테이너 로그를 확인하고 gRPC와 HTTP 서버가 작동 중인지 확인합니다.

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker logs ${CONTAINER_ID}

다음 정보로 끝나는 로그 항목이 표시되면 서버가 요청을 처리할 준비가 된 것입니다.

2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
[warn] getaddrinfo: address family for nodename not supported
2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 245] NET_LOG: Entering the event loop ...

모델 서버에 추론 요청을 보냅니다.

python3 bert_request.py

출력은 다음과 비슷하게 표시됩니다.

For input "The capital of France is [MASK].", the result is ". the capital of france is paris.."
For input "Hello my name [MASK] Jhon, how can I [MASK] you?", the result is ". hello my name is jhon, how can i help you?."

삭제

다른 데모를 실행하기 전에 Docker 컨테이너를 삭제해야 합니다.
```
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}
```
모델 아티팩트를 삭제합니다.
```
sudo rm -rf /tmp/jax/
```

추론을 위한 JAX Stable Diffusion 서빙

Hugging Face에서 사전 학습된 Stable Diffusion 모델을 다운로드할 수 있습니다.

TPU 호환 TF2 저장된 모델 형식으로 Stable Diffusion 모델을 다운로드합니다.
참고: 다운로드에 3분 정도 걸립니다.
```
cd demo/jax/stable_diffusion
python3 export_stable_diffusion_model.py
```

모델의 Cloud TPU 모델 서버 컨테이너를 시작합니다.

docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=/tmp/jax/stable_diffusion_tpu,target=/models/stable_diffusion \
  -e MODEL_NAME=stable_diffusion \
  us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

약 2분 후에 모델 서버 컨테이너 로그를 확인하여 gRPC 및 HTTP 서버가 실행 중인지 확인합니다.

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker logs ${CONTAINER_ID}

다음 정보로 끝나는 로그가 표시되면 서버가 요청을 처리할 준비가 된 것입니다.

2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
[warn] getaddrinfo: address family for nodename not supported
2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 245] NET_LOG: Entering the event loop ...

모델 서버로 요청을 보냅니다.
```
python3 stable_diffusion_request.py
```
이 스크립트는 프롬프트로 'Painting of a squirrel skating in New York'을 전송합니다. 출력 이미지는 현재 디렉터리에 stable_diffusion_images.jpg로 저장됩니다.

삭제

다른 데모를 실행하기 전에 Docker 컨테이너를 삭제해야 합니다.

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}

모델 아티팩트 삭제

sudo rm -rf /tmp/jax/

TensorFlow Serving

다음 안내에서는 TPU VM에서 TensorFlow 모델을 서빙하는 방법을 보여줍니다.

TensorFlow Serving 워크플로

TPU VM용 TensorFlow Serving Docker 이미지를 다운로드합니다.

샘플 환경 변수 설정

export YOUR_LOCAL_MODEL_PATH=model-path
export MODEL_NAME=model-name
# Note: this image name may change later.
export IMAGE_NAME=us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Docker 이미지 다운로드

docker pull ${IMAGE_NAME}

Docker 사용자 인증 정보를 설정하고 추론 변환기 및 TensorFlow Serving Docker 이미지를 가져옵니다.

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
   us-docker.pkg.dev
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

데모 코드를 다운로드합니다.

gsutil -m cp -r \
"gs://cloud-tpu-inference-public/demo" \
.

TensorFlow 데모 종속 항목을 설치합니다.
```
pip install -r ./demo/tf/requirements.txt
```

TPU VM에서 TensorFlow Serving Docker 이미지를 사용하여 TensorFlow 모델을 서빙합니다.

# PORT 8500 is for gRPC model server and 8501 is for HTTP/REST model server.
docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=${YOUR_LOCAL_MODEL_PATH},target=/models/${MODEL_NAME} \
  -e MODEL_NAME=${MODEL_NAME} \
  ${IMAGE_NAME}

Serving Client API를 사용하여 모델을 쿼리합니다.
- REST Client API
- gRPC Client API

TensorFlow ResNet-50 서빙 데모 실행

Keras ResNet-50 모델에서 TPU 호환 TF2 저장된 모델을 내보냅니다.
```
cd demo/tf/resnet-50
python3 export_resnet_model.py
```

모델의 TensorFlow 모델 서버 컨테이너를 시작합니다.

docker run -t --rm --privileged -d \
  -p 8500:8500 -p 8501:8501 \
  --mount type=bind,source=/tmp/tf/resnet_tpu,target=/models/resnet \
  -e MODEL_NAME=resnet \
  us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

모델 서버 컨테이너 로그를 확인하고 gRPC 및 HTTP 서버가 작동 중인지 확인합니다.

CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker logs ${CONTAINER_ID}

다음 정보로 끝나는 로그가 표시되면 서버가 요청을 처리할 준비가 된 것입니다. 30초 정도 걸립니다.

2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
[warn] getaddrinfo: address family for nodename not supported
2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
[evhttp_server.cc : 245] NET_LOG: Entering the event loop ...

모델 서버로 요청을 보냅니다.

요청 이미지는 https://i.imgur.com/j9xCCzn.jpeg의 바나나입니다.
```
python3 resnet_request.py
```
출력은 다음과 비슷하게 표시됩니다.
```
Predict result: [[('n07753592', 'banana', 0.94921875), ('n03532672', 'hook', 0.022338867), ('n07749582', 'lemon', 0.005126953)]]
```
삭제

다른 데모를 실행하기 전에 Docker 컨테이너를 삭제해야 합니다.
```
CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
docker stop ${CONTAINER_ID}
```
모델 아티팩트를 삭제합니다.
```
sudo rm -rf /tmp/tf/
```

PyTorch 모델 추론 및 서빙

PyTorch로 작성된 모델의 워크플로는 다음과 같습니다.

TorchDynamo 및 PyTorch/XLA를 사용하여 로드하고 추론하는 Python 모델 핸들러 작성
TorchModelArchiver를 사용하여 모델 보관 파일 만들기
TorchServe를 사용하여 모델 서빙

TorchDynamo 및 PyTorch/XLA

TorchDynamo(Dynamo)는 PyTorch 프로그램의 속도를 높이도록 설계된 Python 수준 JIT 컴파일러입니다. 컴파일러 백엔드를 연결할 수 있는 깔끔한 API를 제공합니다. 실행 직전에 Python 바이트 코드를 동적으로 수정합니다. PyTorch/XLA 2.0 출시 버전에는 Dynamo를 사용한 추론 및 학습을 위한 실험용 백엔드가 있습니다.

Dynamo는 모델 패턴이 인식되면 Torch FX(FX) 그래프를 제공하며, PyTorch/XLA는 지연 텐서 방식을 사용하여 FX 그래프를 컴파일하고 컴파일된 함수를 반환합니다. Dynamo에 대한 자세한 내용은 다음을 참조하세요.

Pytorch 개발 토론 게시물
TorchDynamo 문서
자세한 내용은 PyTorch 2.0 & XLA를 참조하세요.

다음은 torch.compile로 densenet161 추론을 실행하는 작은 코드 예시입니다.

import torch
import torchvision
import torch_xla.core.xla_model as xm

def eval_model(loader):
  device = xm.xla_device()
  xla_densenet161 = torchvision.models.densenet161().to(device)
  xla_densenet161.eval()
  dynamo_densenet161 = torch.compile(
      xla_densenet161, backend='torchxla_trace_once')
  for data, _ in loader:
    output = dynamo_densenet161(data)

TorchServe

제공된 torchserve-tpu Docker 이미지를 사용하여 TPU VM에서 보관처리된 pytorch 모델을 서빙할 수 있습니다.

Docker 인증을 설정합니다.

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
    us-docker.pkg.dev

Cloud TPU TorchServe Docker 이미지를 TPU VM으로 가져옵니다.

CLOUD_TPU_TORCHSERVE_IMAGE_URL=us-docker.pkg.dev/cloud-tpu-images/inference/torchserve-tpu:v0.9.0-2.1
docker pull ${CLOUD_TPU_TORCHSERVE_IMAGE_URL}

모델 아티팩트 수집

시작하려면 모델 핸들러를 제공해야 합니다. 이 핸들러는 TorchServe 모델 서버 작업자가 모델을 로드하고, 입력 데이터를 처리하고, 추론을 실행하도록 지시합니다. TorchServe 기본 추론 핸들러(소스)를 사용할 수도 있고, base_handler.py에 따라 자체 커스텀 모델 핸들러를 개발할 수도 있습니다. 학습된 모델 및 모델 정의 파일을 제공해야 할 수도 있습니다.

다음 Densenet 161 예시에서는 TorchServe에서 제공하는 모델 아티팩트와 기본 이미지 분류기 핸들러를 사용합니다.

몇 가지 환경 변수를 구성합니다.

CWD="$(pwd)"

WORKDIR="${CWD}/densenet_161"

mkdir -p ${WORKDIR}/model-store
mkdir -p ${WORKDIR}/logs

TorchServe 이미지 분류기 예시에서 모델 아티팩트를 다운로드하고 복사합니다.

git clone https://github.com/pytorch/serve.git

cp ${CWD}/serve/examples/image_classifier/densenet_161/model.py ${WORKDIR}
cp ${CWD}/serve/examples/image_classifier/index_to_name.json ${WORKDIR}

모델 가중치를 다운로드합니다.

wget https://download.pytorch.org/models/densenet161-8d451a50.pth -O densenet161-8d451a50.pth

mv densenet161-8d451a50.pth ${WORKDIR}

Dynamo 백엔드를 사용할 TorchServe 모델 구성 파일을 만듭니다.

echo 'pt2: "torchxla_trace_once"' >> ${WORKDIR}/model_config.yaml

다음 파일 및 디렉터리가 표시됩니다.

>> ls ${WORKDIR}
model_config.yaml
index_to_name.json
logs
model.py
densenet161-8d451a50.pth
model-store

모델 보관 파일 생성

Cloud TPU TorchServe를 사용하여 PyTorch 모델을 서빙하려면 Torch Model Archiver를 사용하여 모델 핸들러와 모든 모델 아티팩트를 모델 아카이브 파일(*.mar)로 패키징해야 합니다.

torch-model-archiver를 사용하여 모델 보관 파일을 생성합니다.

MODEL_NAME=Densenet161

docker run \
    --privileged  \
    --shm-size 16G \
    --name torch-model-archiver \
    -it \
    -d \
    --rm \
    --mount type=bind,source=${WORKDIR},target=/home/model-server/ \
    ${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \
    torch-model-archiver \
        --model-name ${MODEL_NAME} \
        --version 1.0 \
        --model-file model.py \
        --serialized-file densenet161-8d451a50.pth \
        --handler image_classifier \
        --export-path model-store \
        --extra-files index_to_name.json \
        --config-file model_config.yaml

model-store 디렉터리에 모델 보관 파일이 생성되어야 합니다.

>> ls ${WORKDIR}/model-store
Densenet161.mar

추론 요청 처리

이제 모델 보관 파일을 만들었으므로 TorchServe 모델 서버를 시작하고 추론 요청을 처리할 수 있습니다.

TorchServe 모델 서버를 시작합니다.

docker run \
    --privileged  \
    --shm-size 16G \
    --name torchserve-tpu \
    -it \
    -d \
    --rm \
    -p 7070:7070 \
    -p 7071:7071 \
    -p 8080:8080 \
    -p 8081:8081 \
    -p 8082:8082 \
    -p 9001:9001 \
    -p 9012:9012 \
    --mount type=bind,source=${WORKDIR}/model-store,target=/home/model-server/model-store \
    --mount type=bind,source=${WORKDIR}/logs,target=/home/model-server/logs \
    ${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \
    torchserve \
        --start \
        --ncs \
        --models ${MODEL_NAME}.mar \
        --ts-config /home/model-server/config.properties

모델 서버 상태를 쿼리합니다.

curl http://localhost:8080/ping

모델 서버가 실행 중이면 다음과 같이 표시됩니다.

{
  "status": "Healthy"
}

현재 등록된 모델의 기본 버전을 쿼리하려면 다음을 사용합니다.

curl http://localhost:8081/models

등록된 모델이 표시됩니다.

{
  "models": [
    {
      "modelName": "Densenet161",
      "modelUrl": "Densenet161.mar"
    }
  ]
}

추론용 이미지를 다운로드하려면 다음을 사용합니다.

curl -O https://raw.githubusercontent.com/pytorch/serve/master/docs/images/kitten_small.jpg

mv kitten_small.jpg ${WORKDIR}

모델 서버에 추론 요청을 보내려면 다음을 사용합니다.

curl http://localhost:8080/predictions/${MODEL_NAME} -T ${WORKDIR}/kitten_small.jpg

다음과 비슷한 응답이 표시됩니다.

{
  "tabby": 0.47878125309944153,
  "lynx": 0.20393909513950348,
  "tiger_cat": 0.16572578251361847,
  "tiger": 0.061157409101724625,
  "Egyptian_cat": 0.04997897148132324
}

모델 서버 로그

다음 명령어를 사용하여 로그에 액세스합니다.
```
ls ${WORKDIR}/logs/
cat ${WORKDIR}/logs/model_log.log
```
로그에 다음 메시지가 표시됩니다.
```
"Compiled model with backend torchxla\_trace\_once"
```

삭제

Docker 컨테이너를 중지합니다.

rm -rf serve
rm -rf ${WORKDIR}

docker stop torch-model-archiver
docker stop torchserve-tpu

프로파일링

추론을 설정한 후 프로파일러를 사용하여 성능 및 TPU 사용률을 분석할 수 있습니다. 프로파일링에 대한 자세한 내용은 다음을 참조하세요.

Cloud TPU v5e 추론 소개

개요 및 이점

시작하기

Cloud TPU 모델 추론 및 서빙

JAX 모델 추론 및 서빙

SavedModel 형식

추론 변환기

추론 변환기 요구사항

JAX 모델 추론 및 서빙

jax2tf를 사용하여 JAX 모델을 SavedModel 형식으로 직렬화

추론 변환기를 사용하여 저장된 모델을 서빙할 수 있도록 준비

TensorFlow Serving 사용

JAX 모델 서빙 예시

기본 요건

추론을 위한 JAX BERT 모델 서빙

추론을 위한 JAX Stable Diffusion 서빙

TensorFlow Serving

TensorFlow Serving 워크플로

TensorFlow ResNet-50 서빙 데모 실행

PyTorch 모델 추론 및 서빙

TorchDynamo 및 PyTorch/XLA

TorchServe

모델 아티팩트 수집

모델 보관 파일 생성

추론 요청 처리

삭제

프로파일링

`jax2tf`를 사용하여 JAX 모델을 SavedModel 형식으로 직렬화