Entrega modelos abiertos con el contenedor premium Hex-LLM en Cloud TPU

Hex-LLM, un modelo grande de lenguaje (LLM) de alta eficiencia que entrega XLA, es el framework de entrega de Vertex AI LLM que está diseñado y optimizado para el hardware de Cloud TPU. Hex-LLM combina tecnologías de entrega de LLM, como el procesamiento por lotes continuo y PagedAttention con optimizaciones de Vertex AI que están diseñadas para XLA y Cloud TPU. Es un LLM de alta eficiencia y bajo costo que se entrega en Cloud TPU para modelos de código abierto.

Hex-LLM está disponible en Model Garden a través del entorno de juego del modelo, la implementación con un clic y el notebook.

Funciones

Hex-LLM se basa en proyectos de código abierto con las propias optimizaciones de Google para XLA y Cloud TPU. Hex-LLM logra una alta capacidad de procesamiento y baja latencia cuando entrega LLM de uso frecuente.

Hex-LLM incluye las siguientes optimizaciones:

  • Algoritmo de procesamiento por lotes continuo basado en tokens para ayudar a garantizar que los modelos usen por completo el hardware con una gran cantidad de solicitudes simultáneas.
  • Una reescritura completa de los kernels de atención que están optimizados para XLA
  • Estrategias de paralelismo de datos flexibles y compuestas, y paralelismo de tensor con métodos de fragmentación de peso altamente optimizados para ejecutar LLM en varios chips de Cloud TPU de manera eficiente.

Hex-LLM admite una amplia gama de LLM densos y dispersos:

  • Gemma 2B y 7B
  • Gemma 2 9B y 27B
  • Llama 2 7B, 13B y 70B
  • Llama 3 8B y 70B
  • Mistral 7B y Mixtral 8x7B

Hex-LLM también proporciona una variedad de funciones, como las siguientes:

  • Hex-LLM se incluye en un solo contenedor. Hex-LLM empaqueta el servidor de la API, el motor de inferencia y los modelos compatibles en una sola imagen de Docker que se implementará.
  • Compatible con el formato de los modelos de Hugging Face. Hex-LLM puede cargar un modelo de Hugging Face desde el disco local, el Hugging Face Hub y un bucket de Cloud Storage.
  • Cuantización mediante bitsandbytes y AWQ.
  • Carga dinámica de LoRA Hex-LLM puede cargar las ponderaciones de LoRA a través de la lectura del argumento de solicitud durante la entrega.

Primeros pasos con Model Garden

El contenedor de entrega de Cloud TPU de Hex-LLM está integrado en Model Garden. Puedes acceder a esta tecnología de entrega a través de los entornos de prueba, implementación con un clic y ejemplos de notebooks de Colab Enterprise para una variedad de modelos.

Usa la zona de pruebas

La zona de pruebas de Model Garden es un extremo de Vertex AI implementado con anterioridad al que se puede acceder mediante el envío de solicitudes en la tarjeta de modelo.

  1. Ingresa una instrucción y, de manera opcional, incluye argumentos para tu solicitud.

  2. Haz clic en ENVIAR para obtener la respuesta del modelo con rapidez.

Pruébalo con Gemma.

Usa la implementación con un solo clic

Puedes implementar un extremo personalizado de Vertex AI con Hex-LLM mediante una tarjeta de modelo.

  1. Navega a la página de la tarjeta de modelo y haz clic en Implementar.

  2. Para la variación del modelo que deseas usar, selecciona el tipo de máquina de Cloud TPU v5e para la implementación.

  3. Haz clic en Implementar en la parte inferior para comenzar el proceso de implementación. Recibirás dos notificaciones por correo electrónico; uno cuando se sube el modelo y otro cuando el extremo está listo.

Usa el notebook de Colab Enterprise

Si deseas obtener flexibilidad y personalización, puedes usar ejemplos de notebook de Colab Enterprise para implementar un extremo de Vertex AI con Hex-LLM mediante el SDK de Vertex AI para Python.

  1. Navega a la página de la tarjeta de modelo y haz clic en Abrir notebook.

  2. Selecciona el notebook de Vertex Serving. Se abre el notebook en Colab Enterprise.

  3. Ejecuta el notebook para implementar un modelo mediante Hex-LLM y envía solicitudes de predicción al extremo. El fragmento de código para la implementación es el siguiente:

hexllm_args = [
    f"--model=google/gemma-2-9b-it",
    f"--tensor_parallel_size=4",
    f"--hbm_utilization_factor=0.8",
    f"--max_running_seqs=512",
]
hexllm_envs = {
    "PJRT_DEVICE": "TPU",
    "MODEL_ID": "google/gemma-2-9b-it",
    "DEPLOY_SOURCE": "notebook",
}
model = aiplatform.Model.upload(
    display_name="gemma-2-9b-it",
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=[
        "python", "-m", "hex_llm.server.api_server"
    ],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=hexllm_envs,
    serving_container_shared_memory_size_mb=(16 * 1024),
    serving_container_deployment_timeout=7200,
)

endpoint = aiplatform.Endpoint.create(display_name="gemma-2-9b-it-endpoint")
model.deploy(
    endpoint=endpoint,
    machine_type="ct5lp-hightpu-4t",
    deploy_request_timeout=1800,
    service_account="<your-service-account>",
    min_replica_count=1,
    max_replica_count=1,
)

Puedes modificar los siguientes argumentos de inicio del servidor Hex-LLM para la entrega personalizada:

  • --model: El modelo que se cargará. Puedes especificar un ID de modelo de Hugging Face, una ruta de acceso absoluta local o una ruta de bucket de Cloud Storage.
  • --tokenizer: El tokenizador que se cargará. Puede ser un ID de modelo de Hugging Face, una ruta de acceso absoluta local o una ruta de bucket de Cloud Storage. El valor predeterminado es el mismo que el de --model.
  • --enable_jit: Indica si se debe habilitar el modo JIT. El valor predeterminado es True.
  • --enable_lora: Indica si se debe habilitar el modo de carga LoRA. El valor predeterminado es False.
  • --max_lora_rank: La clasificación máxima de LoRA compatible con los adaptadores de LoRA, que se define en las solicitudes. El valor predeterminado es 16.
  • --data_parallel_size: La cantidad de réplicas paralelas de datos. El valor predeterminado es 1.
  • --tensor_parallel_size: La cantidad de réplicas paralelas de tensores. El valor predeterminado es 1.
  • --max_running_seqs: La cantidad máxima de solicitudes que el servidor puede procesar de forma simultánea. Cuanto más grande sea este argumento, mayor será la capacidad de procesamiento que podría alcanzar el servidor, pero con posibles efectos adversos en la latencia. El valor predeterminado es 256.
  • --hbm_utilization_factor: El porcentaje de HSM gratuito de Cloud TPU que se puede asignar para la caché de KV después de que se cargan las ponderaciones del modelo. Establecer este argumento en un valor más bajo puede evitar que Cloud TPU HBM se quede sin memoria. El valor predeterminado es 0.9.
  • --seed: Es el valor inicial para inicializar todos los generadores de números aleatorios. Cambiar este argumento puede afectar el resultado generado para el mismo prompt. El valor predeterminado es 0.

Solicita la cuota de Cloud TPU

En Model Garden, tu cuota predeterminada es de cuatro chips de Cloud TPU v5e en la región us-west1. Estas cuotas se aplican a las implementaciones con un clic y a las implementaciones de notebooks de Colab Enterprise. Para solicitar cuotas adicionales, consulta Solicita una cuota más alta.