Inferência do JetStream MaxText em VMs TPU v6e

Este tutorial mostra como usar o JetStream para disponibilizar modelos MaxText na TPU v6e. O JetStream é um mecanismo otimizado para capacidade de processamento e memória para inferência de modelos de linguagem grandes (LLMs) em dispositivos XLA (TPUs). Neste tutorial, você vai executar o comparativo de inferência para o modelo Llama2-7B.

Antes de começar

Prepare-se para provisionar uma TPU v6e com 4 chips:

  1. Siga o guia Configurar o ambiente do Cloud TPU para configurar um projeto Google Cloud , configurar a Google Cloud CLI, ativar a API Cloud TPU e garantir que você tenha acesso para usar o Cloud TPU.

  2. Faça a autenticação com Google Cloud e configure o projeto e a zona padrão para a Google Cloud CLI.

    gcloud auth login gcloud config set project PROJECT_ID gcloud config set compute/zone ZONE

Capacidade segura

Quando estiver tudo pronto para garantir a capacidade da TPU, consulte Cotas da Cloud TPU para mais informações. Se você tiver outras dúvidas sobre como garantir a capacidade, entre em contato com a equipe de vendas ou de conta do Cloud TPU.

Provisionar o ambiente do Cloud TPU

É possível provisionar VMs do TPU com o GKE, com o GKE e o XPK ou como recursos em fila.

Pré-requisitos

  • Verifique se o projeto tem cota de TPUS_PER_TPU_FAMILY suficiente, que especifica o número máximo de chips que você pode acessar no projetoGoogle Cloud .
  • Verifique se o projeto tem cota suficiente de TPU para:
    • Cota da VM de TPU
    • Cota de endereços IP
    • Quota do Hyperdisk equilibrado
  • Permissões do projeto do usuário

Criar variáveis de ambiente

No Cloud Shell, crie as seguintes variáveis de ambiente:

export PROJECT_ID=your-project-id export TPU_NAME=your-tpu-name export ZONE=us-east5-b export ACCELERATOR_TYPE=v6e-4 export RUNTIME_VERSION=v2-alpha-tpuv6e export SERVICE_ACCOUNT=your-service-account export QUEUED_RESOURCE_ID=your-queued-resource-id

Descrições das variáveis de ambiente

Variável Descrição
PROJECT_ID O ID do seu Google Cloud projeto. Use um projeto existente ou crie um novo.
TPU_NAME O nome da TPU.
ZONE A zona em que a VM TPU será criada. Para mais informações sobre as zonas com suporte, consulte Regiões e zonas de TPU.
ACCELERATOR_TYPE O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.
RUNTIME_VERSION A versão do software do Cloud TPU.
SERVICE_ACCOUNT O endereço de e-mail da sua conta de serviço. Para encontrar essa página, acesse a página de contas de serviço no console Google Cloud .

Por exemplo: tpu-service-account@PROJECT_ID.iam.gserviceaccount.com

QUEUED_RESOURCE_ID O ID de texto atribuído pelo usuário da solicitação de recurso em fila.

Provisionar um TPU v6e

Use o comando a seguir para provisionar uma TPU v6e:

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}

Use os comandos list ou describe para consultar o status do recurso em fila.

gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \     --project ${PROJECT_ID} --zone ${ZONE}

Para mais informações sobre os status de solicitações de recursos em fila, consulte Gerenciar recursos em fila.

Conectar-se à TPU usando SSH

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}

Depois de se conectar à TPU, você pode executar o comparativo de mercado de inferência.

