Notes on Running Llama 3.1 405b

a11ce.com/llama-notes.html

Rough notes on Running Llama 3.1 405b.

Intro

You need an HF token to download the weights from hugging face. Go to the model page, request access, then create an HF token (the permissions don't matter, it's just to identify you).

For reference, the plausible configurations are, vaguely in price order:

Also storing the weights is ~500gb, but in practice you want ~600gb for inference instances (os, drivers, various caches. instance storage is cheap so overkill is fine)

Host options

RunPod has 4x H200 SXM at $18.44/hour. I'm using the pytorch template, you can maybe use the vllm template, not sure what it does. Also I had to use 'SSH over exposed TCP'.

See the end of this page for other neoclouds that I have not yet tried or do not work.

Setup/download

These steps are written for 4x H200 on RunPod, I'll update them as I try new hosts if anything changes. Ask an LLM about vllm parameters for other GPUs. The max-model-len here is suboptimal, haven't looked into it yet.

Download the weights:

python3 -m pip install vllm fastsafetensors "huggingface_hub[cli]"
echo "export HF_TOKEN=hf_xyz" >> ~/.bashrc
echo "export HF_XET_HIGH_PERFORMANCE=1" >> ~/.bashrc
echo "export LLAMA=meta-llama/Llama-3.1-405B-FP8" >> ~/.bashrc
source ~/.bashrc
hf auth login --token "$HF_TOKEN"
hf download "$LLAMA" --exclude "original/*"

This takes about 30 min. RunPod has storage for $70/month so it's break-even if you start 7 sessions/month (assuming runpod storage connects instantly, idk)

Serving

vllm serve $LLAMA --tensor-parallel-size 4 --load-format fastsafetensors --max-model-len 65536 --allow-deprecated-quantization

This takes about 5 min.

The last flag works as of vllm 0.28. Claude says the fix is to use RedHatAI/Meta-Llama-3.1-405B-FP8, haven't tried it. Also I see a bunch of OOM errors but I guess it auto-retries so whatever.

I see roughly 20-30 tok/s per branch (and more branches are ~free which is pretty cool).

To forward local port 8000 to remote port 8000:

ssh -N -L 8000:localhost:8000 root@"<IP> -p <port>

If you want to give other people access, instead start vllm with --api-key SOME_KEY and use cloudflare quick tunnel:

curl -L -o /usr/local/bin/cloudflared https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64
chmod +x /usr/local/bin/cloudflared
cloudflared tunnel --url http://localhost:8000

Then use the random url plus your key as Authorization: Bearer.

Using R2 object storage

Maybe it's faster to use runai to load weights directly to the GPU from S3-like storage? Cloudflare R2 is pretty cheap and has free egress.

(my readonly R2 weights bucket)

access: 41f96e313f23edb542aa2d9f11f27d1f
secret: 1864e92a4c9ac79581489dec9bf2e50b037fc47e35eb581ce46389cacfbdecb4
url: https://47e39f6f60165d6392620a903e13d8b6.r2.cloudflarestorage.com

Upload

Create an R2 bucket on cloudflare and a read+write API key.

After you download the weights for the first time:

curl https://rclone.org/install.sh | bash
rclone config create r2 s3 provider=Cloudflare \
  access_key_id=<ACCESS_KEY> secret_access_key=<SECRET> \
  endpoint=https://<account_id>.r2.cloudflarestorage.com acl=private no_check_bucket=true

Then upload:

rclone copy -L --transfers 16 --s3-upload-concurrency 8 --s3-chunk-size 64M --progress /workspace/.cache/huggingface/hub/models--meta-llama--Llama-3.1-405B-FP8/snapshots/<SOME_LONG_HEX_PATH>/ r2:llama/Llama-3.1-405B-FP8/

This takes about 10 minutes, and will cost about $7/month to store.

Download

Set AWS configs for runai:

export AWS_ACCESS_KEY_ID=$(rclone config show r2 | awk '/access_key_id/{print $3}')
export AWS_SECRET_ACCESS_KEY=$(rclone config show r2 | awk '/secret_access_key/{print $3}')
export AWS_ENDPOINT_URL=$(rclone config show r2 | awk '/endpoint/{print $3}')
export RUNAI_STREAMER_S3_ENDPOINT=$AWS_ENDPOINT_URL
export AWS_REGION=auto AWS_DEFAULT_REGION=auto
export RUNAI_STREAMER_CONCURRENCY=${RUNAI_STREAMER_CONCURRENCY:-32}

Download and serve (note the served-model-name param, otherwise it will be the S3 path):

vllm serve s3://llama/Llama-3.1-405B-FP8/ --load-format runai_streamer --tensor-parallel-size 4 --max-model-len 65536 --allow-deprecated-quantization --served-model-name meta-llama/Llama-3.1-405B-FP8 --api-key KEY

This takes about 10 minutes vs the previous way's 35 (saves ~$8, so it's break-even at one use per month. but again just use mine).

But, the weights never touch disk so it takes 10 minutes every time you start vllm. Tbh this is probably fine, but maybe not if you want to experiment with vllm params?

Downloading to tmpfs

Alternatively, you can pull the weights to an in-RAM tmpfs.

Check the size (needs 500G, runpod has this but others may not):

df -h /dev/shm

Then download:

mkdir -p /dev/shm/llama
rclone copy --transfers 16 --multi-thread-streams 16 --multi-thread-cutoff 64M --progress r2:llama/Llama-3.1-405B-FP8/ /dev/shm/llama/

Takes 3 minutes.

Then serve:

vllm serve /dev/shm/llama --tensor-parallel-size 4 --load-format fastsafetensors --max-model-len 65536 --allow-deprecated-quantization --served-model-name meta-llama/Llama-3.1-405B-FP8 --api-key KEY

That took 3 minutes?? So i guess there's no reason to use runai ok lol.

So it's about 6 minutes ($1.8) for setup, then $18.44/hr. pretty good!

TODO reorganize the steps here and write a one-shot script.

Other neoclouds to try

If you try any of these, please lmk the outcome.

Various neoclouds with issues

No capacity, maybe these are worth checking periodically?:

No clear B2C availability:

'Serverless' options that were still initializing after 2 hours: