Implemente um modelo em VMs do Cloud TPU

Google Cloud oferece acesso a aceleradores de aprendizagem automática concebidos à medida denominados unidades de processamento tensor (TPUs). As TPUs estão otimizadas para acelerar a preparação e a inferência de modelos de aprendizagem automática, o que as torna ideais para uma variedade de aplicações, incluindo o processamento de linguagem natural, a visão computacional e o reconhecimento de voz.

Esta página descreve como implementar os seus modelos num único anfitrião Cloud TPU v5e ou v6e para inferência online na Vertex AI.

Apenas são suportadas as versões v5e e v6e da Cloud TPU. Não são suportadas outras gerações de TPUs do Google Cloud.

Para saber em que localizações estão disponíveis as versões v5e e v6e da Cloud TPU, consulte as localizações.

Importe o seu modelo

Para a implementação em TPUs do Google Cloud, tem de importar o seu modelo para a Vertex AI e configurá-lo para usar um dos seguintes contentores:

Contentor de tempo de execução do TensorFlow otimizado pré-criado

Para importar e executar um TensorFlow SavedModel numa Cloud TPU, o modelo tem de estar otimizado para TPU. Se o seu TensorFlow SavedModel ainda não estiver otimizado para TPUs, existem três formas de otimizar o seu modelo:

  • Otimização manual do modelo: usa o conversor de inferência para otimizar o modelo e guardá-lo. Em seguida, tem de transmitir as flags --saved_model_tags='serve,tpu' e --disable_optimizer=true quando upload o modelo. Por 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 de modelos com partição automática: quando importa um modelo, o Vertex AI tenta otimizar o seu modelo não otimizado através de um algoritmo de partição automática. Esta otimização não funciona em todos os modelos. Se a otimização falhar, tem de otimizar manualmente o modelo ou escolher a otimização automática do modelo com a partição manual. Por 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 partição manual. Especifique a flag --converter_options_string e ajuste o ConverterOptions.TpuFunction de acordo com as suas necessidades. Para ver um exemplo, consulte Converter imagem. Tenha em atenção que apenas o ConverterOptions.TpuFunction, que é tudo o que é necessário para a partição manual, é suportado. Por 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 a importação de modelos, consulte o artigo sobre a importação de modelos para o Vertex AI.

Contentor PyTorch pré-criado

As instruções para importar e executar um modelo do PyTorch na Cloud TPU são as mesmas que as instruções para importar e executar um modelo do PyTorch.

Por exemplo, o TorchServe para inferência de TPU na nuvem v5e demonstra como criar pacotes do modelo Densenet 161 em artefactos de modelo usando o Torch Model Archiver.

Em seguida, carregue os artefactos do modelo para a pasta do Cloud Storage e carregue o modelo da seguinte forma:

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 o artigo Exporte artefactos de modelos para o PyTorch e o bloco de notas do tutorial para Apresente um modelo do PyTorch com um contentor pré-criado.

Contentor personalizado

Para contentores personalizados, o seu modelo não tem de ser um modelo do TensorFlow, mas tem de estar otimizado para a TPU. Para obter informações sobre como produzir um modelo otimizado para TPU, consulte os seguintes guias para frameworks de ML comuns:

Para informações sobre como publicar modelos preparados com JAX, TensorFlow ou PyTorch na Cloud TPU v5e, consulte o artigo Inferência da Cloud TPU v5e.

Certifique-se de que o contentor personalizado cumpre os requisitos do contentor personalizado.

Tem de aumentar o limite de memória bloqueada para que o controlador possa comunicar com os chips da TPU através do acesso direto à memória (DMA). Por exemplo:

Linha de comandos

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 o artigo Use um contentor personalizado para inferência para obter informações sobre a importação de um modelo com um contentor personalizado. Se quiser implementar uma lógica de pré ou pós-processamento, considere usar rotinas de inferência personalizadas.

Crie um ponto final

As instruções para criar um ponto final para as TPUs na nuvem são as mesmas que as instruções para criar qualquer ponto final.

Por exemplo, o comando seguinte cria um recurso endpoint:

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

A resposta contém o ID do novo ponto final, que usa nos passos subsequentes.

Para mais informações sobre como criar um ponto final, consulte o artigo Implemente um modelo num ponto final.

Implemente um modelo

As instruções para implementar um modelo em TPUs na nuvem são as mesmas que as instruções para implementar qualquer modelo, exceto que especifica um dos seguintes tipos de máquinas de TPUs na nuvem suportados:

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

Os aceleradores da TPU estão incorporados no tipo de máquina. Não tem de especificar o tipo de acelerador nem a quantidade de aceleradores.

Por exemplo, o seguinte comando implementa 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 o artigo sobre como implementar um modelo num ponto final.

Obtenha inferências online

A instrução para obter inferências online a partir de uma Cloud TPU é a mesma que a instrução para obter inferências online.

Por exemplo, o comando seguinte envia um pedido de inferência online chamando predict:

deployed_model.predict(...)

Para contentores personalizados, consulte os requisitos de pedidos e respostas de inferência para contentores personalizados.

Garantir capacidade

Para a maioria das regiões, a TPU v5e and v6e cores per region quota para o serviço de modelos personalizados é 0. Em algumas regiões, está limitado.

Para pedir um aumento da quota, consulte o artigo Peça um ajuste da quota.

Preços

Os tipos de máquinas de TPU são faturados por hora, tal como todos os outros tipos de máquinas na Vertex AI Prediction. Para mais informações, consulte a secção Preços de previsão.

O que se segue?