Configurar o ambiente da VM do TPU

  1. Crie um diretório para executar o comparativo de mercado de inferência:

    export MAIN_DIR=your-main-directory mkdir -p ${MAIN_DIR}

  2. Configure um ambiente virtual de Python:

    cd ${MAIN_DIR} sudo apt update sudo apt install python3.10 python3.10-venv python3.10 -m venv venv source venv/bin/activate

  3. Instale o armazenamento de arquivos grandes (LFS) do Git (para dados do OpenOrca):

    sudo apt-get install git-lfs git lfs install

  4. Clone e instale o JetStream:

    cd $MAIN_DIR git clone https://github.com/google/JetStream.git cd JetStream git checkout main pip install -e . cd benchmarks pip install -r requirements.in

  5. Configurar o MaxText:

    cd $MAIN_DIR git clone https://github.com/google/maxtext.git cd maxtext git checkout main bash setup.sh pip install torch --index-url https://download.pytorch.org/whl/cpu

  6. Solicitar acesso aos modelos do Llama para receber uma chave de download do Meta para o modelo Llama 2.

  7. Clone o repositório Llama:

    cd $MAIN_DIR git clone https://github.com/meta-llama/llama cd llama

  8. Execute bash download.sh. Quando solicitado, forneça sua chave de download. Esse script cria um diretório llama-2-7b dentro do diretório llama. 

    bash download.sh

  9. Crie buckets de armazenamento:

    export CHKPT_BUCKET=gs://your-checkpoint-bucket export BASE_OUTPUT_DIRECTORY=gs://your-output-dir export CONVERTED_CHECKPOINT_PATH=gs://bucket-to-store-converted-checkpoints export MAXTEXT_BUCKET_UNSCANNED=gs://bucket-to-store-unscanned-data gcloud storage buckets create ${CHKPT_BUCKET} gcloud storage buckets create ${BASE_OUTPUT_DIRECTORY} gcloud storage buckets create ${CONVERTED_CHECKPOINT_PATH} gcloud storage buckets create ${MAXTEXT_BUCKET_UNSCANNED} gcloud storage cp --recursive llama-2-7b/* ${CHKPT_BUCKET}

Realizar a conversão do checkpoint

  1. Realize a conversão para os checkpoints digitalizados:

    cd $MAIN_DIR/maxtext python3 -m MaxText.llama_or_mistral_ckpt \     --base-model-path $MAIN_DIR/llama/llama-2-7b \     --model-size llama2-7b \     --maxtext-model-path ${CONVERTED_CHECKPOINT_PATH}

  2. Converta em checkpoints não verificados:

    export CONVERTED_CHECKPOINT=${CONVERTED_CHECKPOINT_PATH}/0/items export DIRECT_PARAMETER_CHECKPOINT_RUN=direct_generate_param_only_checkpoint python3 -m MaxText.generate_param_only_checkpoint \     MaxText/configs/base.yml \     base_output_directory=${MAXTEXT_BUCKET_UNSCANNED} \     load_parameters_path=${CONVERTED_CHECKPOINT} \     run_name=${DIRECT_PARAMETER_CHECKPOINT_RUN} \     model_name='llama2-7b' \     force_unroll=true

