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.
  • 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
  • Llama 3.1 8B y 70B
  • Llama 3.2 1B y 3B
  • Llama Guard 3 1B y 8B
  • Mistral 7B
  • Mixtral 8x7B y 8x22B
  • Phi-3 mini y medio
  • Qwen2 0.5B, 1.5B y 7B
  • Qwen2.5 0.5B, 1.5B, 7B, 14B y 32B AWQ

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.

Funciones avanzadas

Hex-LLM admite las siguientes funciones avanzadas:

  • Publicación en varios hosts
  • Publicación desagregada [experimental]
  • Almacenamiento en caché de prefijos
  • Compatibilidad con la cuantificación de 4 bits

Publicación en varios hosts

Hex-LLM ahora admite la entrega de modelos con una porción de TPU de varios hosts. Esta función te permite entregar modelos grandes que no se pueden cargar en una sola VM de TPU de host, que contiene como máximo ocho núcleos v5e.

Para habilitar esta función, establece --num_hosts en los argumentos del contenedor Hex-LLM y --tpu_topology en la solicitud de carga del modelo del SDK de Vertex AI. En el siguiente ejemplo, se muestra cómo implementar el contenedor Hex-LLM con una topología TPU 4x4 v5e que entrega el modelo bfloat16 Llama 3.1 70B:

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Meta-Llama-3.1-70B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=16",
    "--num_hosts=4",
    "--hbm_utilization_factor=0.9",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    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=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    tpu_topology="4x4",
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

Para obtener un instructivo de extremo a extremo para implementar el contenedor Hex-LLM con una topología de TPU de varios hosts, consulta el notebook Llama 3.1 (implementación) de Vertex AI Model Garden.

En general, los únicos cambios necesarios para habilitar la publicación en varios hosts son los siguientes:

  1. Establece el argumento --tensor_parallel_size en la cantidad total de núcleos dentro de la topología de TPU.
  2. Establece el argumento --num_hosts en la cantidad de hosts dentro de la topología de TPU.
  3. Establece --tpu_topology con la API de carga de modelos del SDK de Vertex AI.

Publicación desagregada [experimental]

Hex-LLM ahora admite la publicación desagregada como una función experimental. Solo se puede habilitar en la configuración de un solo host, y el rendimiento está en proceso de optimización.

La publicación desagregada es un método eficaz para equilibrar el tiempo hasta el primer token (TTFT) y el tiempo por token de salida (TPOT) para cada solicitud, y la capacidad de procesamiento general de la publicación. Separa la fase de precarga y la fase de decodificación en cargas de trabajo diferentes para que no interfieran entre sí. Este método es especialmente útil para situaciones que establecen requisitos de latencia estrictos.

Para habilitar esta función, establece --disagg_topo en los argumentos del contenedor Hex-LLM. El siguiente es un ejemplo que muestra cómo implementar el contenedor Hex-LLM en una TPU v5e-8 que entrega el modelo bfloat16 de Llama 3.1 8B:

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=2",
    "--disagg_topo=3,1",
    "--hbm_utilization_factor=0.9",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    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=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

El argumento --disagg_topo acepta una cadena en el formato "number_of_prefill_workers,number_of_decode_workers". En el ejemplo anterior, se establece en "3,1" para configurar tres trabajadores de reabastecimiento y 1 trabajador de decodificación. Cada trabajador usa dos núcleos de TPU v5e.

Almacenamiento en caché de prefijos

La caché de prefijos reduce el tiempo hasta el primer token (TTFT) para las instrucciones que tienen contenido idéntico al principio, como los preámbulos para toda la empresa, las instrucciones comunes del sistema y el historial de conversaciones de varias vueltas. En lugar de procesar los mismos tokens de entrada de forma reiterada, Hex-LLM puede retener una caché temporal de los cálculos de tokens de entrada procesados para mejorar el TTFT.

Para habilitar esta función, establece --enable_prefix_cache_hbm en los argumentos del contenedor de Hex-LLM. El siguiente es un ejemplo que muestra cómo implementar el contenedor de Hex-LLM en la TPU v5e-8 que entrega el modelo bfloat16 de Llama 3.1 8B:

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=2",
    "--hbm_utilization_factor=0.9",
    "--enable_prefix_cache_hbm",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    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=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

Hex-LLM emplea el almacenamiento en caché de prefijos para optimizar el rendimiento de las instrucciones que superan una cierta longitud (512 tokens de forma predeterminada, configurable con prefill_len_padding). Los hits de caché se producen en incrementos de este valor, lo que garantiza que el recuento de tokens almacenados en caché sea siempre un múltiplo de prefill_len_padding. El campo cached_tokens de usage.prompt_tokens_details en la respuesta de la API de finalización de chat indica cuántos de los tokens de instrucciones fueron un acierto de caché.

