SAX no Cloud TPU v5e

Cluster do SAX (célula SAX)

O servidor de administração e o servidor de modelos do SAX são dois componentes essenciais que executam um cluster do SAX.

Servidor de administração do SAX

O servidor de administração do SAX monitora e coordena todos os servidores de modelos do SAX em um cluster do SAX. Nesse cluster, é possível iniciar vários servidores de administração do SAX, em que apenas um deles está ativo por eleição de líder e os outros estão em espera. Quando o servidor de administração ativo falha, um daqueles em espera se torna ativo. O servidor de administração do SAX ativo atribui réplicas de modelo e solicitações de inferência a servidores de modelos do SAX disponíveis.

Bucket de armazenamento do administrador do SAX

Cada cluster do SAX requer um bucket do Cloud Storage para armazenar as configurações e os locais dos servidores de administração e de modelos do SAX que ele contém.

Servidor de modelos do SAX

O servidor de modelos do SAX carrega um checkpoint de modelo e executa a inferência com o GSPMD. Ele é executado em um único worker de VM de TPU. A disponibilização de modelos de TPU de host único exige um único servidor de modelos do SAX em uma VM de TPU de host único. A disponibilização de modelos de TPU de vários hosts exige um grupo de servidores de modelos do SAX em uma fração de TPU de vários hosts. No momento, não é possível usar a disponibilização de modelos de vários hosts, mas este documento oferece um exemplo com um modelo de teste 175B para visualização.

Disponibilização de modelos do SAX

A seção a seguir explica o fluxo de trabalho para disponibilizar modelos de linguagem usando o SAX. Ela usa o modelo GPT-J 6B como exemplo para a disponibilização de modelos de host único.

Antes de começar, instale as imagens do Docker do SAX para Cloud TPU na VM de TPU:

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}

Defina outras variáveis para uso futuro:

export SAX_ADMIN_SERVER_DOCKER_NAME="sax-admin-server"
export SAX_MODEL_SERVER_DOCKER_NAME="sax-model-server"
export SAX_CELL="/sax/test"

Exemplo de disponibilização do modelo de host único GPT-J 6B

A disponibilização de modelos de host único é aplicável a frações de TPU de host único, ou seja, v5litepod-1, v5litepod-4 e v5litepod-8.

  1. Criar um cluster do SAX

    1. Crie um bucket de armazenamento do Cloud Storage para o cluster do SAX:

      SAX_ADMIN_STORAGE_BUCKET=${your_admin_storage_bucket}

gcloud storage buckets create gs://${SAX_ADMIN_STORAGE_BUCKET} \
--project=${PROJECT_ID}

      Talvez seja necessário outro bucket do Cloud Storage para armazenar o checkpoint.

      SAX_DATA_STORAGE_BUCKET=${your_data_storage_bucket}

    2. Conecte-se por SSH à VM de TPU em um terminal para iniciar o servidor de administração do 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}

      Para verificar o registro do Docker, faça o seguinte: 

      docker logs -f ${SAX_ADMIN_SERVER_DOCKER_NAME}

      A saída no registro será parecida com esta:

      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

  2. Iniciar um servidor de modelos do SAX de host único no cluster do SAX

    Até aqui, o cluster do SAX contém apenas o servidor de administração do SAX. É possível se conectar à VM de TPU por SSH em um segundo terminal para iniciar um servidor de modelos do SAX no cluster do 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

  3. Converter o checkpoint do modelo

    É necessário instalar o PyTorch e o Transformers para fazer o download do checkpoint do GPT-J da EleutherAI:

    pip3 install accelerate
pip3 install torch
pip3 install transformers

    Para converter esse checkpoint em um checkpoint do SAX, instale paxml:

    pip3 install paxml==1.1.0

    O script abaixo converte o checkpoint do GPT-J em um checkpoint do SAX:

    python3 -m convert_gptj_ckpt --base EleutherAI/gpt-j-6b --pax pax_6b

    Depois da conversão: 

    ls checkpoint_00000000/

    Você precisa criar um arquivo commit_success e colocá-lo nos subdiretórios:

    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/

  4. Publicar o modelo no cluster do SAX

    Agora você pode publicar o GPT-J com o checkpoint convertido na etapa anterior.

    MODEL_NAME=gptjtokenizedbf16bs32
