Implementa un modelo en las VMs de Cloud TPU

Google Cloud proporciona acceso a aceleradores de aprendizaje automático de diseño personalizado, 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 una TPU v5e de host único de Cloud para la predicción en línea en Vertex AI.

Solo se admite la versión v5e 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 usar uno de los siguientes contenedores:

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

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

  • Optimización manual de modelos: Usas 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 particiones automáticas: Cuando importes un modelo, Vertex AI intentará optimizarlo con un algoritmo de particiones automáticas. Esta optimización no funciona en todos los modelos. Si la optimización falla, debes optimizar el 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 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 todo lo que se necesita para el particionamiento 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 Cómo importar modelos a Vertex AI.

Contenedor de PyTorch compilado previamente

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

Por ejemplo, TorchServe para la inferencia de Cloud TPU v5e muestra cómo empaquetar el modelo Densenet 161 en artefactos de modelos con el archivador de modelos de Torch.

Luego, 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 Cómo exportar artefactos de modelos para PyTorch y el notebook de Jupyter para entregar un modelo de PyTorch con un contenedor compilado previamente.

Contenedor personalizado

En el caso de los contenedores personalizados, tu modelo no necesita ser un modelo de TensorFlow, pero debe estar optimizado para TPU. Para obtener información sobre cómo producir un modelo optimizado para TPU, consulta las siguientes guías de frameworks de AA comunes:

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

Debes aumentar el límite de memoria bloqueada para que el controlador pueda comunicarse con los chips 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 si necesitas información para importar un modelo con un contenedor personalizado. Si quieres implementar una lógica de procesamiento previo o posterior, considera usar las rutinas de predicción personalizadas.

Crea un extremo

Las instrucciones para crear un extremo para Cloud TPU son las mismas que las instrucciones para 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 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 para implementar cualquier modelo, excepto que debes especificar uno de los siguientes tipos de máquinas compatibles con Cloud TPU:

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 ni el recuento 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 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 para obtener predicciones en línea.

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

deployed_model.predict(...)

En el caso de los contenedores personalizados, consulta los requisitos de solicitud y respuesta de predicción para contenedores personalizados.

Cómo proteger la capacidad

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

Para solicitar un aumento, consulta Solicita un límite de cuota más alto.

Precios

Los tipos de máquinas de TPU se facturan por hora, al igual que todos los demás tipo de máquina en Vertex Prediction. Para obtener más información, consulta Precios de la predicción.

¿Qué sigue?