Diffuser des modèles ouverts à l'aide d'un conteneur premium Hex-LLM sur Cloud TPU

Hex-LLM, une solution de diffusion de grands modèles de langage (LLM) à haute efficacité avec XLA, est le framework de diffusion LLM Vertex AI conçu et optimisé pour le matériel Cloud TPU. Hex-LLM combine des technologies de diffusion LLM telles que le traitement par lot continu et l'attention paginée avec des optimisations Vertex AI adaptées à XLA et Cloud TPU. Il s'agit d'une solution de diffusion LLM à haute efficacité et à faible coût sur Cloud TPU pour les modèles Open Source.

Hex-LLM est disponible dans Model Garden via le Playground de modèles, le déploiement en un clic et le notebook.

Fonctionnalités

Hex-LLM est basé sur des projets Open Source avec les optimisations de Google pour XLA et Cloud TPU. Hex-LLM atteint un débit élevé et une faible latence lors de la diffusion de LLM fréquemment utilisés.

Hex-LLM inclut les optimisations suivantes :

  • Algorithme de traitement par lot continu basé sur des jetons pour s'assurer que les modèles utilisent pleinement le matériel avec un grand nombre de requêtes simultanées.
  • Réécriture complète des noyaux d'attention optimisés pour XLA.
  • Stratégies flexibles et composables de parallélisme des données et de parallélisme des tensors avec des méthodes de segmentation des pondérations hautement optimisées pour exécuter efficacement les LLM sur plusieurs puces Cloud TPU.

Hex-LLM est compatible avec un large éventail de LLM denses et creux :

  • Gemma 2B et 7B
  • Gemma-2 9B et 27B
  • Llama-2 7B, 13B et 70B
  • Llama-3 8B et 70B
  • Llama-3.1 8B et 70B
  • Llama-3.2 1B et 3B
  • Llama-3.3 70B
  • Llama-Guard-3 1B et 8B
  • Llama-4 Scout-17B-16E
  • Mistral 7B
  • Mixtral 8x7B et 8x22B
  • Phi-3 mini et medium
  • Phi-4, Phi-4 reasoning et Phi-4 reasoning plus
  • Qwen-2 0.5B, 1.5B et 7B
  • Qwen-2.5 0.5B, 1.5B, 7B, 14B et 32B

Hex-LLM fournit également diverses fonctionnalités, telles que les suivantes :

  • Hex-LLM est inclus dans un seul conteneur. Hex-LLM regroupe le serveur d'API, le moteur d'inférence et les modèles compatibles dans une seule image Docker à déployer.
  • Compatible avec le format de modèles Hugging Face. Hex-LLM peut charger un modèle Hugging Face à partir d'un disque local, du hub Hugging Face et d'un bucket Cloud Storage.
  • Quantification avec bitsandbytes et AWQ.
  • Chargement LoRA dynamique. Hex-LLM est capable de charger les pondérations LoRA en lisant l'argument de requête lors de la diffusion.

Fonctionnalités avancées

Hex-LLM est compatible avec les fonctionnalités avancées suivantes :

  • Diffusion multi-hôtes
  • Diffusion désagrégée [expérimental]
  • Mise en cache des préfixes
  • Compatibilité avec la quantification 4 bits

Diffusion multi-hôtes

Hex-LLM permet désormais de diffuser des modèles avec une tranche TPU multi-hôtes. Cette fonctionnalité vous permet de diffuser de grands modèles qui ne peuvent pas être chargés dans une VM TPU à hôte unique, qui contient au maximum huit cœurs v5e.

Pour activer cette fonctionnalité, définissez --num_hosts dans les arguments du conteneur Hex-LLM et définissez --tpu_topology dans la requête d'importation du modèle du SDK Vertex AI. L'exemple suivant montre comment déployer le conteneur Hex-LLM avec une topologie TPU 4x4 v5e qui diffuse le modèle 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,
)