MODEL_CONFIG_PATH=saxml.server.pax.lm.params.gptj.GPTJ4TokenizedBF16BS32
REPLICA=1

    Para publicar o GPT-J (e as etapas seguintes), conecte-se por SSH à VM de TPU usando um terceiro terminal:

    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}

    Você vai notar muita atividade no registro do Docker do servidor de modelos até perceber algo como o seguinte, que indica que o modelo foi carregado:

    I0829 01:33:49.287459 139865140229696 servable_model.py:697] loading completed.

  5. Gerar os resultados de inferência

    Para o GPT-J, a entrada e a saída precisam ser formatadas como uma string de IDs de token separados por vírgulas. Você vai precisar tokenizar a entrada de texto.

    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\:

    Você pode usar o tokenizador EleutherAI/gpt-j-6b para ter acesso à string de IDs de token:

    from transformers import GPT2Tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-j-6b")                  :

    Tokenize o texto de entrada:

    encoded_example = tokenizer(TEXT)
input_ids = encoded_example.input_ids
INPUT_STR = ",".join([str(input_id) for input_id in input_ids])

    Você vai receber uma string de IDs de token parecida com esta:

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

    Para gerar um resumo do artigo:

    docker run \
  ${SAX_UTIL_IMAGE_URL} \
    --sax_root=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \
    lm.generate \
      ${SAX_CELL}/${MODEL_NAME} \
      ${INPUT_STR}

    Você vai receber algo parecido com isto:

    +----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+
|                                                                                                                                                    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 |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+

    Para retirar a tokenização da string de IDs de tokens de saída, faça o seguinte:

    output_token_ids = [int(token_id) for token_id in OUTPUT_STR.split(',')]
OUTPUT_TEXT = tokenizer.decode(output_token_ids, skip_special_tokens=True)

    O texto sem tokenização será assim:

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

  6. Limpar os contêineres do Docker e os buckets de armazenamento do Cloud Storage

Pré-lançamento da disponibilização do modelo de vários hosts 175B

Alguns modelos de linguagem grandes exigem uma fração de TPU de vários hosts, ou seja, v5litepod-16 ou mais recente. Nesses casos, todos os hosts de TPU de vários hosts precisam ter uma cópia de um servidor de modelos do SAX, e esses servidores funcionam como um grupo de servidores de modelos do SAX para disponibilizar o modelo grande em uma fração de TPU de vários hosts.

  1. Criar um cluster do SAX

    Siga a mesma etapa do tutorial do GPT-Jque ensina a criar um cluster do SAX para criar um novo cluster do SAX e um servidor de administração do SAX.

    Se você já tiver um cluster do SAX, inicie um servidor de modelos de vários hosts nele.

  2. Iniciar um servidor de modelos do SAX de vários hosts em um cluster do SAX

    Use o mesmo comando que é usado para criar uma fração de TPU de host único para criar uma de vários hosts. Para isso, especifique o tipo de acelerador de vários hosts adequado:

    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} \
  --reserved

    Para extrair a imagem do servidor de modelos do SAX para todos os hosts/workers de TPU e iniciá-los, faça o seguinte:

    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"

  3. Publicar o modelo no cluster do SAX

    Este exemplo usa um modelo LmCloudSpmd175B32Test:

    MODEL_NAME=lmcloudspmd175b32test
MODEL_CONFIG_PATH=saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test
CHECKPOINT_PATH=None
REPLICA=1

    Para publicar o modelo de teste, faça o seguinte:

    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}

  4. Gerar os resultados de inferência

    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\: "

    Como o exemplo usa um modelo de teste com pesos aleatórios, a saída pode não ser relevante.

  5. Limpeza

    Interrompa os contêineres do Docker:

    docker stop ${SAX_ADMIN_SERVER_DOCKER_NAME}
docker stop ${SAX_MODEL_SERVER_DOCKER_NAME}

    Exclua o bucket de armazenamento do administrador do Cloud Storage e qualquer bucket de armazenamento de dados usando a gcloud CLI, conforme mostrado abaixo.

    gcloud storage rm gs://${SAX_ADMIN_STORAGE_BUCKET} --recursive --continue-on-error
gcloud storage rm gs://${SAX_DATA_STORAGE_BUCKET} --recursive --continue-on-error