Implementa un modelo en las VM de Cloud TPU

Google Cloud proporciona acceso a aceleradores de aprendizaje automático personalizados llamados unidades de procesamiento tensorial (TPU). TPUs están optimizados para acelerar el entrenamiento y la inferencia de modelos de aprendizaje automático, lo que los hace ideales para una variedad de aplicaciones, incluyendo procesamiento de lenguaje natural, visión por computadora y reconocimiento de voz.

En esta página, se describe cómo implementar tus modelos en un host único de Cloud TPU v5e para la predicción en línea en Vertex AI.

Solo se admite la versión 5 de Cloud TPU. No se admiten otras generaciones de Cloud TPU.

Importa tu modelo

Para la implementación en Cloud TPU, debes importar tu modelo a Vertex AI y configurarlo para que use uno de los siguientes contenedores:

Contenedor de entorno de ejecución optimizado de TensorFlow compilado previamente

Para importar y ejecutar un SavedModel en Cloud TPU, el modelo debe estar optimizado para TPU. Si tu SavedModel de TensorFlow aún no está optimizado para TPU, hay tres formas de optimizar tu modelo:

  • Optimización manual del modelo: Usa el Convertidor de inferencia para optimizar tu modelo y guardarlo. Luego, 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 intentará optimizar tu modelo no optimizado con un algoritmo de partición automática. Esta optimización no funciona en todos los modelos. Si la optimización falla, debes optimizar tu modelo de forma manual o elegir la optimización automática del modelo con particiones manuales. 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 el ConverterOptions.TpuFunction según tus necesidades. Para ver un ejemplo, consulta Imagen de convertidor. Ten en cuenta que solo se admite ConverterOptions.TpuFunction, que es todo lo 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 la importación de modelos, consulta Importa modelos a Vertex AI.

Contenedor PyTorch compilado con anterioridad

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

Por ejemplo, TorchServe para Cloud TPU v5e inferencia muestra cómo empaquetar el modelo Densenet 161 en artefactos de modelo mediante Torch Model Archiver.

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

Si deseas obtener más información, consulta Exporta artefactos de modelo para PyTorch y el Notebook de Jupyter para Entrega un modelo de PyTorch mediante un contenedor compilado previamente.

Contenedor personalizado

No es necesario que tu modelo sea de TensorFlow, pero debe estar optimizado para TPU. Si deseas obtener información a fin de producir un modelo optimizado para TPU, consulta las siguientes guías para marcos de trabajo de AA comunes:

Para obtener información sobre cómo entregar 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 con los requisitos de contenedores personalizados.

Debes aumentar el límite de memoria bloqueada para que el controlador pueda comunicarse con los chips de TPU a través del 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
    ),
  )

Luego, consulta Usa un contenedor personalizado para la predicción a fin de obtener información sobre cómo importar un modelo con un contenedor personalizado. Si deseas implementar la lógica de procesamiento previo o posterior, considera usar rutinas de predicción personalizadas.

Crea un extremo

Las instrucciones a fin de crear un extremo para Cloud TPU son las mismas que las instrucciones a fin de crear cualquier extremo.

Por ejemplo, el siguiente comando crea un recurso endpoint:

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

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

Para obtener más información sobre cómo crear un extremo, consulta Implementa un modelo en un extremo.

Implementar un modelo

Las instrucciones para implementar un modelo en Cloud TPU son las mismas que las instrucciones a fin de implementar cualquier modelo, excepto que especificas uno de los siguientes tipos de máquinas de Cloud TPU compatibles:

Tipo de máquina Cantidad de chips TPU
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 es necesario que especifiques el tipo de acelerador ni el recuento de aceleradores.

Por ejemplo, el siguiente comando implementa un modelo mediante una llamada 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 Implementa un modelo en un extremo.

Obtén predicciones en línea

La instrucción para obtener predicciones en línea desde una Cloud TPU es la misma que la instrucción que se usa a fin de obtener predicciones en línea.

Por ejemplo, el siguiente comando envía una solicitud de predicción en línea mediante una llamada a predict:

deployed_model.predict(...)

Para contenedores personalizados, consulta los requisitos de solicitud de predicción y respuesta para contenedores personalizados.

Protege la capacidad

De forma predeterminada, la cuota para Custom model serving TPU v5e cores per region es 0.

Para solicitar un aumento, consulta Solicita un límite de cuota mayor.

Precios

Los tipos de máquinas 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 Precios de predicción.

¿Qué sigue?