Pour obtenir un tutoriel de bout en bout sur le déploiement du conteneur Hex-LLM avec une topologie TPU multihôte, consultez le notebook Vertex AI Model Garden – Llama 3.1 (Deployment).

En général, les seules modifications nécessaires pour activer la diffusion multi-hôte sont les suivantes :

  1. Définissez l'argument --tensor_parallel_size sur le nombre total de cœurs dans la topologie TPU.
  2. Définissez l'argument --num_hosts sur le nombre d'hôtes dans la topologie TPU.
  3. Définissez --tpu_topology avec l'API d'importation de modèles du SDK Vertex AI.

Diffusion désagrégée [expérimental]

Hex-LLM est désormais compatible avec la diffusion désagrégée en tant que fonctionnalité expérimentale. Il ne peut être activé que sur la configuration à hôte unique et les performances sont en cours d'optimisation.

La diffusion désagrégée est une méthode efficace pour équilibrer le délai avant le premier jeton (TTFT) et le temps par jeton de sortie (TPOT) pour chaque requête, ainsi que le débit de diffusion global. Il sépare la phase de préremplissage et la phase de décodage en différentes charges de travail afin qu'elles n'interfèrent pas les unes avec les autres. Cette méthode est particulièrement utile pour les scénarios qui définissent des exigences de latence strictes.

Pour activer cette fonctionnalité, définissez --disagg_topo dans les arguments du conteneur Hex-LLM. Voici un exemple montrant comment déployer le conteneur Hex-LLM sur TPU v5e-8 qui diffuse le modèle 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,
)

L'argument --disagg_topo accepte une chaîne au format "number_of_prefill_workers,number_of_decode_workers". Dans l'exemple précédent, il est défini sur "3,1" pour configurer trois nœuds de calcul de préremplissage et un nœud de calcul de décodage. Chaque nœud de calcul utilise deux cœurs TPU v5e.

Mise en cache des préfixes

La mise en cache des préfixes réduit le délai avant le premier jeton (TTFT) pour les requêtes dont le contenu est identique au début de la requête, comme les préambules à l'échelle de l'entreprise, les instructions système courantes et l'historique des conversations multitours. Au lieu de traiter les mêmes jetons d'entrée à plusieurs reprises, Hex-LLM peut conserver un cache temporaire des calculs de jetons d'entrée traités pour améliorer le TTFT.

Pour activer cette fonctionnalité, définissez --enable_prefix_cache_hbm dans les arguments du conteneur Hex-LLM. Voici un exemple montrant comment déployer le conteneur Hex-LLM sur TPU v5e-8 qui diffuse le modèle 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 utilise la mise en cache des préfixes pour optimiser les performances des requêtes dépassant une certaine longueur (512 jetons par défaut, configurable à l'aide de prefill_len_padding). Les accès au cache se produisent par incréments de cette valeur, ce qui garantit que le nombre de jetons mis en cache est toujours un multiple de prefill_len_padding. Le champ cached_tokens de usage.prompt_tokens_details dans la réponse de l'API Chat Completion indique le nombre de jetons d'invite qui ont été succès de cache (hit).

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

Préremplissage par blocs

Le préremplissage par blocs divise le préremplissage d'une requête en blocs plus petits et combine le préremplissage et le décodage en une seule étape de traitement par lot. Hex-LLM implémente le préremplissage par blocs pour équilibrer le délai avant le premier jeton (TTFT) et le délai par jeton de sortie (TPOT), et améliorer le débit.

Pour activer cette fonctionnalité, définissez --enable_chunked_prefill dans les arguments du conteneur Hex-LLM. Voici un exemple montrant comment déployer le conteneur Hex-LLM sur TPU v5e-8 qui diffuse le modèle 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,
)

Compatibilité avec la quantification 4 bits

La quantification est une technique permettant de réduire les coûts de calcul et de mémoire liés à l'exécution de l'inférence en représentant les poids ou les activations avec des types de données de faible précision tels que INT8 ou INT4 au lieu des types BF16 ou FP32 habituels.

