Tutorial on how to deploy the Llama 3.1 405B model on GKE Autopilot with 8 x A100 80GB GPUs using KubeAI.
We're using fp8 (8 bits) precision for this model. This allows us to reduce GPU memory required and allows us to serve the model on a single machine.
Create a GKE Autopilot cluster
gcloud container clusters create-auto cluster-1 \
--location=us-central1
Add the helm repo for KubeAI:
helm repo add kubeai https://www.kubeai.org
helm repo update
Create a values file for KubeAI with required settings:
cat <<EOF > kubeai-values.yaml
resourceProfiles:
nvidia-gpu-a100-80gb:
imageName: "nvidia-gpu"
limits:
nvidia.com/gpu: "1"
requests:
nvidia.com/gpu: "1"
# Each A100 80GB GPU gets 10 CPU and 12Gi memory
cpu: 10
memory: 12Gi
tolerations:
- key: "nvidia.com/gpu"
operator: "Equal"
value: "present"
effect: "NoSchedule"
nodeSelector:
cloud.google.com/gke-accelerator: "nvidia-a100-80gb"
cloud.google.com/gke-spot: "true"
EOF
Install KubeAI with Helm:
helm upgrade --install kubeai kubeai/kubeai \
-f ./kubeai-values.yaml \
--wait
Deploy Llama 3.1 405B by creating a KubeAI Model object:
kubectl apply -f - <<EOF
apiVersion: kubeai.org/v1
kind: Model
metadata:
name: llama-3.1-405b-instruct-fp8-a100
spec:
features: [TextGeneration]
owner:
url: hf://neuralmagic/Meta-Llama-3.1-405B-Instruct-FP8
engine: VLLM
env:
VLLM_ATTENTION_BACKEND: FLASHINFER
args:
- --max-model-len=65536
- --max-num-batched-token=65536
- --gpu-memory-utilization=0.98
- --tensor-parallel-size=8
- --enable-prefix-caching
- --disable-log-requests
- --max-num-seqs=128
- --kv-cache-dtype=fp8
- --enforce-eager
- --enable-chunked-prefill=false
- --num-scheduler-steps=8
targetRequests: 128
minReplicas: 1
maxReplicas: 1
resourceProfile: nvidia-gpu-a100-80gb:8
EOF
The pod takes about 15 minutes to startup. Wait for the model pod to be ready:
kubectl get pods -w
Once the pod is ready, the model is ready to serve requests.
Setup a port-forward to the KubeAI service on localhost port 8000:
kubectl port-forward service/kubeai 8000:80
Send a request to the model to test:
curl -v http://localhost:8000/openai/v1/completions \
-H "Content-Type: application/json" \
-d '{"model": "llama-3.1-405b-instruct-fp8-a100", "prompt": "Who was the first president of the United States?", "max_tokens": 40}'
Now let's run a benchmarking using the vLLM benchmarking script:
git clone https://github.com/vllm-project/vllm.git
cd vllm/benchmarks
wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
python3 benchmark_serving.py --backend openai \
--base-url http://localhost:8000/openai \
--dataset-name=sharegpt --dataset-path=ShareGPT_V3_unfiltered_cleaned_split.json \
--model llama-3.1-405b-instruct-fp8-a100 \
--seed 12345 --tokenizer neuralmagic/Meta-Llama-3.1-405B-Instruct-FP8
This was the output of the benchmarking script on 8 x A100 80GB GPUs:
============ Serving Benchmark Result ============
Successful requests: 1000
Benchmark duration (s): 410.49
Total input tokens: 232428
Total generated tokens: 173391
Request throughput (req/s): 2.44
Output token throughput (tok/s): 422.40
Total Token throughput (tok/s): 988.63
---------------Time to First Token----------------
Mean TTFT (ms): 136607.47
Median TTFT (ms): 125998.27
P99 TTFT (ms): 335309.25
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 302.24
Median TPOT (ms): 267.34
P99 TPOT (ms): 1427.52
---------------Inter-token Latency----------------
Mean ITL (ms): 249.94
Median ITL (ms): 128.63
P99 ITL (ms): 1240.35
==================================================
Hope this is helpful to other folks struggling to get Llama 3.1 405B up and running on GKE. Similar steps would work for GKE standard as long as you create your a2-ultragpu-8g nodepools in advance.
Top comments (0)