Desplegar un modelo en VMs de TPU de Cloud

Google Cloud proporciona acceso a aceleradores de aprendizaje automático diseñados a medida, denominados unidades de procesamiento de tensor (TPUs). Las TPUs se han optimizado para acelerar la preparación y la inferencia de los modelos de aprendizaje automático, por lo que son ideales para diversas aplicaciones, como el procesamiento del lenguaje natural, la visión artificial y el reconocimiento de voz.

En esta página se describe cómo desplegar tus modelos en un solo host de Cloud TPU v5e o v6e para la inferencia online en Vertex AI.

Solo se admiten las versiones 5e y 6e de TPU de Cloud. No se admiten otras generaciones de Cloud TPU.

Para saber en qué ubicaciones están disponibles las versiones 5e y 6e de TPU de Cloud, consulta las ubicaciones.

Importar el modelo

Para desplegar tu modelo en Cloud TPUs, debes importarlo a Vertex AI y configurarlo para que use uno de los siguientes contenedores:

Contenedor de entorno de ejecución de TensorFlow optimizado prediseñado

Para importar y ejecutar un modelo de TensorFlowSavedModel en una TPU de Cloud, el modelo debe estar optimizado para TPUs. Si tu TensorFlow SavedModel aún no está optimizado para TPU, puedes optimizar tu modelo automáticamente. Para ello, importa tu modelo y, a continuación, Vertex AI optimizará tu modelo no optimizado mediante un algoritmo de partición automática. Esta optimización no funciona en todos los modelos. Si la optimización falla, debes optimizar el modelo manualmente.

El siguiente código de muestra muestra cómo usar la optimización automática de modelos con particiones automáticas:

  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=[
      ]
  )

Para obtener más información sobre cómo importar modelos, consulta el artículo sobre importar modelos a Vertex AI.

Contenedor de PyTorch prediseñado

Las instrucciones para importar y ejecutar un modelo de PyTorch en una TPU de Cloud son las mismas que las instrucciones para importar y ejecutar un modelo de PyTorch.

Por ejemplo, TorchServe para la inferencia de TPU de Cloud v5e A continuación, sube los artefactos del modelo a tu carpeta de Cloud Storage y sube tu modelo como se muestra a continuación:

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 obtener más información, consulta el artículo sobre cómo exportar artefactos de modelo para PyTorch y el cuaderno del tutorial sobre cómo servir un modelo de PyTorch con un contenedor precompilado.

Contenedor personalizado

En el caso de los contenedores personalizados, no es necesario que tu modelo sea un modelo de TensorFlow, pero debe estar optimizado para TPUs. Para obtener información sobre cómo crear un modelo optimizado para TPUs, consulta las siguientes guías de los frameworks de aprendizaje automático más habituales:

Para obtener información sobre cómo servir modelos entrenados con JAX, TensorFlow o PyTorch en Cloud TPU v5e, consulta Inferencia de Cloud TPU v5e.

Asegúrate de que tu contenedor personalizado cumpla los requisitos de los contenedores personalizados.

Debes aumentar el límite de memoria bloqueada para que el controlador pueda comunicarse con los chips de TPU mediante el acceso directo a la memoria (DMA). Por ejemplo:

Línea 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
    ),
  )

A continuación, consulta Usar un contenedor personalizado para la inferencia para obtener información sobre cómo importar un modelo con un contenedor personalizado. Si quieres implementar una lógica de preprocesamiento o postprocesamiento, te recomendamos que uses rutinas de inferencia personalizadas.

Crear un punto final

Las instrucciones para crear un endpoint para las TPUs de Cloud son las mismas que las instrucciones para crear cualquier endpoint.

Por ejemplo, el siguiente comando crea un recurso endpoint:

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

La respuesta contiene el ID del nuevo endpoint, que usarás en los pasos posteriores.

Para obtener más información sobre cómo crear un endpoint, consulta el artículo Implementar un modelo en un endpoint.

Desplegar un modelo

Las instrucciones para implementar un modelo en las TPUs de Cloud son las mismas que las instrucciones para implementar cualquier modelo, excepto que debes especificar uno de los siguientes tipos de máquinas de TPU de Cloud admitidos:

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

Los aceleradores de TPU están integrados en el tipo de máquina. No tienes que especificar el tipo ni el número de aceleradores.

Por ejemplo, el siguiente comando implementa un modelo llamando a 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 obtener más información, consulta Desplegar un modelo en un endpoint.

Obtener inferencias online

Las instrucciones para obtener inferencias online de una TPU de Cloud son las mismas que las instrucciones para obtener inferencias online.

Por ejemplo, el siguiente comando envía una solicitud de inferencia online llamando a predict:

deployed_model.predict(...)

En el caso de los contenedores personalizados, consulta los requisitos de las solicitudes y respuestas de inferencia de los contenedores personalizados.

Asegurar la capacidad

En la mayoría de las regiones, la TPU v5e and v6e cores per region cuota de servicio de modelos personalizados es 0. En algunas zonas, está limitada.

Para solicitar un aumento de cuota, consulta Solicitar un ajuste de cuota.

Precios

Los tipos de máquinas de TPU se facturan por hora, al igual que todos los demás tipos de máquinas de Vertex Prediction. Para obtener más información, consulta los precios de las predicciones.

Siguientes pasos