"usage": {
  "prompt_tokens": 643,
  "total_tokens": 743,
  "completion_tokens": 100,
  "prompt_tokens_details": {
    "cached_tokens": 512
  }
}

Compatibilidad con la cuantización de 4 bits

La cuantización es una técnica para reducir los costos computacionales y de memoria de la ejecución de inferencias mediante la representación de los pesos o las activaciones con tipos de datos de baja precisión, como INT8 o INT4, en lugar de los BF16 o FP32 habituales.

Hex-LLM admite la cuantización de solo pesos INT8. La compatibilidad extendida incluye modelos con pesos INT4 cuantificados con la cuantificación de punto cero de AWQ. Hex-LLM admite variantes INT4 de las familias de modelos Mistral, Mixtral y Llama.

No se requiere ninguna marca adicional para publicar modelos cuantificados.

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

Estos son algunos ejemplos de notebooks de Colab Enterprise:

Configura los argumentos del servidor y las variables de entorno

Puedes establecer los siguientes argumentos para iniciar el servidor Hex-LLM. Puedes adaptar los argumentos para que se ajusten mejor a tu caso de uso y a tus requisitos. Ten en cuenta que los argumentos están predefinidos para la implementación con un clic para habilitar la experiencia de implementación más fácil. Para personalizar los argumentos, puedes basarte en los ejemplos de notebooks como referencia y configurar los argumentos según corresponda.