Hex-LLM est compatible avec la quantification INT8 des poids uniquement. L'assistance étendue inclut les modèles avec des poids INT4 quantifiés à l'aide de la quantification AWQ à point zéro. Hex-LLM est compatible avec les variantes INT4 des familles de modèles Mistral, Mixtral et Llama.

Aucun flag supplémentaire n'est requis pour diffuser des modèles quantifiés.

Premiers pas dans Model Garden

Le conteneur de diffusion Cloud TPU Hex-LLM est intégré à Model Garden. Vous pouvez accéder à cette technologie de diffusion via Playground, le déploiement en un clic et les exemples de notebooks Colab Enterprise pour divers modèles.

Utiliser Playground

Model Garden Playground est un point de terminaison Vertex AI prédéployé accessible en envoyant des requêtes dans la fiche de modèle.

  1. Saisissez une requête et, éventuellement, incluez des arguments pour votre requête.

  2. Cliquez sur ENVOYER pour obtenir rapidement la réponse du modèle.

Faites l'essai avec Gemma !

Utiliser le déploiement en un clic

Vous pouvez déployer un point de terminaison Vertex AI personnalisé avec Hex-LLM à l'aide d'une fiche de modèle.

  1. Accédez à la page Fiche de modèle, puis cliquez sur Déployer.

  2. Pour la variante de modèle que vous souhaitez utiliser, sélectionnez le type de machine Cloud TPU v5e pour le déploiement.

  3. Cliquez sur Déployer en bas pour lancer le processus de déploiement. Vous recevez deux notifications par e-mail : une lorsque le modèle est importé et une autre lorsque le point de terminaison est prêt.

Utiliser le notebook Colab Enterprise

Pour plus de flexibilité et de personnalisation, vous pouvez utiliser les exemples de notebooks Colab Enterprise pour déployer un point de terminaison Vertex AI avec Hex-LLM à l'aide du SDK Vertex AI pour Python.

  1. Accédez à la page de la fiche de modèle, puis cliquez sur Ouvrir le notebook.

  2. Sélectionnez le notebook Vertex Serving. Le notebook s'ouvre dans Colab Enterprise.

  3. Parcourez le notebook pour déployer un modèle à l'aide de Hex-LLM et envoyez des requêtes de prédiction au point de terminaison. L'extrait de code pour le déploiement se présente comme suit :

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

Voici quelques exemples de notebooks Colab Enterprise :

Configurer les arguments du serveur et les variables d'environnement

Vous pouvez définir les arguments suivants pour lancer le serveur Hex-LLM. Vous pouvez adapter les arguments pour qu'ils correspondent le mieux à votre cas d'utilisation et à vos besoins. Notez que les arguments sont prédéfinis pour un déploiement en un clic afin de faciliter au maximum le déploiement. Pour personnaliser les arguments, vous pouvez vous appuyer sur les exemples de notebooks pour référence et définir les arguments en conséquence.

