GPU Memory Optimization for PyTorch Inference: Reduce VRAM Usage by Half

GPUs are expensive. A T4 has 16GB VRAM. If your model and overhead use 15GB, you can't increase batch size, add a KV cache, or run anything else on that GPU.

Here's how to cut PyTorch inference memory usage without changing your model architecture.

1. Use FP16 Instead of FP32

Half precision cuts memory usage by 50% with negligible accuracy loss for most models.

model = AutoModel.from_pretrained("model-name")
model = model.half()  # FP32 → FP16
model = model.to("cuda")

Memory: 14GB → 7GB for a 7B parameter model.

Verify accuracy:

with torch.no_grad():
    fp32_out = fp32_model(input_ids)
    fp16_out = fp16_model(input_ids)
    
    diff = (fp32_out - fp16_out.float()).abs().max().item()
    print(f"Max difference: {diff:.6f}")

For most transformer models, FP16 max difference is under 1e-3. Acceptable for production.

2. Use torch.no_grad()

PyTorch builds a computation graph during forward passes for backpropagation. During inference, you don't need gradients. The graph wastes memory.

# Bad: computes gradients (memory grows with each forward pass)
output = model(input_ids)

# Good: no gradient computation
with torch.no_grad():
    output = model(input_ids)

Memory savings: ~30-40% depending on model size.

3. Disable KV Cache for Short Sequences

For short inputs (under 128 tokens), the KV cache overhead is larger than the benefit. Disable it:

output = model.generate(
    input_ids,
    use_cache=False,  # disable KV cache
    max_new_tokens=50,
)

Memory savings: 100-500MB for small models, 1-2GB for large models.

4. Use torch.inference_mode() (PyTorch 1.9+)

inference_mode() is a stricter version of no_grad() that additionally disables autograd version tracking. Slightly faster and uses less memory:

with torch.inference_mode():
    output = model(input_ids)

Functionally identical to no_grad() for inference but with slightly lower overhead. Use it for production inference.

5. Offload to CPU When GPU Is Idle

If your GPU sits idle between requests, move the model to CPU and only load it to GPU when a request arrives:

import threading
import time

class LazyGPUModel:
    def __init__(self, model_name):
        self.model = AutoModel.from_pretrained(model_name)
        self.on_gpu = False
        self.last_used = time.time()
        self._idle_timeout = 300  # 5 minutes
    
    def predict(self, inputs):
        if not self.on_gpu:
            self.model = self.model.to("cuda")
            self.on_gpu = True
        self.last_used = time.time()
        with torch.inference_mode():
            return self.model(inputs)
    
    def check_idle(self):
        if self.on_gpu and time.time() - self.last_used > self._idle_timeout:
            self.model = self.model.to("cpu")
            torch.cuda.empty_cache()
            self.on_gpu = False

This frees the GPU for other processes when there's no traffic. Adds a ~500ms penalty on the first request after idle (loading model back to GPU).

6. Use Flash Attention

Flash Attention computes attention without materializing the full attention matrix. For long sequences (2K+ tokens), it cuts memory by up to 80%.

pip install flash-attn --no-build-isolation

model = AutoModel.from_pretrained(
    "model-name",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.float16,
)

Only works on Ampere+ GPUs (A100, A10, RTX 3090+). T4 doesn't support it.

7. Quantize to INT8

INT8 quantization reduces memory by 4x (32 bits → 8 bits per weight):

from transformers import BitsAndBytesConfig

quant_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0,
)

model = AutoModel.from_pretrained(
    "model-name",
    quantization_config=quant_config,
)

Memory: 14GB → 3.5GB for a 7B model.

Tradeoff: INT8 inference is slower than FP16 on most GPUs. The memory savings let you run larger models on smaller GPUs, but per-token latency increases. For a 7B model on T4: FP16 = 15 tokens/sec, INT8 = 8 tokens/sec. Worth it if you're VRAM-constrained, not if you're latency-sensitive.

8. Delete Unused Tensors

PyTorch doesn't garbage-collect GPU memory immediately. Tensors you think are freed are still allocated:

# After each request:
del outputs, hidden_states, logits
torch.cuda.empty_cache()

empty_cache() releases cached memory back to the GPU. Without it, after 100 requests you might hit OOM even though each request uses only 5GB.

Real-World Impact

For a 7B parameter model on a T4 (16GB):

Optimization	VRAM Usage	Tokens/sec
Baseline (FP32, with gradients)	OOM (14GB + overhead)	—
+ no_grad()	12.5 GB	12
+ FP16	7.2 GB	22
+ INT8 (4-bit)	3.8 GB	8
+ Flash Attention (2K seq)	3.2 GB	18
+ All combined	3.2 GB	18

From OOM to 3.2GB with 18 tokens/sec. On the same hardware.

When Not to Optimize

If your model fits comfortably in VRAM with 30%+ headroom, don't optimize. The complexity isn't worth it. Only optimize when you're within 80% of VRAM capacity or when you need to increase batch size.

Roptal detects your model's VRAM requirements during repo analysis and applies appropriate optimizations — FP16, no_grad, Flash Attention — based on the target GPU.