Allocating GPU Resources Efficiently

GPU memory is the bottleneck I run into more often than any other on ML workloads. Whether you're fine-tuning a 7B parameter model or serving a 70B model in production, understanding how VRAM is consumed — and how to squeeze the most out of it — is the difference between a workflow that runs and one that crashes with CUDA out of memory ten minutes in.

VRAM requirements

The memory footprint depends on parameter count and the numerical precision you store them in:

Memory (bytes) = Parameters × Bytes per Parameter

Precision	Bytes per Param	7B Model	13B Model	70B Model
FP32	4	28 GB	52 GB	280 GB
FP16/BF16	2	14 GB	26 GB	140 GB
INT8	1	7 GB	13 GB	70 GB
INT4	0.5	3.5 GB	6.5 GB	35 GB

Weight memory by precision and model size. Each precision step roughly halves VRAM — and 70B at INT4 (35 GB) finally fits on a single 40 GB consumer-class GPU.

That table only accounts for weights. During training you also need memory for several other things. Optimizer states — Adam keeps two extra copies of every parameter (momentum and variance), roughly 2x the model size in FP32. Gradients — one full copy of all parameters. Activations — intermediate values stored for the backward pass, which scale with batch size and sequence length. And the KV cache during inference, which grows with sequence length and batch size.

A rough rule of thumb for training with Adam in mixed precision: total VRAM ends up around 4x the model weight size. A 7B model in BF16 (14 GB of weights) needs about 56-60 GB to train at a reasonable batch size.

Batch size vs memory

Batch size has a roughly linear relationship with activation memory. Double the batch, double the activations. So you have a tension: bigger batches give better throughput and more stable gradients, but you may not have the VRAM.

The practical move is to find the largest batch size that fits, then use gradient accumulation to hit a larger effective batch size.

# Instead of batch_size=32 (which may OOM):
# Use batch_size=8 with gradient_accumulation_steps=4
# Effective batch size = 8 × 4 = 32, peak memory = batch_size of 8

from transformers import TrainingArguments

training_args = TrainingArguments(
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    # effective batch size: 8 * 4 = 32
    fp16=True,
    output_dir="./output",
)

Gradient accumulation runs multiple forward-backward passes at a smaller batch size, accumulating gradients before doing a single optimizer step. The math comes out the same as a large batch, but peak memory matches the smaller per-step batch.

Mixed precision training

Mixed precision uses FP16 or BF16 for the forward and backward passes while keeping a master copy of weights in FP32 for the optimizer update. Roughly half the memory and better throughput on modern GPUs with tensor cores.

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch in dataloader:
    optimizer.zero_grad()

    with autocast(dtype=torch.bfloat16):
        outputs = model(batch["input_ids"])
        loss = loss_fn(outputs, batch["labels"])

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

On FP16 vs BF16: BF16 has become the default for training because it keeps the same exponent range as FP32 (8 bits), which avoids the overflow and underflow problems that make FP16 painful. The loss-scaling gymnastics you need with FP16 just go away with BF16. If your GPU supports it (A100, H100, or newer), use BF16.

Multi-GPU strategies

When a single GPU isn't enough — either the model doesn't fit or training is too slow — you have three primary parallelism strategies.

Data parallelism

The simplest. Replicate the entire model on each GPU and split the data. Each GPU runs a different mini-batch, computes gradients, and synchronizes via all-reduce.

# PyTorch DistributedDataParallel
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")
model = DDP(model.to(local_rank), device_ids=[local_rank])

Use it when the model fits on a single GPU and you want more throughput. This is the default and your first choice.

The limitation is that every GPU has to hold a full copy of the model. For a 70B model in BF16 (140 GB), data parallelism on 80 GB A100s is not an option.

Model parallelism (tensor parallelism)

Split individual layers across GPUs. A large matmul Y = XW can be partitioned column-wise across GPUs, each one computing a slice of Y, then combining the results.

Use it when the model is too large for a single GPU and you want to minimize latency. Tensor parallelism keeps pipeline depth at 1 (no pipeline bubbles) but it needs high-bandwidth interconnects like NVLink, because every layer involves communication.

Pipeline parallelism

Assign different layers to different GPUs. GPU 0 runs layers 0-15, GPU 1 runs layers 16-31, and so on. Micro-batches flow through the pipeline.

Use it when the model is too large for a single GPU and you have many GPUs. Pipeline parallelism has lower communication overhead than tensor parallelism (only activations move between stages) but you pay for it with pipeline bubbles — idle time when a GPU is waiting on the previous stage.

Combining strategies

Production training of large models usually combines all three:

Typical 70B model training on 64 GPUs:
├── 8 nodes × 8 GPUs each
├── Tensor Parallelism: 8-way within each node (NVLink)
├── Pipeline Parallelism: 2-way across node pairs
└── Data Parallelism: 4-way across remaining groups

DeepSpeed, Megatron-LM, and FSDP (Fully Sharded Data Parallel) handle the complexity of combining these. FSDP has been the one I reach for most often lately because it shards optimizer states, gradients, and parameters across GPUs at the same time — basically data parallelism plus memory savings from sharded model state.

Practical CUDA memory management

Beyond the high-level strategies, a handful of techniques help day to day.

Activation checkpointing

Instead of storing every intermediate activation for the backward pass, recompute them on the fly during backprop. About 30% more compute time for 60-70% less activation memory.

from torch.utils.checkpoint import checkpoint

# Instead of:
# output = self.expensive_layer(input)

# Use:
output = checkpoint(self.expensive_layer, input, use_reentrant=False)

Most training frameworks expose this as a single flag:

# HuggingFace Transformers
training_args = TrainingArguments(
    gradient_checkpointing=True,
    # ...
)

Monitoring VRAM

You can't diagnose memory issues without visibility. nvidia-smi is fine for a quick look, but for fine-grained analysis:

import torch

# Current allocation
print(f"Allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
print(f"Reserved:  {torch.cuda.memory_reserved() / 1e9:.2f} GB")

# Enable memory snapshot for debugging
torch.cuda.memory._record_memory_history()
# ... run your code ...
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")
# Visualize at: pytorch.org/memory_viz

Clearing cache

PyTorch's CUDA allocator caches freed memory for reuse. Usually that's good, but it can sometimes block memory from being available for different-sized allocations:

torch.cuda.empty_cache()  # Releases cached memory back to CUDA

empty_cache() does not free tensors that are still referenced. If you're leaking memory, it's almost always Python references keeping tensors alive, not the CUDA allocator.

Setting memory limits

In shared GPU environments you can cap a process's VRAM usage:

# Limit to 40 GB on an 80 GB GPU
torch.cuda.set_per_process_memory_fraction(0.5, device=0)

A decision flow

When you hit a GPU memory wall, work through this progression in order:

Reduce precision. Switch to BF16/FP16 if you haven't (2x memory).
Enable gradient checkpointing. Trade compute for memory.
Reduce batch size and increase gradient accumulation. Effective batch stays the same.
Use FSDP or DeepSpeed ZeRO to shard model state across GPUs.
Apply quantization for inference (INT8 or INT4).
Use pipeline or tensor parallelism to split the model itself.
Rent more or bigger GPUs. Sometimes the right answer is more hardware.

Each step adds complexity, so start with the simplest one that solves the problem. A lot of real ML projects can be handled with a single GPU, mixed precision, and gradient checkpointing. You don't always need a multi-node cluster.