An interactive tool to estimate the VRAM requirements for training and inference with Paligemma models.
This section provides a visual and tabular breakdown of the VRAM required for your selected configuration. The bar chart offers a quick comparison of the memory consumed by each component, while the table details the precise calculations. Interact with the controls on the left to see how each parameter affects the total memory footprint.
| Component | Memory (GB) | Calculation |
|---|---|---|
| Total Estimated Memory | (Subtotal + 15% Overhead) |
Numerical precision determines the number of bytes used to store each parameter in the model. Lower precision reduces memory usage and can speed up computation, but may come with a trade-off in accuracy.
| Data Type | Bytes per Parameter | Typical Use Case | Key Considerations |
|---|---|---|---|
| FP32 | 4 bytes | Baseline calculations | Highest memory cost, but most numerically stable. |
| BF16 | 2 bytes | Standard for training LLMs | Best balance of range and precision for stable training. |
| FP16 | 2 bytes | Training & Inference | Can suffer from instability (underflow/overflow). |
| FP8 | 1 byte | Accelerated Training | Requires modern hardware (e.g., H100). |
| INT8 | 1 byte | Inference (Quantization) | 4x memory reduction vs FP32. Small potential accuracy loss. |
| INT4 | 0.5 bytes | Inference (Quantization) | 8x memory reduction. Higher risk of accuracy degradation. |
Memory = Weights + Gradients + Optimizer + Activations
(P * B_w) + (P_t * B_g) + (P_t * B_o) + (L*B*S*H_d*K*B_a)
Memory = Weights + KV Cache
(P * B_w) + (B*S*L*H_d*2*B_kv)
P: Total ParametersP_t: Trainable ParametersB_w, B_g, B_o, B_a, B_kv: Bytes per Weight, Gradient, Optimizer, Activation, KV CacheL: Number of LayersB: Batch SizeS: Sequence LengthH_d: Hidden Dimension SizeK: Activation Constant (K_act)This is the memory needed to store the model's parameters—the learned numbers that encode its knowledge. It represents a static, foundational memory cost. Once the model is loaded onto the GPU, this memory is occupied and does not change. The calculation is straightforward: the total number of parameters multiplied by the bytes required for the chosen precision. For example, a 3 billion parameter model at BF16 precision (2 bytes/param) requires 3B × 2 = 6 GB of VRAM just to be loaded.
Gradients are a training-only memory component. During the backward pass of training, the model calculates the derivative of the loss function with respect to every single trainable parameter. These derivatives, or gradients, indicate how to adjust each weight to reduce the model's error. A buffer must be allocated to store these gradients. The size of this buffer is directly proportional to the number of trainable parameters. For full fine-tuning, this means the gradient memory is equal to the model weight memory. For parameter-efficient fine-tuning (PEFT) like LoRA, only gradients for the small set of trainable adapter weights are stored, leading to massive memory savings.
This is another training-only component, and often the largest. Modern optimizers are "stateful," meaning they track historical information about gradients to make smarter weight updates. This history, or "state," consumes VRAM for every trainable parameter. The memory cost varies drastically by optimizer:
The standard for transformers, but memory-heavy. It stores two states for each parameter: a 1st moment (momentum) and a 2nd moment (variance). Both are typically kept in full FP32 (4 bytes each) for stability, resulting in 4 + 4 = 8 bytes per parameter. This is why the optimizer states can consume 2x more memory than the model weights and gradients combined.
A highly efficient version that quantizes the Adam states. By storing momentum and variance in an 8-bit format (1 byte each), it achieves a 4x memory reduction (1 + 1 = 2 bytes) compared to standard Adam, often with minimal impact on performance.
These optimizers are leaner. SGD with Momentum stores only one FP32 state (the momentum), costing 4 bytes/param. Adafactor approximates Adam's second moment, so its memory cost is also dominated by the single FP32 momentum state it keeps, making it comparable to SGD.
The simplest form of SGD is "stateless" and tracks no historical gradient information. It consumes no extra memory but may lead to slower or less stable training compared to its stateful counterparts.
This is the primary dynamic memory component, as its size depends on the input data (batch size and sequence length), not just the model size.
Sequence Length (denoted as S) represents the number of tokens the model processes at once. It's also known as the "context window." This is a critical parameter that directly controls the dynamic memory footprint of the model.
S) is often limited by the VRAM available to store this cache.As shown in the dynamic memory formulas, S is a direct multiplier. Therefore, adjusting the sequence length is one of the most effective ways to manage memory when dealing with "Out of Memory" errors, especially in long-context scenarios.
This script demonstrates a more realistic training setup for Paligemma using PyTorch, `transformers`, and the `datasets` library. The comments highlight key parameters you can change to manage performance and memory, directly relating to the concepts explained in this calculator.
import torch
from datasets import load_dataset
from transformers import (
PaliGemmaForConditionalGeneration,
AutoProcessor,
TrainingArguments,
Trainer
)
from PIL import Image
import requests
# 1. --- MODEL, PRECISION, AND PERFORMANCE ---
# Parameter: Model Size
# Change this string to load different model sizes. E.g., "google/paligemma-3b-pt-224".
# This is the largest factor in STATIC memory usage (Weights, Gradients, Optimizer States).
model_id = "google/paligemma-3b-pt-224"
# Parameter: Precision
# `torch_dtype` controls the precision of model weights.
# torch.bfloat16 (BF16) is recommended for training on modern GPUs (Ampere+).
model_precision = torch.bfloat16
# Performance Factor: Flash Attention 2
# `attn_implementation="flash_attention_2"` drastically reduces activation memory.
# It is crucial for training with long sequences or large batches.
use_flash_attention = True
processor = AutoProcessor.from_pretrained(model_id)
model = PaliGemmaForConditionalGeneration.from_pretrained(
model_id,
torch_dtype=model_precision,
attn_implementation="flash_attention_2" if use_flash_attention else "eager",
device_map="auto"
)
# 2. --- DATASET AND PREPROCESSING ---
# Load a small sample of a real dataset (e.g., GQA for visual question answering)
# For a real run, you would use the full dataset: `load_dataset("gqa", "all")`
ds = load_dataset('graphcore/gqa-tiny', split='train')
# We need to format the data into a prompt-response structure for the model.
# The prompt for PaliGemma is typically a task prefix like "answer" or "caption".
def preprocess_data(examples):
# Parameter: Sequence Length (implicitly controlled by max_length in processor)
# This is a major factor in DYNAMIC memory usage.
max_length = 128
# Download and open images. In a real use case, images would be local.
try:
image_url = examples['image_url']
image = Image.open(requests.get(image_url, stream=True).raw).convert("RGB")
except Exception:
# Handle cases where image fails to load
image = Image.new('RGB', (224, 224))
prompt = "answer " + examples['question']
inputs = processor(
text=prompt,
images=image,
return_tensors="pt",
padding='max_length',
max_length=max_length,
truncation=True
)
# The 'labels' are the tokenized answers.
labels = processor(
text=examples['answer'],
return_tensors="pt",
padding='max_length',
max_length=max_length,
truncation=True
).input_ids
inputs['labels'] = labels
return inputs
processed_ds = ds.map(preprocess_data, remove_columns=ds.column_names)
# 3. --- TRAINING CONFIGURATION ---
training_args = TrainingArguments(
output_dir="./paligemma-finetuned-gqa",
# Parameter: Per-Device Batch Size (B)
# Number of samples on one GPU. Directly impacts dynamic memory.
per_device_train_batch_size=8,
# Performance Factor: Gradient Accumulation
# Effective batch size = per_device_train_batch_size * gradient_accumulation_steps.
# Allows large effective batch sizes by trading compute for memory.
gradient_accumulation_steps=4,
# Parameter: Optimizer
# 'adamw_torch' is the default (8 bytes/param). For memory savings, consider
# `optim="paged_adamw_8bit"` which requires `bitsandbytes`.
optim="adamw_torch",
# Parameter: Precision (for training loop)
# Use BF16 mixed-precision training. Requires Ampere+ GPU. Set to `fp16=True` for older GPUs.
bf16=True,
# Performance Factor: Gradient Checkpointing
# Trades compute for memory by not storing all activations.
# Dramatically reduces activation memory, allowing larger batches or longer sequences.
gradient_checkpointing=True,
# Performance Factor: Dataloader Workers
# `dataloader_num_workers` uses multiple CPU processes to prepare data batches.
# Prevents the GPU from waiting for data, improving throughput. Uses CPU RAM.
dataloader_num_workers=4,
# Other standard arguments
num_train_epochs=1,
logging_steps=10,
save_strategy="epoch",
learning_rate=2e-5,
report_to="none", # Can be "wandb", "tensorboard"
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=processed_ds,
)
# Start the training process
trainer.train()