Cloud TPU v5e의 SAX
SAX 클러스터(SAX 셀)
SAX 관리 서버 및 SAX 모델 서버는 SAX 클러스터를 실행하는 두 가지 필수 구성요소입니다.
SAX 관리 서버
SAX 관리 서버는 SAX 클러스터의 모든 SAX 모델 서버를 모니터링하고 조정합니다. SAX 클러스터에서 여러 SAX 관리 서버를 실행할 수 있으며, 여기서 SAX 관리 서버 중 하나만 리더 선택을 통해 활성화되고 나머지는 대기 서버입니다. 활성 관리 서버가 실패하면 대기 관리 서버가 활성화됩니다. 활성 SAX 관리 서버가 사용 가능한 SAX 모델 서버에 모델 복제본 및 추론 요청을 할당합니다.
SAX 관리 스토리지 버킷
각 SAX 클러스터에는 SAX 관리 서버와 SAX 모델 서버의 구성 및 위치를 SAX 클러스터에 저장하기 위해 Cloud Storage 버킷이 필요합니다.
SAX 모델 서버
SAX 모델 서버가 모델 체크포인트를 로드하고 GSPMD로 추론을 실행합니다. SAX 모델 서버는 단일 TPU VM 작업자에서 실행됩니다. 단일 호스트 TPU 모델 서빙을 위해서는 단일 호스트 TPU VM에 단일 SAX 모델 서버가 필요합니다. 멀티 호스트 TPU 모델 서빙을 위해서는 멀티 호스트 TPU 슬라이스에서 SAX 모델 서버 그룹이 필요합니다. 멀티 호스트 모델 서빙은 현재 제공되지 않지만, 이 문서에서는 미리보기를 위해 175B 테스트 모델이 포함된 예시를 제공합니다.
SAX 모델 서빙
다음 섹션에서는 SAX를 사용하여 언어 모델을 서빙하는 워크플로를 안내합니다. 단일 호스트 모델 서빙의 예로 GPT-J 6B 모델을 사용합니다.
시작하기 전에 TPU VM에 Cloud TPU SAX Docker 이미지를 설치합니다.
sudo usermod -a -G docker ${USER} newgrp docker gcloud auth configure-docker us-docker.pkg.dev SAX_ADMIN_SERVER_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server" SAX_MODEL_SERVER_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server" SAX_UTIL_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/inference/sax-util" SAX_VERSION=v1.0.0 export SAX_ADMIN_SERVER_IMAGE_URL=${SAX_ADMIN_SERVER_IMAGE_NAME}:${SAX_VERSION} export SAX_MODEL_SERVER_IMAGE_URL=${SAX_MODEL_SERVER_IMAGE_NAME}:${SAX_VERSION} export SAX_UTIL_IMAGE_URL="${SAX_UTIL_IMAGE_NAME}:${sax_version}" docker pull ${SAX_ADMIN_SERVER_IMAGE_URL} docker pull ${SAX_MODEL_SERVER_IMAGE_URL} docker pull ${SAX_UTIL_IMAGE_URL}
나중에 사용할 다른 변수를 설정합니다.
export SAX_ADMIN_SERVER_DOCKER_NAME="sax-admin-server" export SAX_MODEL_SERVER_DOCKER_NAME="sax-model-server" export SAX_CELL="/sax/test"
GPT-J 6B 단일 호스트 모델 서빙 예시
단일 호스트 모델 서빙은 단일 호스트 TPU 슬라이스, 즉 v5litepod-1, v5litepod-4, v5litepod-8에 적용할 수 있습니다.
SAX 클러스터 만들기
SAX 클러스터용 Cloud Storage 스토리지 버킷을 만듭니다.
SAX_ADMIN_STORAGE_BUCKET=${your_admin_storage_bucket} gcloud storage buckets create gs://${SAX_ADMIN_STORAGE_BUCKET} \ --project=${PROJECT_ID}
체크포인트를 저장하기 위해 다른 Cloud Storage 스토리지 버킷이 필요할 수 있습니다.
SAX_DATA_STORAGE_BUCKET=${your_data_storage_bucket}
터미널에서 TPU VM에 SSH를 통해 SAX 관리 서버를 실행합니다.
docker run \ --name ${SAX_ADMIN_SERVER_DOCKER_NAME} \ -it \ -d \ --rm \ --network host \ --env GSBUCKET=${SAX_ADMIN_STORAGE_BUCKET} \ ${SAX_ADMIN_SERVER_IMAGE_URL}
다음 방법으로 Docker 로그를 확인할 수 있습니다.
docker logs -f ${SAX_ADMIN_SERVER_DOCKER_NAME}
로그의 출력이 다음과 비슷하게 표시됩니다.
I0829 01:22:31.184198 7 config.go:111] Creating config fs_root: "gs://test_sax_admin/sax-fs-root" I0829 01:22:31.347883 7 config.go:115] Created config fs_root: "gs://test_sax_admin/sax-fs-root" I0829 01:22:31.360837 24 admin_server.go:44] Starting the server I0829 01:22:31.361420 24 ipaddr.go:39] Skipping non-global IP address 127.0.0.1/8. I0829 01:22:31.361455 24 ipaddr.go:39] Skipping non-global IP address ::1/128. I0829 01:22:31.361462 24 ipaddr.go:39] Skipping non-global IP address fe80::4001:aff:fe8e:fc8/64. I0829 01:22:31.361469 24 ipaddr.go:39] Skipping non-global IP address fe80::42:bfff:fef9:1bd3/64. I0829 01:22:31.361474 24 ipaddr.go:39] Skipping non-global IP address fe80::20fb:c3ff:fe5b:baac/64. I0829 01:22:31.361482 24 ipaddr.go:56] IPNet address 10.142.15.200 I0829 01:22:31.361488 24 ipaddr.go:56] IPNet address 172.17.0.1 I0829 01:22:31.456952 24 admin.go:305] Loaded config: fs_root: "gs://test_sax_admin/sax-fs-root" I0829 01:22:31.609323 24 addr.go:105] SetAddr /gcs/test_sax_admin/sax-root/sax/test/location.proto "10.142.15.200:10000" I0829 01:22:31.656021 24 admin.go:325] Updated config: fs_root: "gs://test_sax_admin/sax-fs-root" I0829 01:22:31.773245 24 mgr.go:781] Loaded manager state I0829 01:22:31.773260 24 mgr.go:784] Refreshing manager state every 10s I0829 01:22:31.773285 24 admin.go:350] Starting the server on port 10000 I0829 01:22:31.773292 24 cloud.go:506] Starting the HTTP server on port 8080
단일 호스트 SAX 모델 서버를 SAX 클러스터에 실행합니다.
현재 SAX 클러스터에는 SAX 관리 서버만 포함됩니다. 두 번째 터미널에서 SSH를 통해 TPU VM에 연결하여 SAX 클러스터에서 SAX 모델 서버를 실행할 수 있습니다.
docker run \ --privileged \ -it \ -d \ --rm \ --network host \ --name ${SAX_MODEL_SERVER_DOCKER_NAME} \ --env SAX_ROOT=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \ ${SAX_MODEL_SERVER_IMAGE_URL} \ --sax_cell=${SAX_CELL} \ --port=10001 \ --platform_chip=tpuv4 \ --platform_topology=1x1
모델 체크포인트 변환:
EleutherAI에서 GPT-J 체크포인트를 다운로드하려면 PyTorch 및 Transformer를 설치해야 합니다.
pip3 install accelerate pip3 install torch pip3 install transformers
체크포인트를 SAX 체크포인트로 변환하려면
paxml을 설치해야 합니다.pip3 install paxml==1.1.0
다음 스크립트는 GPT-J 체크포인트를 SAX 체크포인트로 변환합니다.
python3 -m convert_gptj_ckpt --base EleutherAI/gpt-j-6b --pax pax_6b
변환이 완료된 후:
ls checkpoint_00000000/commit_success 파일을 만들어 하위 디렉터리에 배치해야 합니다.
gcloud storage cp checkpoint_00000000 ${CHECKPOINT_PATH} --recursive touch commit_success.txt gcloud storage cp commit_success.txt ${CHECKPOINT_PATH}/ gcloud storage cp commit_success.txt ${CHECKPOINT_PATH}/metadata/ gcloud storage cp commit_success.txt ${CHECKPOINT_PATH}/state/SAX 클러스터에 모델 게시
이제 이전 단계에서 변환된 체크포인트로 GPT-J를 게시할 수 있습니다.
MODEL_NAME=gptjtokenizedbf16bs32 MODEL_CONFIG_PATH=saxml.server.pax.lm.params.gptj.GPTJ4TokenizedBF16BS32 REPLICA=1
GPT-J(및 이후 단계)를 게시하려면 세 번째 터미널에서 TPU VM에 SSH를 통해 연결합니다.
docker run \ ${SAX_UTIL_IMAGE_URL} \ --sax_root=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \ publish \ ${SAX_CELL}/${MODEL_NAME} \ ${MODEL_CONFIG_PATH} \ ${CHECKPOINT_PATH} \ ${REPLICA}
모델이 성공적으로 로드되었음을 나타내는 다음과 같은 출력이 표시될 때까지 모델 서버 Docker 로그에 많은 활동이 표시됩니다.
I0829 01:33:49.287459 139865140229696 servable_model.py:697] loading completed.
추론 결과 생성
GPT-J의 경우 입력 및 출력이 쉼표로 구분된 토큰 ID 문자열 형식이어야 합니다. 텍스트 입력을 토큰화해야 합니다.
TEXT = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction\:\nSummarize the following news article\:\n\n### Input\:\nMarch 10, 2015 . We're truly international in scope on Tuesday. We're visiting Italy, Russia, the United Arab Emirates, and the Himalayan Mountains. Find out who's attempting to circumnavigate the globe in a plane powered partially by the sun, and explore the mysterious appearance of craters in northern Asia. You'll also get a view of Mount Everest that was previously reserved for climbers. On this page you will find today's show Transcript and a place for you to request to be on the CNN Student News Roll Call. TRANSCRIPT . Click here to access the transcript of today's CNN Student News program. Please note that there may be a delay between the time when the video is available and when the transcript is published. CNN Student News is created by a team of journalists who consider the Common Core State Standards, national standards in different subject areas, and state standards when producing the show. ROLL CALL . For a chance to be mentioned on the next CNN Student News, comment on the bottom of this page with your school name, mascot, city and state. We will be selecting schools from the comments of the previous show. You must be a teacher or a student age 13 or older to request a mention on the CNN Student News Roll Call! Thank you for using CNN Student News!\n\n### Response\:
EleutherAI/gpt-j-6b 토크나이저를 통해 토큰 ID 문자열을 가져올 수 있습니다.
from transformers import GPT2Tokenizer tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-j-6b") :
입력 텍스트를 토큰화합니다.
encoded_example = tokenizer(TEXT) input_ids = encoded_example.input_ids INPUT_STR = ",".join([str(input_id) for input_id in input_ids])
토큰 ID 문자열은 다음과 비슷할 수 있습니다.
>>> INPUT_STR '21106,318,281,12064,326,8477,257,4876,11,20312,351,281,5128,326,3769,2252,4732,13,19430,257,2882,326,20431,32543,262,2581,13,198,198,21017,46486,25,198,13065,3876,1096,262,1708,1705,2708,25,198,198,21017,23412,25,198,16192,838,11,1853,764,775,821,4988,3230,287,8354,319,3431,13,775,821,10013,8031,11,3284,11,262,1578,4498,24880,11,290,262,42438,22931,21124,13,9938,503,508,338,9361,284,2498,4182,615,10055,262,13342,287,257,6614,13232,12387,416,262,4252,11,290,7301,262,11428,5585,286,1067,8605,287,7840,7229,13,921,1183,635,651,257,1570,286,5628,41336,326,373,4271,10395,329,39311,13,1550,428,2443,345,481,1064,1909,338,905,42978,290,257,1295,329,345,284,2581,284,307,319,262,8100,13613,3000,8299,4889,13,48213,6173,46023,764,6914,994,284,1895,262,14687,286,1909,338,8100,13613,3000,1430,13,4222,3465,326,612,743,307,257,5711,1022,262,640,618,262,2008,318,1695,290,618,262,14687,318,3199,13,8100,13613,3000,318,2727,416,257,1074,286,9046,508,2074,262,8070,7231,1812,20130,11,2260,5423,287,1180,2426,3006,11,290,1181,5423,618,9194,262,905,13,15107,3069,42815,764,1114,257,2863,284,307,4750,319,262,1306,8100,13613,3000,11,2912,319,262,4220,286,428,2443,351,534,1524,1438,11,37358,11,1748,290,1181,13,775,481,307,17246,4266,422,262,3651,286,262,2180,905,13,921,1276,307,257,4701,393,257,3710,2479,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,0,6952,345,329,1262,8100,13613,3000,0,198,198,21017,18261,25'
자료 요약을 생성하려면 다음 안내를 따르세요.
docker run \ ${SAX_UTIL_IMAGE_URL} \ --sax_root=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \ lm.generate \ ${SAX_CELL}/${MODEL_NAME} \ ${INPUT_STR}다음과 비슷하게 표시됩니다.
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+ | GENERATE | SCORE | +----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+ | 1212,2443,3407,262,905,42978,764,198,11041,262,42978,284,1037,2444,351,3555,35915,290,25818,764,198,2953,262,4220,286,262,2443,11,2912,329,257,2863,284,307,4750,319,8100,13613,3000,13,220,921,1276,307,257,4701,393,257,3710,2479,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,13,50256 | -0.023136413 | | 1212,2443,3407,262,905,42978,764,198,11041,262,42978,284,1037,2444,351,3555,35915,290,25818,764,198,2953,262,4220,286,262,2443,11,2912,329,257,2863,284,307,4750,319,8100,13613,3000,13,220,921,1276,307,257,4701,393,257,3710,2479,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,0,50256 | -0.91842502 | | 1212,2443,3407,262,905,42978,764,198,11041,262,42978,284,1037,2444,351,3555,35915,290,25818,764,198,2953,262,4220,286,262,2443,11,2912,329,257,2863,284,307,4750,319,8100,13613,3000,13,921,1276,307,257,4701,393,257,3710,2479,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,13,50256 | -1.1726116 | | 1212,2443,3407,262,905,42978,764,198,11041,262,42978,284,1037,2444,351,3555,35915,290,25818,764,198,2953,262,4220,286,262,2443,11,2912,329,257,2863,284,307,4750,319,8100,13613,3000,13,220,921,1276,307,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,13,50256 | -1.2472695 | +----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+
출력 토큰 ID 문자열을 역토큰화하려면 다음을 실행합니다.
output_token_ids = [int(token_id) for token_id in OUTPUT_STR.split(',')] OUTPUT_TEXT = tokenizer.decode(output_token_ids, skip_special_tokens=True)역토큰화된 텍스트는 다음과 같이 표시될 수 있습니다.
>>> OUTPUT_TEXT 'This page includes the show Transcript.\nUse the Transcript to help students with reading comprehension and vocabulary.\nAt the bottom of the page, comment for a chance to be mentioned on CNN Student News. You must be a teacher or a student age 13 or older to request a mention on the CNN Student News Roll Call.'
Docker 컨테이너 및 Cloud Storage 스토리지 버킷을 삭제합니다.
175B 멀티 호스트 모델 서빙 미리보기
일부 대규모 언어 모델에는 멀티 호스트 TPU 슬라이스, 즉 v5litepod-16 이상이 필요합니다. 이러한 경우 모든 멀티 호스트 TPU 호스트는 SAX 모델 서버 사본이 있어야 하며 모든 모델 서버는 SAX 모델 서버 그룹으로 작동하여 멀티 호스트 TPU 슬라이스에서 대규모 모델을 서빙합니다.
새 SAX 클러스터 만들기
GPT-J 둘러보기에서 SAX 클러스터 만들기 단계를 똑같이 수행하여 새 SAX 클러스터와 SAX 관리 서버를 만들 수 있습니다.
또는 기존 SAX 클러스터가 이미 있으면 SAX 클러스터에 멀티 호스트 모델 서버를 실행할 수 있습니다.
SAX 클러스터로 멀티 호스트 SAX 모델 서버 실행
단일 호스트 TPU 슬라이스에 사용하는 것과 동일한 명령어를 사용하고 적절한 멀티 호스트 가속기 유형을 지정하여 멀티 호스트 TPU 슬라이스를 만듭니다.
ACCELERATOR_TYPE=v5litepod-32 ZONE=us-east1-c gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \ --node-id ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --accelerator-type ${ACCELERATOR_TYPE} \ --runtime-version ${RUNTIME_VERSION} \ --service-account ${SERVICE_ACCOUNT} \ --reservedSAX 모델 서버 이미지를 모든 TPU 호스트/작업자로 가져와서 실행하려면 다음 안내를 따르세요.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE} \ --worker=all \ --command=" gcloud auth configure-docker \ us-docker.pkg.dev # Pull sax model server image docker pull ${SAX_MODEL_SERVER_IMAGE_URL} # Run model server docker run \ --privileged \ -it \ -d \ --rm \ --network host \ --name ${SAX_MODEL_SERVER_DOCKER_NAME} \ --env SAX_ROOT=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \ ${SAX_MODEL_SERVER_IMAGE_URL} \ --sax_cell=${SAX_CELL} \ --port=10001 \ --platform_chip=tpuv4 \ --platform_topology=1x1"SAX 클러스터에 모델 게시
이 예시에서는 LmCloudSpmd175B32Test 모델을 사용합니다.
MODEL_NAME=lmcloudspmd175b32test MODEL_CONFIG_PATH=saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test CHECKPOINT_PATH=None REPLICA=1
테스트 모델을 게시하려면 다음 안내를 따르세요.
docker run \ ${SAX_UTIL_IMAGE_URL} \ --sax_root=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \ publish \ ${SAX_CELL}/${MODEL_NAME} \ ${MODEL_CONFIG_PATH} \ ${CHECKPOINT_PATH} \ ${REPLICA}
추론 결과 생성
docker run \ ${SAX_UTIL_IMAGE_URL} \ --sax_root=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \ lm.generate \ ${SAX_CELL}/${MODEL_NAME} \ "Q: Who is Harry Porter's mother? A\: "이 예시에서는 임의의 가중치가 있는 테스트 모델을 사용하므로 출력이 의미 없을 수 있습니다.
삭제
Docker 컨테이너를 중지합니다.
docker stop ${SAX_ADMIN_SERVER_DOCKER_NAME} docker stop ${SAX_MODEL_SERVER_DOCKER_NAME}
아래와 같이 gcloud CLI를 사용하여 Cloud Storage 관리자 스토리지 버킷과 모든 데이터 스토리지 버킷을 삭제합니다.
gcloud storage rm gs://${SAX_ADMIN_STORAGE_BUCKET} --recursive --continue-on-error gcloud storage rm gs://${SAX_DATA_STORAGE_BUCKET} --recursive --continue-on-error