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'PagedAttention 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
  • Mistral 7B et Mixtral 8x7B

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.

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

Vous pouvez modifier les arguments de lancement du serveur Hex-LLM suivants pour la diffusion personnalisée :

  • --model : modèle à charger. Vous pouvez spécifier un ID de modèle Hugging Face, un chemin absolu local ou un chemin d'accès à un bucket Cloud Storage.
  • --tokenizer : tokenizer à charger. Il peut s'agir d'un ID de modèle Hugging Face, d'un chemin absolu local ou d'un chemin d'accès à un bucket Cloud Storage. La valeur par défaut est identique à celle de --model.
  • --enable_jit : indique si le mode JIT doit être activé. La valeur par défaut est True.
  • --enable_lora : indique si le mode de chargement LoRA doit être activé. La valeur par défaut est False.
  • --max_lora_rank : rang LoRA maximal compatible avec les adaptateurs LoRA, défini dans les requêtes. La valeur par défaut est 16.
  • --data_parallel_size : nombre de réplications parallèles de données. La valeur par défaut est 1.
  • --tensor_parallel_size : nombre de réplications parallèles de Tensors. La valeur par défaut est 1.
  • --max_running_seqs : nombre maximal de requêtes que le serveur peut traiter simultanément. Plus cet argument est grand, plus le débit du serveur peut être élevé, mais avec des effets négatifs potentiels sur la latence. La valeur par défaut est 256.
  • --hbm_utilization_factor : pourcentage de mémoire HBM Cloud TPU sans frais pouvant être alloué pour le cache KV après le chargement des pondérations du modèle. La définition d'une valeur inférieure pour cet argument peut efficacement éviter la saturation de la mémoire HBM Cloud TPU. La valeur par défaut est 0.9.
  • --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. La valeur par défaut est 0.

Demander un quota Cloud TPU

Dans Model Garden, votre quota par défaut est de quatre 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 des quotas supplémentaires, consultez la section Demander une augmentation de quota.