Modelo

  • --model: El modelo que se cargará. Puedes especificar un ID de modelo de Hugging Face, una ruta de bucket de Cloud Storage (gs://my-bucket/my-model) o una ruta local. Se espera que los artefactos del modelo sigan el formato de Hugging Face y usen archivos safetensors para los pesos del modelo. Los artefactos de modelos cuantificados de BitsAndBytes int8 y AWQ son compatibles con Llama, Gemma 2 y Mistral/Mixtral.
  • --tokenizer: El tokenizador que se cargará. Puede ser un ID de modelo de Hugging Face, una ruta de bucket de Cloud Storage (gs://my-bucket/my-model) o una ruta local. Si no se establece este argumento, se establece de manera predeterminada en el valor de --model.
  • --tokenizer_mode: Es el modo del tokenizador. Las opciones posibles son ["auto", "slow"]. El valor predeterminado es "auto". Si se establece como "auto", se usa el analizador rápido si está disponible. Los tokenizers lentos se escriben en Python y se proporcionan en la biblioteca de Transformers, mientras que los tokenizers rápidos que ofrecen una mejora del rendimiento se escriben en Rust y se proporcionan en la biblioteca de Tokenizers. Para obtener más información, consulta la documentación de Hugging Face.
  • --trust_remote_code: Indica si se permiten archivos de código remotos definidos en los repositorios de modelos de Hugging Face. El valor predeterminado es False.
  • --load_format: Es el formato de los puntos de control del modelo que se cargarán. Las opciones posibles son ["auto", "dummy"]. El valor predeterminado es "auto". Si se establece en "auto", las ponderaciones del modelo se cargan en formato safetensors. Si se establece en "dummy", los pesos del modelo se inicializan de forma aleatoria. Configurarlo en "dummy" es útil para la experimentación.
  • --max_model_len: Es la longitud máxima del contexto (longitud de entrada más la longitud de salida) que se entregará para el modelo. El valor predeterminado se lee del archivo de configuración del modelo en formato Hugging Face: config.json. Una longitud de contexto máxima más grande requiere más memoria de TPU.
  • --sliding_window: Si se establece, este argumento anula el tamaño de la ventana del modelo para la atención de ventana deslizante. Si configuras este argumento en un valor más alto, el mecanismo de atención incluirá más tokens y se acercará al efecto de la atención automática estándar. Este argumento solo se diseñó para uso experimental. En los casos de uso generales, te recomendamos que uses el tamaño de ventana original del modelo.
  • --seed: Es el valor inicial para inicializar todos los generadores de números aleatorios. Cambiar este argumento podría afectar el resultado generado para el mismo prompt a través de los tokens que se muestrean como tokens siguientes. El valor predeterminado es 0.

Motor de inferencia

  • --num_hosts: Es la cantidad de hosts que se ejecutarán. El valor predeterminado es 1. Para obtener más detalles, consulta la documentación sobre la configuración de TPU v5e.
  • --disagg_topo: Define la cantidad de trabajadores de precarga y decodificación con la publicación desagregada de atributos experimentales. El valor predeterminado es None. El argumento sigue el formato "number_of_prefill_workers,number_of_decode_workers".
  • --data_parallel_size: La cantidad de réplicas paralelas de datos. El valor predeterminado es 1. Si estableces este valor en N desde 1, se mejora aproximadamente la productividad en N y, al mismo tiempo, se mantiene la misma latencia.
  • --tensor_parallel_size: La cantidad de réplicas paralelas de tensores. El valor predeterminado es 1. En general, aumentar la cantidad de réplicas paralelas de tensores mejora la latencia, ya que acelera la multiplicación de matrices reduciendo su tamaño.
  • --worker_distributed_method: Es el método distribuido para iniciar el trabajador. Usa mp para el módulo de multiprocesamiento o ray para la biblioteca de Ray. El valor predeterminado es mp.
  • --enable_jit: Indica si se debe habilitar el modo JIT (compilación just-in-time). El valor predeterminado es True. Configurar --no-enable_jit la inhabilita. Habilitar el modo JIT mejora el rendimiento de la inferencia a costa de requerir tiempo adicional en la compilación inicial. En general, los beneficios de rendimiento de la inferencia superan la sobrecarga.
  • --warmup: Indica si se debe activar el servidor con solicitudes de muestra durante la inicialización. El valor predeterminado es True. Configurar --no-warmup lo inhabilita. Se recomienda el calentamiento, ya que las solicitudes iniciales activan una compilación más pesada y, por lo tanto, serán más lentas.
  • --max_prefill_seqs: Es la cantidad máxima de secuencias que se pueden programar para el reabastecimiento previo por iteración. El valor predeterminado es 1. Cuanto más grande sea este valor, mayor será la capacidad de procesamiento que puede alcanzar el servidor, pero con posibles efectos adversos en la latencia.
  • --prefill_seqs_padding: El servidor agrega padding al tamaño del lote de reabastecimiento a un múltiplo de este valor. El valor predeterminado es 8. Aumentar este valor reduce los tiempos de recomposición del modelo, pero aumenta el procesamiento desperdiciado y la sobrecarga de inferencia. La configuración óptima depende del tráfico de solicitudes.
  • --prefill_len_padding: El servidor agrega padding a la longitud de la secuencia para que sea un múltiplo de este valor. El valor predeterminado es 512. Aumentar este valor reduce los tiempos de recomposición del modelo, pero aumenta la sobrecarga de inferencia y procesamiento desperdiciado. La configuración óptima depende de la distribución de datos de las consultas.
  • --max_decode_seqs/--max_running_seqs: Es la cantidad máxima de secuencias que se pueden programar para la decodificación por iteración. El valor predeterminado es 256. Cuanto más grande sea este valor, mayor será la capacidad de procesamiento que puede alcanzar el servidor, pero con posibles efectos adversos en la latencia.
  • --decode_seqs_padding: El servidor agrega padding al tamaño del lote de decodificación para que sea un múltiplo de este valor. El valor predeterminado es 8. Aumentar este valor reduce los tiempos de recomposición del modelo, pero aumenta la sobrecarga de inferencia y procesamiento desperdiciado. La configuración óptima depende del tráfico de solicitudes.
  • --decode_blocks_padding: El servidor agrega padding a la cantidad de bloques de memoria que se usan para la caché de par clave-valor (caché de KV) de una secuencia a un múltiplo de este valor durante la decodificación. El valor predeterminado es 128. Aumentar este valor reduce los tiempos de recomposición del modelo, pero aumenta la sobrecarga de inferencia y procesamiento desperdiciado. La configuración óptima depende de la distribución de datos de las consultas.
  • --enable_prefix_cache_hbm: Indica si se debe habilitar la caché de prefijos en HBM. El valor predeterminado es False. Configurar este argumento puede mejorar el rendimiento reutilizando los cálculos de los prefijos compartidos de solicitudes anteriores.

Administración de memoria

  • --hbm_utilization_factor: Es el porcentaje de memoria de ancho de banda alto (HBM) de Cloud TPU libre que se puede asignar para la caché de KV después de que se cargan las ponderaciones del modelo. El valor predeterminado es 0.9. Si configuras este argumento en un valor más alto, se incrementa el tamaño de la caché de KV y se puede mejorar la capacidad de procesamiento, pero aumenta el riesgo de quedarse sin HBM de Cloud TPU durante la inicialización y el tiempo de ejecución.
  • --num_blocks: Es la cantidad de bloques de dispositivos que se asignarán para la caché de KV. Si se establece este argumento, el servidor ignora --hbm_utilization_factor. Si no se configura este argumento, el servidor perfila el uso de HBM y calcula la cantidad de bloques de dispositivos que se asignarán en función de --hbm_utilization_factor. Si configuras este argumento en un valor más alto, se aumenta el tamaño de la caché de KV y se puede mejorar la capacidad de procesamiento, pero aumenta el riesgo de quedarse sin HBM de Cloud TPU durante la inicialización y el tiempo de ejecución.
  • --block_size: Es la cantidad de tokens almacenados en un bloque. Las opciones posibles son [8, 16, 32, 2048, 8192]. El valor predeterminado es 32. Si estableces este argumento en un valor más alto, se reduce la sobrecarga en la administración de bloques, a costa de más desperdicio de memoria. El impacto exacto en el rendimiento se debe determinar empíricamente.

LoRA dinámico

  • --enable_lora: Indica si se debe habilitar la carga de adaptadores LoRA dinámicos desde Cloud Storage. El valor predeterminado es False. Esto es compatible con la familia de modelos Llama.
  • --max_lora_rank: Es la clasificación máxima de LoRA compatible con los adaptadores de LoRA que se define en las solicitudes. El valor predeterminado es 16. Establecer este argumento en un valor más alto permite una mayor flexibilidad en los adaptadores LoRA que se pueden usar con el servidor, pero aumenta la cantidad de HBM de Cloud TPU asignado para las ponderaciones de LoRA y disminuye la capacidad de procesamiento.
  • --enable_lora_cache: Indica si se debe habilitar el almacenamiento en caché de los adaptadores LoRA dinámicos. El valor predeterminado es True. Configurar --no-enable_lora_cache la inhabilita. El almacenamiento en caché mejora el rendimiento porque elimina la necesidad de volver a descargar los archivos del adaptador LoRA que se usaron anteriormente.
  • --max_num_mem_cached_lora: Es la cantidad máxima de adaptadores LoRA almacenados en la caché de memoria de la TPU.El valor predeterminado es 16. Si estableces este argumento en un valor más alto, mejora la probabilidad de que se produzca un acierto de caché, pero aumenta la cantidad de uso de HBM de Cloud TPU.

También puedes configurar el servidor con las siguientes variables de entorno:

  • HEX_LLM_LOG_LEVEL: Controla la cantidad de información de registro que se genera. El valor predeterminado es INFO. Establece este valor en uno de los niveles de registro estándar de Python definidos en el módulo de registro.
  • HEX_LLM_VERBOSE_LOG: Indica si se debe habilitar el resultado de registro detallado. Los valores permitidos son true o false. El valor predeterminado es false.

Argumentos del servidor de Tune

Los argumentos del servidor están relacionados entre sí y tienen un efecto colectivo en el rendimiento de la publicación. Por ejemplo, un parámetro de configuración más alto de --max_model_len=4096 genera un mayor uso de memoria de TPU y, por lo tanto, requiere una asignación de memoria más grande y menos procesamiento por lotes. Además, el caso de uso determina algunos argumentos, mientras que otros se pueden ajustar. Este es un flujo de trabajo para configurar el servidor Hex-LLM.

  1. Determina la familia y la variante del modelo que te interesan. Por ejemplo, Llama 3.1 8B Instruct.
  2. Estima el límite inferior de la memoria de TPU necesaria según el tamaño y la precisión del modelo: model_size * (num_bits / 8). Para un modelo de 8B y una precisión de bfloat16, el límite inferior de la memoria de TPU necesaria sería 8 * (16 / 8) = 16 GB.
  3. Estima la cantidad de chips TPU v5e necesarios, en los que cada chip v5e ofrece 16 GB: tpu_memory / 16. Para un modelo de 8B y precisión bfloat16, necesitas más de 1 chip. Entre las configuraciones de 1, 4 y 8 chips, la configuración más pequeña que ofrece más de 1 chip es la de 4 chips: ct5lp-hightpu-4t. Luego, puedes configurar --tensor_parallel_size=4.
  4. Determina la longitud máxima de contexto (longitud de entrada + longitud de salida) para el caso de uso previsto. Por ejemplo, 4096. Luego, puedes configurar --max_model_len=4096.
  5. Ajusta la cantidad de memoria TPU libre asignada para la caché de KV al valor máximo que se puede lograr según el modelo, el hardware y la configuración del servidor (--hbm_utilization_factor). Comienza con 0.95. Implementa el servidor Hex-LLM y pruébalo con instrucciones largas y alta simultaneidad. Si el servidor se queda sin memoria, reduce el factor de utilización según corresponda.

Un conjunto de argumentos de muestra para implementar Llama 3.1 8B Instruct es el siguiente:

python -m hex_llm.server.api_server \
    --model=meta-llama/Llama-3.1-8B-Instruct \
    --tensor_parallel_size=4 \
    --max_model_len=4096
    --hbm_utilization_factor=0.95

Un conjunto de argumentos de muestra para implementar Llama 3.1 70B Instruct AWQ en ct5lp-hightpu-4t es el siguiente:

python -m hex_llm.server.api_server \
    --model=hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4 \
    --tensor_parallel_size=4 \
    --max_model_len=4096
    --hbm_utilization_factor=0.45

Solicita la cuota de Cloud TPU

En Model Garden, tu cuota predeterminada es de 4 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.