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

Hex-LLM, un modelo de lenguaje grande (LLM) de alta eficiencia que entrega XLA, es el framework de entrega de LLM de Gemini Enterprise Agent Platform 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 Gemini Enterprise Agent Platform 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-3.3 70B
  • Llama-Guard-3 1B y 8B
  • Llama-4 Scout-17B-16E
  • Mistral 7B
  • Mixtral 8x7B y 8x22B
  • Phi-3 mini y medium
  • Phi-4, Phi-4 reasoning y Phi-4 reasoning plus
  • Qwen-2 0.5B, 1.5B y 7B
  • Qwen-2.5 0.5B, 1.5B, 7B, 14B y 32B

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:

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

Entrega de 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 VM de TPU de host único, que contiene como máximo ocho núcleos v5e.

Para habilitar esta función, establece --num_hosts en los argumentos del contenedor de Hex-LLM y establece --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 de Hex-LLM con una topología de TPU 4x4 v5e que entrega el modelo Llama 3.1 70B bfloat16:

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 sobre cómo implementar el contenedor de Hex-LLM con una topología de TPU de varios hosts, consulta el notebook de Agent Platform Model Garden - Llama 3.1 (Deployment).

En general, los únicos cambios necesarios para habilitar la entrega de 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 de anuncios desagregada [experimental]

Hex-LLM ahora admite la entrega desagregada como una función experimental. Solo se puede habilitar en la configuración de host único y el rendimiento se está ajustando.

La entrega desagregada es un método eficaz para equilibrar el tiempo hasta el primer token (TTFT) y el tiempo por token de salida (TPOT) de cada solicitud, y la capacidad de procesamiento general de la entrega. Separa la fase de precompletado y la fase de decodificación en diferentes cargas de trabajo 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 de Hex-LLM. En el siguiente ejemplo, se muestra cómo implementar el contenedor de Hex-LLM en TPU v5e-8 que entrega el modelo Llama 3.1 8B bfloat16:

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 con el formato "number_of_prefill_workers,number_of_decode_workers". En el ejemplo anterior, se establece en "3,1" para configurar tres trabajadores de precompletado y 1 trabajador de decodificación. Cada trabajador usa dos núcleos de TPU v5e.

Almacenamiento en caché de prefijos

El almacenamiento en caché de prefijos reduce el tiempo hasta el primer token (TTFT) para las instrucciones que tienen contenido idéntico al principio de la instrucción, como preámbulos de toda la empresa, instrucciones comunes del sistema y el historial de conversaciones de varios turnos. En lugar de procesar los mismos tokens de entrada de forma repetida, 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. En el siguiente ejemplo, se muestra cómo implementar el contenedor de Hex-LLM en TPU v5e-8 que entrega el modelo Llama 3.1 8B bfloat16:

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=4",
    "--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 los mensajes que superan una longitud determinada (512 tokens de forma predeterminada, configurable con prefill_len_padding). Los aciertos 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 tokens de la instrucción fueron un acierto de caché.

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

Autocompletado fragmentado

El relleno previo fragmentado divide el relleno previo de una solicitud en fragmentos más pequeños y mezcla el relleno previo y la decodificación en un paso por lotes. Hex-LLM implementa el prellenado fragmentado para equilibrar el tiempo hasta el primer token (TTFT) y el tiempo por token de salida (TPOT), y mejora la capacidad de procesamiento.

Para habilitar esta función, establece --enable_chunked_prefill en los argumentos del contenedor de Hex-LLM. En el siguiente ejemplo, se muestra cómo implementar el contenedor de Hex-LLM en la TPU v5e-8 que entrega el modelo 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=4",
    "--hbm_utilization_factor=0.9",
    "--enable_chunked_prefill",
]

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

Compatibilidad con la cuantización de 4 bits

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

Hex-LLM admite la cuantización solo de pesos INT8. La compatibilidad extendida incluye modelos con pesos INT4 cuantizados con la cuantizació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 entregar modelos cuantizados.

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 Agent Platform implementado con anterioridad al que se puede acceder enviando 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 Agent Platform con Hex-LLM usando 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 Agent Platform con Hex-LLM con el SDK de Agent Platform 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:

Configura los argumentos del servidor y las variables de entorno

Puedes establecer los siguientes argumentos para iniciar el servidor de Hex-LLM. Puedes adaptar los argumentos para que se ajusten mejor a tu caso de uso y requisitos previstos. Ten en cuenta que los argumentos están predefinidos para la implementación con un solo clic, lo que permite que la experiencia de implementación sea lo más sencilla posible. Para personalizar los argumentos, puedes usar los ejemplos de notebooks como referencia y establecer los argumentos según corresponda.

Modelo

  • --model: El modelo que se cargará. Puedes especificar un ID de modelo de Hugging Face, una ruta de acceso de bucket de Cloud Storage (gs://my-bucket/my-model) o una ruta de acceso 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 int8 de BitsAndBytes y AWQ son compatibles con Llama, Gemma 2 y Mistral/Mixtral.
  • --tokenizer: Es el tokenizador que se cargará. Puede ser un ID de modelo de Hugging Face, una ruta de acceso al bucket de Cloud Storage (gs://my-bucket/my-model) o una ruta de acceso local. Si no se establece este argumento, se establece de forma 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 en "auto", se usa el asignador de token rápido si está disponible. Los tokenizadores lentos se escriben en Python y se proporcionan en la biblioteca de Transformers, mientras que los tokenizadores 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", los pesos del modelo se cargan en formato safetensors. Si se establece en "dummy", los pesos del modelo se inicializan de forma aleatoria. Establecer este valor en "dummy" es útil para experimentar.
  • --max_model_len: Es la longitud máxima de 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 de 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 con un valor mayor, el mecanismo de atención incluirá más tokens y se acercará al efecto de la autoatención estándar. Este argumento solo se usa con fines experimentales. En los casos de uso generales, recomendamos usar 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 puede afectar el resultado generado para el mismo mensaje, ya que cambia 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 precompletado y de decodificación con la función experimental de entrega desagregada. 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. Establecer este valor en N desde 1 mejora el rendimiento en aproximadamente N, mientras se mantiene la misma latencia.
  • --tensor_parallel_size: La cantidad de réplicas paralelas de tensores. El valor predeterminado es 1. Aumentar la cantidad de réplicas paralelas de tensores suele mejorar la latencia, ya que acelera la multiplicación de matrices reduciendo el tamaño de la matriz.
  • --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 justo a tiempo). El valor predeterminado es True. Si configuras --no-enable_jit, se inhabilita. Habilitar el modo JIT mejora el rendimiento de la inferencia, pero requiere tiempo adicional para la compilación inicial. En general, los beneficios de rendimiento de la inferencia superan la sobrecarga.
  • --warmup: Indica si se debe preparar el servidor con solicitudes de muestra durante la inicialización. El valor predeterminado es True. Si configuras --no-warmup, se 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 llenado previo por iteración. El valor predeterminado es 1. Cuanto más grande sea este valor, mayor será la capacidad de procesamiento que podría alcanzar el servidor, pero con posibles efectos adversos en la latencia.
  • --prefill_seqs_padding: El servidor agrega padding al tamaño del lote de precompletado para que sea un múltiplo de este valor. El valor predeterminado es 8. Aumentar este valor reduce los tiempos de recompilació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 recompilación del modelo, pero aumenta el procesamiento desperdiciado y la sobrecarga de inferencia. La configuración óptima depende de la distribución de datos de las solicitudes.
  • --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 podría 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 recompilación del modelo, pero aumenta el procesamiento desperdiciado y la sobrecarga de inferencia. 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 recompilación del modelo, pero aumenta el procesamiento desperdiciado y la sobrecarga de inferencia. La configuración óptima depende de la distribución de datos de las solicitudes.
  • --enable_prefix_cache_hbm: Indica si se debe habilitar el almacenamiento en caché de prefijos en HBM. El valor predeterminado es False. Establecer este argumento puede mejorar el rendimiento, ya que se reutilizan los cálculos de los prefijos compartidos de las solicitudes anteriores.
  • --enable_chunked_prefill: Indica si se debe habilitar el relleno previo fragmentado. El valor predeterminado es False. Establecer este argumento puede admitir una longitud de contexto más larga y mejorar el rendimiento.

Administración de la memoria

  • --hbm_utilization_factor: Es el porcentaje de memoria de gran ancho de banda (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. Establecer este argumento en un valor más alto aumenta el tamaño de la caché de KV y puede mejorar la capacidad de procesamiento, pero aumenta el riesgo de quedarse sin HBM de Cloud TPU durante la inicialización y en el entorno de ejecución.
  • --num_blocks: Es la cantidad de bloques de dispositivos que se asignarán a la caché de KV. Si se establece este argumento, el servidor ignora --hbm_utilization_factor. Si no se establece este argumento, el servidor genera perfiles del uso de HBM y calcula la cantidad de bloques de dispositivos que se asignarán en función de --hbm_utilization_factor. Configurar este argumento con un valor más alto aumenta el tamaño de la caché de KV y puede mejorar la capacidad de procesamiento, pero aumenta el riesgo de quedarse sin HBM de Cloud TPU durante la inicialización y en el entorno 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. Establecer este argumento en un valor más alto reduce la sobrecarga en la administración de bloques, a costa de un mayor desperdicio de memoria. El impacto exacto en el rendimiento se debe determinar de forma empírica.

LoRA dinámico

  • --enable_lora: Indica si se debe habilitar la carga dinámica de adaptadores de LoRA desde Cloud Storage. El valor predeterminado es False. Esto es compatible con la familia de modelos de 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 de LoRA que se pueden usar con el servidor, pero aumenta la cantidad de HBM de Cloud TPU asignada para los pesos 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. Si configuras --no-enable_lora_cache, se inhabilita. El almacenamiento en caché mejora el rendimiento, ya que elimina la necesidad de volver a descargar los archivos del adaptador de LoRA que se usaron antes.
  • --max_num_mem_cached_lora: Es la cantidad máxima de adaptadores de LoRA almacenados en la caché de la memoria de TPU TPU.El valor predeterminado es 16. Establecer este argumento en un valor más alto mejora la probabilidad de 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 de ajuste del servidor

Los argumentos del servidor están interrelacionados y tienen un efecto colectivo en el rendimiento de la entrega. Por ejemplo, un parámetro de configuración más grande 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, algunos argumentos se determinan según el caso de uso, mientras que otros se pueden ajustar. Este es un flujo de trabajo para configurar el servidor de Hex-LLM.

  1. Determina la familia y la variante del modelo de interés. 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 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 una precisión bfloat16, necesitas más de 1 chip. Entre los parámetros de configuración de 1, 4 y 8 chips, el más pequeño que ofrece más de 1 chip es el de 4 chips: ct5lp-hightpu-4t. Luego, puedes configurar --tensor_parallel_size=4.
  4. Determina la longitud de contexto máxima (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 asignación de memoria de TPU libre para la caché de KV al valor máximo que se puede alcanzar según las configuraciones del modelo, el hardware y el servidor (--hbm_utilization_factor). Comienza con 0.95. Implementa el servidor de 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.

A continuación, se muestra un conjunto de argumentos de ejemplo para implementar Llama 3.1 8B Instruct:

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 32 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 un valor de cuota más alto, consulta Solicita un ajuste de cuota.