Loading models bigger than your GPU

Setup#

pip install transformers torch accelerate

Tested with transformers 5.6 and accelerate 1.13. The device_map machinery lives in accelerate. Model: distilgpt2, ~356 MB download.

Automatic placement, and reading the plan#

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "distilgpt2",
    device_map="auto",          # let accelerate place the layers
    dtype=torch.float16,        # half precision: half the memory
)
print("device:", model.device)
print("footprint:", round(model.get_memory_footprint() / 1e6), "MB")

device: cuda:0
footprint: 164 MB

This output is from a machine with one NVIDIA GPU. On a laptop without one you will see device: cpu — same code, different plan. Small model, roomy hardware: everything landed in one place.

Forcing a split, to watch the machinery#

Real splits happen when the model outgrows the GPU. You can simulate that by lying about how much GPU memory exists:

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "distilgpt2",
    device_map="auto",
    dtype=torch.float16,
    max_memory={0: "100MB", "cpu": "10GB"},   # pretend the GPU is tiny
)
for module, place in model.hf_device_map.items():
    print(f"{module:20} -> {place}")

transformer.wte      -> 0
lm_head              -> 0
transformer.wpe      -> 0
transformer.drop     -> 0
transformer.ln_f     -> cpu
transformer.h        -> cpu

A 164 MB model against a fake 100 MB budget: embeddings went to GPU 0, the transformer blocks spilled to CPU RAM. A warning also appears — "Some parameters are on the meta device because they were offloaded to the cpu" — which is the split announcing itself. This snippet needs a CUDA GPU; without one there is nothing to split against.

The walkthrough#

dtype=torch.float16 halved the bill before placement started. 82M parameters at 4 bytes would be ~328 MB; at 2 bytes, 164 MB. Half precision is the first move for inference, before any offloading. In transformers 5 this argument is named dtype — you will still see the older name torch_dtype in tutorials written against v4.

hf_device_map exists only when it has something to say. When everything fits on one device, transformers 5 skips the attribute — check model.device instead. After a real split, the map lists every module's home.

Send inputs to where the model starts. inputs.to(model.device) handles both worlds; hard-coding "cuda" breaks the code on CPU machines.

Disk is the third tier. Add offload_folder="offload" and layers beyond RAM spill to disk. Generation then crawls — weights stream from disk on every single token. It is a way to try a model, not to serve one. For actually shrinking the bill, see loading in 4-bit.

Common mistakes#

Calling .to("cuda") after device_map. The placement plan already decided homes; yanking everything to one device either crashes on memory or defeats the offload. Pick one mechanism, not both.

Expecting offloaded speed. Layers on CPU compute at CPU speed, and travel costs are paid per forward pass. If generation feels 50x slower with a split, that is arithmetic, not a bug.

Forgetting dtype and loading fp32. Double the memory for no inference benefit — the difference between fitting and spilling for many models.

Using device_map for training. The mechanism targets inference. Training needs gradients and optimizer states placed too — a different problem, handled by Accelerate configs, DeepSpeed and friends.

Try it yourself#

Tighten the fake budget to "50MB" and print the map — which extra modules moved out? Then time generate() for 20 tokens in the fits-on-GPU and forced-split versions, and put a number on what offloading costs you on your machine.

What to learn next#

• Loading a model in 4-bit — shrinking the model instead of relocating it.
• Accelerate — the library doing the placement, met directly.
• Moving tensors between CPU and GPU — the mechanics one level down.