Modèle

  • --model : modèle à charger. Vous pouvez spécifier un ID de modèle Hugging Face, un chemin d'accès à un bucket Cloud Storage (gs://my-bucket/my-model) ou un chemin local. Les artefacts du modèle doivent suivre le format Hugging Face et utiliser des fichiers safetensors pour les pondérations du modèle. Les artefacts de modèles quantifiés BitsAndBytes int8 et AWQ sont compatibles avec Llama, Gemma 2 et Mistral/Mixtral.
  • --tokenizer : tokenizer à charger. Il peut s'agir d'un ID de modèle Hugging Face, d'un chemin d'accès à un bucket Cloud Storage (gs://my-bucket/my-model) ou d'un chemin d'accès local. Si cet argument n'est pas défini, la valeur par défaut est celle de --model.
  • --tokenizer_mode : mode du tokenizer. Les choix possibles sont ["auto", "slow"]. La valeur par défaut est "auto". Si cette valeur est définie sur "auto", le tokenizer rapide est utilisé s'il est disponible. Les tokenizers lents sont écrits en Python et fournis dans la bibliothèque Transformers, tandis que les tokenizers rapides offrant une amélioration des performances sont écrits en Rust et fournis dans la bibliothèque Tokenizers. Pour en savoir plus, consultez la documentation Hugging Face.
  • --trust_remote_code : indique s'il faut autoriser les fichiers de code à distance définis dans les dépôts de modèles Hugging Face. La valeur par défaut est False.
  • --load_format : format des points de contrôle du modèle à charger. Les choix possibles sont ["auto", "dummy"]. La valeur par défaut est "auto". Si cette valeur est définie sur "auto", les pondérations du modèle sont chargées au format safetensors. Si cette valeur est définie sur "dummy", les pondérations du modèle sont initialisées de manière aléatoire. Définir cette valeur sur "dummy" est utile pour les tests.
  • --max_model_len : longueur maximale du contexte (longueur de l'entrée plus la longueur de la sortie) à utiliser pour le modèle. La valeur par défaut est lue à partir du fichier de configuration du modèle au format Hugging Face : config.json. Une longueur de contexte maximale plus importante nécessite plus de mémoire TPU.
  • --sliding_window : si cette valeur est définie, cet argument remplace la taille de la fenêtre du modèle pour l'attention à fenêtre glissante. Si vous définissez cet argument sur une valeur plus élevée, le mécanisme d'attention inclut plus de jetons et se rapproche de l'effet de l'auto-attention standard. Cet argument est destiné à un usage expérimental uniquement. Dans les cas d'utilisation généraux, nous vous recommandons d'utiliser la taille de fenêtre d'origine du modèle.
  • --seed : valeur source pour l'initialisation de tous les générateurs de nombres aléatoires. La modification de cet argument peut affecter la sortie générée pour la même requête en modifiant les jetons qui sont échantillonnés en tant que jetons suivants. La valeur par défaut est 0.

Moteur d'inférence

  • --num_hosts : nombre d'hôtes à exécuter. La valeur par défaut est 1. Pour en savoir plus, consultez la documentation sur la configuration TPU v5e.
  • --disagg_topo : définit le nombre de nœuds de calcul de préremplissage et de décodage avec la fonctionnalité expérimentale de diffusion dissociée. La valeur par défaut est None. L'argument suit le format suivant : "number_of_prefill_workers,number_of_decode_workers".
  • --data_parallel_size : nombre de réplications parallèles de données. La valeur par défaut est 1. Si vous définissez cette valeur sur N au lieu de 1, le débit est approximativement amélioré de N, tout en conservant la même latence.
  • --tensor_parallel_size : nombre de réplications parallèles de Tensors. La valeur par défaut est 1. L'augmentation du nombre de répliques parallèles de tenseurs améliore généralement la latence, car elle accélère la multiplication matricielle en réduisant la taille de la matrice.
  • --worker_distributed_method : méthode distribuée pour lancer le nœud de calcul. Utilisez mp pour le module multitraitement ou ray pour la bibliothèque Ray. La valeur par défaut est mp.
  • --enable_jit : indique si le mode JIT (Just-in-Time Compilation) doit être activé. La valeur par défaut est True. Si la propriété --no-enable_jit est définie sur , la détection automatique est désactivée. L'activation du mode JIT améliore les performances d'inférence, mais nécessite un temps supplémentaire pour la compilation initiale. En général, les avantages en termes de performances d'inférence l'emportent sur les frais généraux.
  • --warmup : indique s'il faut réchauffer le serveur avec des exemples de requêtes lors de l'initialisation. La valeur par défaut est True. Si vous définissez --no-warmup, vous le désactivez. L'échauffement est recommandé, car les requêtes initiales déclenchent une compilation plus lourde et sont donc plus lentes.
  • --max_prefill_seqs : nombre maximal de séquences pouvant être planifiées pour le préremplissage par itération. La valeur par défaut est 1. Plus cette valeur est élevée, plus le débit du serveur peut être élevé, mais avec des effets négatifs potentiels sur la latence.
  • --prefill_seqs_padding : le serveur ajuste la taille du lot de préremplissage à un multiple de cette valeur. La valeur par défaut est 8. L'augmentation de cette valeur réduit les temps de recompilation du modèle, mais augmente le gaspillage de calcul et la surcharge d'inférence. Le paramètre optimal dépend du trafic de requêtes.
  • --prefill_len_padding : le serveur ajuste la longueur de la séquence à un multiple de cette valeur. La valeur par défaut est 512. L'augmentation de cette valeur réduit les temps de recompilation du modèle, mais augmente le gaspillage de calcul et la surcharge d'inférence. Le paramètre optimal dépend de la distribution des données des requêtes.
  • --max_decode_seqs/--max_running_seqs : nombre maximal de séquences pouvant être planifiées pour le décodage par itération. La valeur par défaut est 256. Plus cette valeur est élevée, plus le débit du serveur peut être élevé, mais avec des effets négatifs potentiels sur la latence.
  • --decode_seqs_padding : le serveur ajuste la taille du lot de décodage à un multiple de cette valeur. La valeur par défaut est 8. L'augmentation de cette valeur réduit les temps de recompilation du modèle, mais augmente le gaspillage de calcul et la surcharge d'inférence. Le paramètre optimal dépend du trafic de requêtes.
  • --decode_blocks_padding : le serveur ajuste le nombre de blocs de mémoire utilisés pour le cache clé-valeur (cache KV) d'une séquence à un multiple de cette valeur lors du décodage. La valeur par défaut est 128. L'augmentation de cette valeur réduit les temps de recompilation du modèle, mais augmente le gaspillage de calcul et la surcharge d'inférence. Le paramètre optimal dépend de la distribution des données des requêtes.
  • --enable_prefix_cache_hbm : indique si la mise en cache des préfixes doit être activée dans HBM. La valeur par défaut est False. La définition de cet argument peut améliorer les performances en réutilisant les calculs des préfixes partagés des requêtes précédentes.
  • --enable_chunked_prefill : indique s'il faut activer le préremplissage par blocs. La valeur par défaut est False. Définir cet argument peut prendre en charge une longueur de contexte plus longue et améliorer les performances.

Gestion de la mémoire

  • --hbm_utilization_factor : pourcentage de mémoire à haut débit (HBM) Cloud TPU disponible pouvant être alloué pour le cache KV après le chargement des pondérations du modèle. La valeur par défaut est 0.9. Définir cet argument sur une valeur plus élevée augmente la taille du cache KV et peut améliorer le débit, mais cela augmente le risque de manquer de HBM Cloud TPU lors de l'initialisation et de l'exécution.
  • --num_blocks : nombre de blocs d'appareil à allouer au cache KV. Si cet argument est défini, le serveur ignore --hbm_utilization_factor. Si cet argument n'est pas défini, le serveur profile l'utilisation de la mémoire HBM et calcule le nombre de blocs d'appareils à allouer en fonction de --hbm_utilization_factor. Si vous définissez cet argument sur une valeur plus élevée, la taille du cache KV augmente, ce qui peut améliorer le débit. Toutefois, cela augmente le risque de manquer de HBM Cloud TPU lors de l'initialisation et de l'exécution.
  • --block_size : nombre de jetons stockés dans un bloc. Les choix possibles sont [8, 16, 32, 2048, 8192]. La valeur par défaut est 32. Définir cet argument sur une valeur plus élevée réduit les frais généraux liés à la gestion des blocs, au détriment d'une plus grande perte de mémoire. L'impact exact sur les performances doit être déterminé de manière empirique.

LoRA dynamique

  • --enable_lora : indique si le chargement dynamique des adaptateurs LoRA depuis Cloud Storage doit être activé. La valeur par défaut est False. Cette fonctionnalité est compatible avec la famille de modèles Llama.
  • --max_lora_rank : rang LoRA maximal compatible avec les adaptateurs LoRA définis dans les requêtes. La valeur par défaut est 16. Définir une valeur plus élevée pour cet argument permet une plus grande flexibilité dans les adaptateurs LoRA qui peuvent être utilisés avec le serveur, mais augmente la quantité de mémoire HBM Cloud TPU allouée aux pondérations LoRA et diminue le débit.
  • --enable_lora_cache : indique si la mise en cache des adaptateurs LoRA dynamiques doit être activée. La valeur par défaut est True. Si la propriété --no-enable_lora_cache est définie sur , la détection automatique est désactivée. La mise en cache améliore les performances, car elle élimine la nécessité de télécharger à nouveau les fichiers d'adaptateur LoRA précédemment utilisés.
  • --max_num_mem_cached_lora : nombre maximal d'adaptateurs LoRA stockés dans le cache de mémoire TPU valeur par défaut est 16. Définir une valeur plus élevée pour cet argument augmente les chances d'un succès de cache (hit), mais augmente la quantité de mémoire HBM Cloud TPU utilisée.

Vous pouvez également configurer le serveur à l'aide des variables d'environnement suivantes :

  • HEX_LLM_LOG_LEVEL : contrôle la quantité d'informations de journalisation générées. La valeur par défaut est INFO. Définissez la valeur sur l'un des niveaux de journalisation Python standards définis dans le module de journalisation.
  • HEX_LLM_VERBOSE_LOG : indique si la sortie de journalisation détaillée doit être activée. Les valeurs autorisées sont true ou false. La valeur par défaut est false.

Ajuster les arguments du serveur

Les arguments du serveur sont interdépendants et ont un effet collectif sur les performances de diffusion. Par exemple, un paramètre --max_model_len=4096 plus élevé entraîne une utilisation plus importante de la mémoire TPU et nécessite donc une allocation de mémoire plus importante et moins de regroupements. De plus, certains arguments sont déterminés par le cas d'utilisation, tandis que d'autres peuvent être ajustés. Voici un workflow pour configurer le serveur Hex-LLM.

  1. Déterminez la famille et la variante de modèle qui vous intéressent. Par exemple, Llama 3.1 8B Instruct.
  2. Estimez la limite inférieure de la mémoire TPU nécessaire en fonction de la taille et de la précision du modèle : model_size * (num_bits / 8). Pour un modèle 8B et une précision bfloat16, la limite inférieure de la mémoire TPU requise serait de 8 * (16 / 8) = 16 GB.
  3. Estimez le nombre de puces TPU v5e nécessaires, sachant que chaque puce v5e offre 16 Go : tpu_memory / 16. Pour un modèle de 8 milliards de paramètres et une précision bfloat16, vous avez besoin de plus d'un chip. Parmi les configurations à 1, 4 et 8 puces, la plus petite configuration offrant plus d'une puce est la configuration à 4 puces : ct5lp-hightpu-4t. Vous pourrez ensuite définir --tensor_parallel_size=4.
  4. Déterminez la longueur maximale du contexte (longueur de l'entrée + longueur de la sortie) pour le cas d'utilisation prévu. Par exemple, 4096. Vous pourrez ensuite définir --max_model_len=4096.
  5. Ajustez la quantité de mémoire TPU libre allouée au cache KV à la valeur maximale possible en fonction des configurations du modèle, du matériel et du serveur (--hbm_utilization_factor). Commencez par 0.95. Déployez le serveur Hex-LLM et testez-le avec de longues requêtes et une forte simultanéité. Si le serveur manque de mémoire, réduisez le facteur d'utilisation en conséquence.

Voici un exemple d'ensemble d'arguments pour déployer 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

Voici un exemple d'ensemble d'arguments pour déployer Llama 3.1 70B Instruct AWQ sur ct5lp-hightpu-4t :

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

Demander un quota Cloud TPU

Dans Model Garden, votre quota par défaut est de 32 puces Cloud TPU v5e dans la région us-west1. Ces quotas s'appliquent aux déploiements en un clic et aux déploiements de notebooks Colab Enterprise. Pour demander une augmentation de quota, consultez Demander un ajustement de quota.