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, hay tres formas de optimizar tu modelo:

  • Optimización manual del modelo: usas Inference Converter para optimizar el modelo y guardarlo. Después, debes pasar las marcas --saved_model_tags='serve,tpu' y --disable_optimizer=true cuando upload tu modelo. Por ejemplo:

    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",
    ]
)

  • Optimización automática de modelos con partición automática: cuando importas un modelo, Vertex AI intenta optimizarlo mediante un algoritmo de partición automática. Esta optimización no funciona en todos los modelos. Si la optimización falla, debes optimizar manualmente el modelo o elegir la optimización automática del modelo con partición manual. Por ejemplo:

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

  • Optimización automática de modelos con partición manual. Especifica la marca --converter_options_string y ajusta ConverterOptions.TpuFunction para que se adapte a tus necesidades. Para ver un ejemplo, consulta Imagen del convertidor. Ten en cuenta que solo se admite ConverterOptions.TpuFunction, que es lo único que se necesita para la partición manual. Por ejemplo:

    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 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, en TorchServe para la inferencia de Cloud TPU v5e se muestra cómo empaquetar el modelo Densenet 161 en artefactos de modelo mediante Torch Model Archiver.

A continuación, sube los artefactos del modelo a tu carpeta de Cloud Storage y sube el 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