En este instructivo, se muestra cómo entregar un modelo de lenguaje grande (LLM) Gemma con unidades de procesamiento tensorial (TPU) en Google Kubernetes Engine (GKE). Implementas un contenedor prediseñado con JetStream y MaxText en GKE. También configurarás GKE para que cargue las ponderaciones de Gemma 7B desde Cloud Storage en el tiempo de ejecución.
Este instructivo está dirigido a ingenieros de aprendizaje automático (AA), administradores y operadores de plataformas, y especialistas en datos y en IA que estén interesados en usar las capacidades de organización de contenedores de Kubernetes para entregar LLM. Para obtener más información sobre los roles comunes y las tareas de ejemplo a las que hacemos referencia en el contenido deGoogle Cloud , consulta Roles de usuario y tareas comunes de GKE.
Antes de leer esta página, asegúrate de estar familiarizado con lo siguiente:
- Modo Autopilot y modo Standard
- Disponibilidad de la versión actual de TPU con la arquitectura del sistema de Cloud TPU
- TPUs en GKE
Fondo
En esta sección, se describen las tecnologías clave que se usan en este instructivo.
Gemma
Gemma es un conjunto de modelos de Inteligencia Artificial (IA) básicos y de disponibilidad general que se lanzan con una licencia abierta. Estos modelos de IA están disponibles para ejecutarse en tus aplicaciones, hardware, dispositivos móviles o servicios alojados. Puedes usar los modelos de Gemma para la generación de texto, pero también puedes ajustar estos modelos en el caso de tareas especializadas.
Para obtener más información, consulta la documentación de Gemma.
TPU
Las TPU son circuitos integrados personalizados específicos de aplicaciones (ASIC) de Google que se usan para acelerar el aprendizaje automático y los modelos de IA compilados con frameworks como el siguiente:TensorFlow, PyTorch yJAX.
En este instructivo se aborda la entrega del modelo Gemma 7B. GKE implementa el modelo en los nodos TPUv5e de host único con topologías de TPU configuradas según los requisitos del modelo para entregar mensajes con baja latencia.
JetStream
JetStream es un framework de entrega de inferencia de código abierto que desarrolla Google. JetStream permite la inferencia de alto rendimiento, alta capacidad de procesamiento y con optimización de memoria en TPU y GPU. Proporciona optimizaciones de rendimiento avanzadas, incluidas técnicas de procesamiento por lotes y de cuantización continuas, para facilitar la implementación de LLM. JetStream permite que PyTorch/XLA y JAX TPU entregue un rendimiento óptimo.
Para obtener más información sobre estas optimizaciones, consulta los repositorios de proyectos de JetStream PyTorch y JetStream MaxText.
MaxText
MaxText es una implementación de LLM JAX escalable y adaptable, compilada en bibliotecas JAX de código abierto como Flax, Orbax y Optax. La implementación de LLM solo de decodificador de MaxText se escribe en Python. Aprovecha en gran medida el compilador XLA para lograr un alto rendimiento sin necesidad de compilar kernels personalizados.
Para obtener más información sobre los modelos y tamaños de parámetros más recientes que admite MaxText, consulta el repositorio del proyecto de MaxtText.