Hugging Face "CUDA out of memory" loading a model in CI
The model weights plus overhead exceed the GPU memory on the runner, so torch raises CUDA out of memory during from_pretrained. In CI the fix is usually to load on CPU for a smoke test, pick a smaller checkpoint, or load quantized. (See the CUDA OOM reference for the general case.)
What this error means
from_pretrained fails with "torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate ... GiB" while placing weights on the GPU.
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0;
14.76 GiB total capacity; 13.10 GiB already allocated; 1.20 GiB free; ...)Common causes
The checkpoint is larger than the runner GPU
A multi-billion-parameter model in full precision needs more VRAM than the CI GPU provides.
Full-precision load without offload
Loading in fp32 without device_map or quantization doubles the memory versus fp16/8-bit.
How to fix it
Load on CPU for a CI smoke test
For a correctness check, load on CPU so no GPU memory is needed.
from transformers import AutoModel
AutoModel.from_pretrained("org/model", device_map="cpu")Load in lower precision or quantized
Use fp16 or 8-bit loading (needs accelerate and bitsandbytes) to fit the model on a smaller GPU.
import torch
from transformers import AutoModelForCausalLM
AutoModelForCausalLM.from_pretrained("org/model", torch_dtype=torch.float16, device_map="auto")How to prevent it
- Use small checkpoints for CI correctness tests.
- Load on CPU or in fp16/8-bit where the GPU is small.
- Reserve full-size GPU loads for dedicated large runners.