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-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:

  • Publicación en varios hosts
  • Publicación de anuncios desglosada [experimental]
  • Almacenamiento en caché de prefijos
  • Compatibilidad con la cuantizació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 de 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 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 para implementar el contenedor de Hex-LLM con una topología de TPU de varios hosts, consulta el notebook de Vertex AI Model Garden - Llama 3.1 (Deployment).

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 de anuncios desglosada [experimental]

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

El servicio desagregado 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, así como el rendimiento general del servicio. 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 estrictos de latencia.

Para habilitar esta función, establece --disagg_topo en los argumentos del contenedor Hex-LLM. A continuación, se muestra un ejemplo de cómo implementar el contenedor Hex-LLM en la 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 relleno previo 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, 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 conservar una caché temporal de los cálculos de los 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. A continuación, se muestra un ejemplo de 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 las instrucciones que superan una cierta longitud (512 tokens de forma predeterminada, que se pueden configurar 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é siempre sea un múltiplo de prefill_len_padding. El campo cached_tokens de usage.prompt_tokens_details en la respuesta de la API de Chat Completion indica cuántos de los 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
  }
}

Completado previo en fragmentos

El llenado previo en fragmentos divide el llenado previo de una solicitud en fragmentos más pequeños y combina el llenado previo y la decodificación en un solo paso por lotes. Hex-LLM implementa el prellenado en fragmentos 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. A continuación, se muestra un ejemplo de 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 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,
)

Los notebooks de ejemplo de Colab Enterprise incluyen los siguientes:

Configura argumentos del servidor y 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 requisitos previstos. Ten en cuenta que los argumentos están predefinidos para la implementación con un solo clic, lo que permite una experiencia de implementación más sencilla. Para personalizar los argumentos, puedes basarte en 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. Se admiten los artefactos de modelos cuantificados BitsAndBytes int8 y AWQ para 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 usará de forma predeterminada el valor de --model.
  • --tokenizer_mode: Es el modo del tokenizador. Las opciones posibles son ["auto", "slow"]. El valor predeterminado es "auto". Si se configura como "auto", se usa el tokenizador rápido si está disponible. Los tokenizadores lentos están escritos en Python y se proporcionan en la biblioteca de Transformers, mientras que los tokenizadores rápidos que ofrecen una mejora en el rendimiento están escritos 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. Establecer este parámetro 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) para la que se puede usar 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 mayor requiere más memoria de TPU.
  • --sliding_window: Si se configura, este argumento anula el tamaño de la ventana del modelo para la atención de ventana deslizante. Si se establece este argumento en 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 está destinado solo al uso experimental. 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 prompt, ya que cambia los tokens que se muestrean como los 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 prefill 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. Si se configura en N en lugar de 1, se mejora el rendimiento en aproximadamente N y 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 de paralelismo de tensores mejora la latencia, ya que acelera la multiplicación de matrices al reducir 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 de 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 se establece --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 prellenado por iteración. El valor predeterminado es 1. Cuanto mayor 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 relleno previo hasta un múltiplo de este valor. El valor predeterminado es 8. Aumentar este valor reduce los tiempos de recompilación del modelo, pero incrementa el desperdicio de procesamiento y la sobrecarga de la inferencia. El parámetro de configuración óptimo 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 desperdicio de procesamiento y la sobrecarga de inferencia. El parámetro de configuración óptimo 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 mayor 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 recompilación del modelo, pero aumenta el desperdicio de procesamiento y la sobrecarga de la inferencia. El parámetro de configuración óptimo depende del tráfico de solicitudes.
  • --decode_blocks_padding: Durante la decodificación, el servidor agrega padding a la cantidad de bloques de memoria que se usan para la caché de clave-valor (caché de KV) de una secuencia hasta que sea un múltiplo de este valor. El valor predeterminado es 128. Aumentar este valor reduce los tiempos de recompilación del modelo, pero aumenta el desperdicio de procesamiento y la sobrecarga de inferencia. El parámetro de configuración óptimo 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 la HBM. El valor predeterminado es False. Configurar este argumento puede mejorar el rendimiento, ya que se reutilizan los cálculos de los prefijos compartidos de solicitudes anteriores.
  • --enable_chunked_prefill: Indica si se debe habilitar el relleno previo en fragmentos. El valor predeterminado es False. Configurar este argumento puede admitir una longitud de contexto más larga y mejorar el rendimiento.

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 con un valor más alto, se incrementará el tamaño de la caché de KV y se puede mejorar el rendimiento, pero también aumentará el riesgo de quedarse sin HBM de TPU de Cloud 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, los perfiles del servidor usan la HBM y calculan la cantidad de bloques de dispositivos que se asignarán según --hbm_utilization_factor. Si configuras este argumento con un valor más alto, se incrementa el tamaño de la caché de KV y se puede mejorar la capacidad de procesamiento, pero también aumenta el riesgo de quedarse sin HBM de TPU de Cloud 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 se establece este argumento en un valor más grande, se reduce la sobrecarga en la administración de bloques, pero se desperdicia más memoria. El impacto exacto en el rendimiento debe determinarse de forma empírica.

Dynamic LoRA

  • --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 Llama.
  • --max_lora_rank: Es la clasificación máxima de LoRA compatible con los adaptadores de LoRA definidos 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 dinámicos de LoRA. El valor predeterminado es True. Si configuras --no-enable_lora_cache, se inhabilita. El almacenamiento en caché mejora el rendimiento porque elimina la necesidad de volver a descargar los archivos del adaptador de LoRA que se usaron anteriormente.
  • --max_num_mem_cached_lora: Es la cantidad máxima de adaptadores de LoRA almacenados en la caché de memoria de TPU TPU.El valor predeterminado es 16. Establecer este argumento en un valor más grande 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.

Ajusta los argumentos del servidor

Los argumentos del servidor están interrelacionados y tienen un efecto colectivo en el rendimiento de la publicación. Por ejemplo, un valor más grande de --max_model_len=4096 genera un mayor uso de la memoria de la TPU y, por lo tanto, requiere una mayor asignación de memoria y menos procesamiento por lotes. Además, algunos argumentos se determinan según el caso de uso, mientras que otros se pueden ajustar. A continuación, se muestra 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 la que cada chip v5e ofrece 16 GB: tpu_memory / 16. Para un modelo de 8B y una precisión de bfloat16, necesitas más de 1 chip. Entre las configuraciones de 1, 4 y 8 chips, la más pequeña que ofrece más de 1 chip es la de 4 chips: ct5lp-hightpu-4t. Luego, puedes establecer --tensor_parallel_size=4.
  4. Determina la longitud máxima del contexto (longitud de entrada + longitud de salida) para el caso de uso previsto. Por ejemplo, 4096. Luego, puedes establecer --max_model_len=4096.
  5. Ajusta la cantidad de memoria de TPU libre asignada para la caché de KV al valor máximo que se puede alcanzar con 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 muestra 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.