Implantar um modelo nas VMs do Cloud TPU

O Google Cloud fornece acesso a aceleradores de machine learning personalizados, chamados Unidades de Processamento de Tensor (TPUs). As TPUs são otimizadas para acelerar o treinamento e a inferência de modelos de machine learning, o que as torna ideais para uma variedade de aplicativos, como processamento de linguagem natural, visão computacional e reconhecimento de fala.

Nesta página, você vai aprender a implantar modelos em um host único do Cloud TPU v5e para previsão on-line na Vertex AI.

Só há suporte para o Cloud TPU v5e. Outras gerações do Cloud TPU não têm suporte.

Importar o modelo

Para a implantação em Cloud TPUs, é preciso importar o modelo para a Vertex AI e configurá-lo a fim de usar um dos seguintes contêineres:

Contêiner de ambiente de execução otimizado e predefinido do TensorFlow

Para importar e executar um TensorFlow SavedModel em um Cloud TPU, o modelo precisa ser otimizado para TPU. Se o SavedModel do TensorFlow ainda não está otimizado para TPU, há três maneiras de otimizar o modelo:

  • Otimização manual do modelo: use o conversor de inferência para otimizar e salvar o modelo. Em seguida, transmita as flags --saved_model_tags='serve,tpu' e --disable_optimizer=true ao upload o modelo. Exemplo:

    model = aiplatform.Model.upload(
    display_name='Manually optimized model',
    artifact_uri="gs://model-artifact-uri",
    serving_container_image_uri="us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest",
    serving_container_args=[
        "--saved_model_tags=serve,tpu",
        "--disable_optimizer=true",
    ]
)

  • Otimização automática do modelo com particionamento automático: quando você importa um modelo, a Vertex AI tenta otimizar esse modelo não otimizado usando um algoritmo de particionamento automático. Essa otimização não funciona em todos os modelos. Se a otimização falhar, otimize manualmente o modelo ou escolha a otimização automática com particionamento manual. Exemplo:

    model = aiplatform.Model.upload(
    display_name='TPU optimized model with automatic partitioning',
    artifact_uri="gs://model-artifact-uri",
    serving_container_image_uri="us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest",
    serving_container_args=[
    ]
)

  • Otimização automática do modelo com particionamento manual. Especifique a flag --converter_options_string e ajuste o ConverterOptions.TpuFunction de acordo com suas necessidades. Para um exemplo, consulte Converter uma imagem. Observe que só há suporte para ConverterOptions.TpuFunction, que é tudo o que é necessário para o particionamento manual. Exemplo:

    model = aiplatform.Model.upload(
display_name='TPU optimized model with manual partitioning',
  artifact_uri="gs://model-artifact-uri",
  serving_container_image_uri="us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-tpu.2-15:latest",
  serving_container_args=[
      "--converter_options_string='tpu_functions { function_alias: \"partitioning function name\" }'"
  ]
)

Para mais informações sobre como importar modelos, consulte Como importar modelos para a Vertex AI.

Contêiner PyTorch predefinido

As instruções para importar e executar um modelo PyTorch no Cloud TPU são iguais às seguidas para importar e executar um modelo PyTorch.

Por exemplo, a inferência do TorchServe para Cloud TPU v5e demonstra como empacotar o modelo Densenet 161 em artefatos de modelo usando o arquivador de modelo do Torch.

Em seguida, faça o upload dos artefatos do modelo para a pasta do Cloud Storage e do modelo conforme abaixo:

model = aiplatform.Model.upload(
    display_name='DenseNet TPU model from SDK PyTorch 2.1',
    artifact_uri="gs://model-artifact-uri",
    serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/pytorch-tpu.2-1:latest",
    serving_container_args=[],
    serving_container_predict_route="/predictions/model",
    serving_container_health_route="/ping",
    serving_container_ports=[8080]
)

Para mais informações, consulte Exportar artefatos de modelo para o PyTorch e o Notebook do Jupyter para Exibir um modelo PyTorch usando um contêiner predefinido.

Contêiner personalizado

No caso de contêineres personalizados, o modelo não precisa ser do TensorFlow, mas precisa ser otimizado para TPU. Para informações sobre como produzir um modelo otimizado para TPU, consulte os seguintes guias de frameworks de ML comuns:

Para informações sobre a exibição de modelos treinados com JAX, TensorFlow ou PyTorch no Cloud TPU v5e, consulte Inferência do Cloud TPU v5e.

Verifique se o contêiner personalizado atende aos requisitos de contêiner personalizado.

É preciso aumentar o limite de memória bloqueada para que o driver possa se comunicar com os chips de TPU por acesso direto à memória (DMA). Exemplo:

Linha de comando

ulimit -l 68719476736

Python

import resource

resource.setrlimit(
    resource.RLIMIT_MEMLOCK,
    (
        68_719_476_736_000,  # soft limit
        68_719_476_736_000,  # hard limit
    ),
  )

Em seguida, consulte Usar um contêiner personalizado para previsão a fim de conferir informações sobre como importar um modelo com um contêiner personalizado. Para implementar a lógica de pré ou pós-processamento, use Rotinas de previsão personalizadas.

Crie um endpoint

As instruções para criar um endpoint para Cloud TPUs são iguais às seguidas a fim de criar qualquer endpoint.

Por exemplo, o seguinte comando cria um recurso endpoint:

endpoint = aiplatform.Endpoint.create(display_name='My endpoint')

A resposta contém o ID do novo endpoint, que você vai usar nas próximas etapas.

Para mais informações sobre como criar um endpoint, consulte Implantar um modelo em um endpoint.

Implantar um modelo

As instruções para implantar um modelo em Cloud TPUs são iguais às seguidas com qualquer modelo, com a exceção de que é preciso especificar um dos seguintes tipos de máquina do Cloud TPU com suporte:

Tipo de máquina Número de chips do TPU
ct5lp-hightpu-1t 1
ct5lp-hightpu-4t 4
ct5lp-hightpu-8t 8

Os aceleradores de TPU são integrados ao tipo de máquina. Você não precisa especificar o tipo ou a contagem de aceleradores.

Por exemplo, o seguinte comando implanta um modelo chamando deployModel:

machine_type = 'ct5lp-hightpu-1t'

deployed_model = model.deploy(
    endpoint=endpoint,
    deployed_model_display_name='My deployed model',
    machine_type=machine_type,
    traffic_percentage=100,
    min_replica_count=1
    sync=True,
)

Para mais informações, consulte Implantar um modelo em um endpoint.

Receber previsões on-line

A instrução para receber previsões on-line de um Cloud TPU é igual à seguida para receber previsões on-line.

Por exemplo, o seguinte comando envia uma solicitação de previsão on-line chamando predict:

deployed_model.predict(...)

Para contêineres personalizados, consulte os requisitos de solicitação e resposta de previsão para contêineres personalizados.

Como garantir a capacidade

Por padrão, a cota para Custom model serving TPU v5e cores per region é 0.

Para solicitar um aumento, consulte Solicitar um limite de cota maior.

Preços

Os tipos de máquina de TPU são cobrados por hora, assim como todos os outros tipos de máquina na Vertex Prediction. Para mais informações, consulte Preços da previsão.

A seguir