Deploying Meta Llama 3.1 405B on Google Cloud Vertex AI

Meta Llama 3.1 405B can be deployed on Google Cloud Vertex AI using Hugging Face's Deep Learning Containers (DLCs) and Text Generation Inference (TGI). This allows users to run the largest open LLM from Meta in a managed environment, leveraging A3 accelerator-optimized machines for high-performance inference.

Hardware Requirements and Memory Management

Deploying Meta Llama 3.1 405B requires significant GPU VRAM, which varies based on the precision of the model weights. To load the model checkpoint, the approximate memory requirements are:

Model Size FP16 FP8 INT4
8B 16 GB 8 GB 4 GB
70B 140 GB 70 GB 35 GB
405B 810 GB 405 GB 203 GB

Because a single H100 node (8 x 80GB GPUs) provides approximately 640GB of VRAM, the 405B model must be run at a lower precision, such as FP8, or in a multi-node setup. For this deployment, the A3 High GPU machine type is recommended, providing 8 x NVIDIA H100 80GB GPUs, 208 vCPUs, and 1872 GB of memory.

Deployment Process on Vertex AI

Deploying the model involves registering the model in the Vertex AI Model Registry and then deploying it to a Vertex AI Endpoint.

Model Registration

Registration is performed using the google-cloud-aiplatform Python SDK. The process requires a Hugging Face TGI DLC that supports the Meta Llama 3.1 architecture (which uses a different RoPE scaling method than Llama 3). Key configuration environment variables include:

  • MODEL_ID: Set to meta-llama/Meta-Llama-3.1-405B-Instruct-FP8.
  • HUGGING_FACE_HUB_TOKEN: A read-access token for the gated Meta Llama repository.
  • NUM_SHARD: Set to 8 to utilize all GPUs on the A3 instance.
  • HF_XET_HIGH_PERFORMANCE: Recommended to be set to 1 to accelerate the download of the ~400 GiB model weights.

Endpoint Deployment

Once registered, the model is deployed to a Vertex AI Endpoint using the a3-highgpu-8g machine type and 8 NVIDIA H100 80GB accelerators. The full deployment process typically takes 25-30 minutes, accounting for resource allocation, weight downloading (10 minutes), and loading into TGI (2 minutes).

Running Online Predictions

Vertex AI exposes a /predict route that serves text generation via the TGI DLC. Because the /generate endpoint is used, inputs must be formatted with the appropriate chat template before being sent to the endpoint.

Formatting Inputs

Users should use the apply_chat_template method from the transformers library to ensure the conversation is formatted correctly for Meta Llama 3.1. A typical formatted prompt includes specific tokens such as <|begin_of_text|>, <|start_header_id|>, and <|eot_id|>.

Prediction Methods

Predictions can be executed through three primary methods:

  1. Same-session Python: Using the deployed_model.predict method immediately after deployment.
  2. Different-session Python: Instantiating the aiplatform.Endpoint via its resource name (projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}).
  3. Vertex AI UI: Using the Online Prediction UI by providing a JSON payload containing the formatted inputs and generation parameters (e.g., max_new_tokens, temperature, top_p).

Resource Management

To avoid unnecessary costs, resources should be cleaned up after use. This involves undeploying the model from all endpoints using deployed_model.undeploy_all(), deleting the endpoint via deployed_model.delete(), and finally removing the model from the registry using model.delete().

Sources