Realizar inferência

  1. Execute um teste de validação:

    export UNSCANNED_CKPT_PATH=${MAXTEXT_BUCKET_UNSCANNED}/${DIRECT_PARAMETER_CHECKPOINT_RUN}/checkpoints/0/items python3 -m MaxText.decode \     MaxText/configs/base.yml \     load_parameters_path=${UNSCANNED_CKPT_PATH} \     run_name=runner_decode_unscanned_${idx} \     base_output_directory=${BASE_OUTPUT_DIRECTORY} \     per_device_batch_size=1 \     model_name='llama2-7b' \     ici_autoregressive_parallelism=4 \     max_prefill_predict_length=4 \     max_target_length=16 \     prompt="I love to" \     attention=dot_product \     scan_layers=false

  2. Execute o servidor no terminal atual:

    export TOKENIZER_PATH=assets/tokenizer.llama2 export LOAD_PARAMETERS_PATH=${UNSCANNED_CKPT_PATH} export MAX_PREFILL_PREDICT_LENGTH=1024 export MAX_TARGET_LENGTH=2048 export MODEL_NAME=llama2-7b export ICI_FSDP_PARALLELISM=1 export ICI_AUTOREGRESSIVE_PARALLELISM=1 export ICI_TENSOR_PARALLELISM=-1 export SCAN_LAYERS=false export WEIGHT_DTYPE=bfloat16 export PER_DEVICE_BATCH_SIZE=11  cd $MAIN_DIR/maxtext python3 -m MaxText.maxengine_server \     MaxText/configs/base.yml \     tokenizer_path=${TOKENIZER_PATH} \     load_parameters_path=${LOAD_PARAMETERS_PATH} \     max_prefill_predict_length=${MAX_PREFILL_PREDICT_LENGTH} \     max_target_length=${MAX_TARGET_LENGTH} \     model_name=${MODEL_NAME} \     ici_fsdp_parallelism=${ICI_FSDP_PARALLELISM} \     ici_autoregressive_parallelism=${ICI_AUTOREGRESSIVE_PARALLELISM} \     ici_tensor_parallelism=${ICI_TENSOR_PARALLELISM} \     scan_layers=${SCAN_LAYERS} \     weight_dtype=${WEIGHT_DTYPE} \     per_device_batch_size=${PER_DEVICE_BATCH_SIZE}

  3. Abra uma nova janela de terminal, conecte-se à TPU e mude para o mesmo ambiente virtual usado na primeira janela de terminal:

    source venv/bin/activate

  4. Execute os comandos a seguir para executar o comparativo de mercado do JetStream.

    export MAIN_DIR=your-main-directory cd $MAIN_DIR  python JetStream/benchmarks/benchmark_serving.py \     --tokenizer $MAIN_DIR/maxtext/assets/tokenizer.llama2 \     --warmup-mode sampled \     --save-result \     --save-request-outputs \     --request-outputs-file-path outputs.json \     --num-prompts 1000 \     --max-output-length 1024 \     --dataset openorca \     --dataset-path $MAIN_DIR/JetStream/benchmarks/open_orca_gpt4_tokenized_llama.calibration_1000.pkl

Resultados

A saída a seguir foi gerada ao executar o comparativo de mercado usando a v6e-8. Os resultados vão variar de acordo com o hardware, o software, o modelo e a rede.

Mean output size: 929.5959798994975 Median output size: 1026.0 P99 output size: 1026.0 Successful requests: 995 Benchmark duration: 195.533269 s Total input tokens: 217011 Total generated tokens: 924948 Request throughput: 5.09 requests/s Input token throughput: 1109.84 tokens/s Output token throughput: 4730.39 tokens/s Overall token throughput: 5840.23 tokens/s Mean ttft: 538.49 ms Median ttft: 95.66 ms P99 ttft: 13937.86 ms Mean ttst: 1218.72 ms Median ttst: 152.57 ms P99 ttst: 14241.30 ms Mean TPOT: 91.83 ms Median TPOT: 16.63 ms P99 TPOT: 363.37 ms

Limpar

  1. Desconecte da TPU:

    $ (vm) exit

  2. Exclua a TPU:

    gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \     --project ${PROJECT_ID} \     --zone ${ZONE} \     --force \     --async

  3. Exclua os buckets e o conteúdo deles:

    export CHKPT_BUCKET=gs://your-checkpoint-bucket export BASE_OUTPUT_DIRECTORY=gs://your-output-dir export CONVERTED_CHECKPOINT_PATH=gs://bucket-to-store-converted-checkpoints export MAXTEXT_BUCKET_UNSCANNED=gs://bucket-to-store-unscanned-data gcloud storage rm -r ${CHKPT_BUCKET} gcloud storage rm -r ${BASE_OUTPUT_DIRECTORY} gcloud storage rm -r ${CONVERTED_CHECKPOINT_PATH} gcloud storage rm -r ${MAXTEXT_BUCKET_UNSCANNED}