# How to Solve Memory Bottlenecks in Training Foundation Models: 6 Proven Strategies

> Solve memory bottlenecks in foundation model training with 6 expert strategies. Learn PEFT, gradient checkpointing, quantization, and more to optimize your training.

- Repository: [Chip Huyen/aie-book](https://github.com/chiphuyen/aie-book)
- Tags: best-practices
- Published: 2026-04-24

---

**Memory bottlenecks in training foundation models are caused by the simultaneous storage of full model weights, optimizer states, and intermediate activations, which can be mitigated through Parameter-Efficient Fine-Tuning (PEFT), gradient checkpointing, quantized training, and activation offloading.**

Training modern foundation models pushes GPU/TPU memory to its absolute limits. According to Chip Huyen’s *AI Engineering* book (repository `chiphuyen/aie-book`), memory constraints often determine whether a model can be fine-tuned on available hardware or remains inaccessible. Understanding the architectural sources of these bottlenecks—and the specific techniques to alleviate them—is essential for practitioners working with large language models and multimodal systems.

## Root Causes of Memory Scarce Training

Training foundation models requires holding multiple tensor types simultaneously in accelerator memory. As noted in [[`chapter-summaries.md`](https://github.com/chiphuyen/aie-book/blob/main/chapter-summaries.md)](https://github.com/chiphuyen/aie-book/blob/main/chapter-summaries.md#L158), "the more parameters to update during finetuning, the more memory finetuning needs."

### Optimizer State Overhead

Optimizers like **Adam** maintain per-parameter momentum tensors—typically two additional floating-point values for every weight. This effectively **triples the memory footprint** beyond the raw model weights, as the optimizer must store first and second moment estimates alongside the parameters themselves.

### Activation Memory and Backpropagation

Each transformer layer caches **intermediate activations** during the forward pass to compute gradients during backpropagation. This requirement scales with model depth, hidden dimension size, and batch size, creating a linear growth in memory consumption that often exceeds the weight storage itself.

### KV-Cache and Long Contexts

While primarily an inference concern, the **KV-cache** illustrates how transformer architectures consume memory proportionally to sequence length. As stated in [[`resources.md`](https://github.com/chiphuyen/aie-book/blob/main/resources.md)](https://github.com/chiphuyen/aie-book/blob/main/resources.md#L338), "KV cache optimization … is one of the most memory‑heavy parts of transformer inference," a principle that extends to training when handling long-context windows.

## Architectural Strategies to Reduce Memory Usage

The book outlines several architectural approaches designed to lower peak GPU memory usage while preserving model performance.

### Parameter-Efficient Fine-Tuning (PEFT) and LoRA

**PEFT** techniques freeze the base model weights and train only a small set of additional parameters. **LoRA** (Low-Rank Adaptation) injects trainable rank-decomposition matrices into attention layers, reducing trainable parameters by orders of magnitude. This cuts both weight-update memory and optimizer state, as the book notes regarding PEFT’s ability to reduce finetuning memory requirements (see [[`chapter-summaries.md`](https://github.com/chiphuyen/aie-book/blob/main/chapter-summaries.md)](https://github.com/chiphuyen/aie-book/blob/main/chapter-summaries.md#L160)).

### Gradient Checkpointing

Instead of storing all activations during the forward pass, **gradient checkpointing** trades compute for memory by recomputing activations during the backward pass. This technique selectively retains activations for specific layers while recalculating others, significantly reducing peak memory consumption at the cost of approximately 20-30% additional computation time.

### Quantized Training with 8-bit and 4-bit Precision

**Quantized training** reduces the bit-width of weights and optimizer states from 32-bit floats to 8-bit or 4-bit integers. As mentioned in [[`chapter-summaries.md`](https://github.com/chiphuyen/aie-book/blob/main/chapter-summaries.md)](https://github.com/chiphuyen/aie-book/blob/main/chapter-summaries.md#L160), this directly mitigates memory bottlenecks "by reducing the number of bits needed to represent each value," though it requires careful handling of activation and gradient precision to maintain training stability.

### Activation Offloading and ZeRO-Offload

**Activation offloading** moves intermediate tensors to CPU memory or NVMe storage when not immediately needed for computation. Frameworks like DeepSpeed implement **ZeRO-Offload** (Zero Redundancy Optimizer), which partitions optimizer states and gradients across CPU and GPU memory hierarchies, effectively expanding the available memory pool beyond VRAM limits.

### Model Parallelism and Distributed Training

**Model parallelism** distributes layers or tensor slices across multiple GPUs, treating the aggregate memory of several devices as a single pool. This approach is essential when individual model parameters exceed the capacity of a single accelerator, though it introduces communication overhead between devices.

## Practical Implementation: Code Examples

Below are concise, runnable snippets illustrating three primary memory-saving techniques discussed in the book.

### Implementing LoRA with Hugging Face PEFT

This example demonstrates how to freeze a base model and train only low-rank adapters:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model

model_name = "meta-llama/Meta-Llama-3-8B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")

# LoRA config: inject low‑rank adapters into attention layers

lora_cfg = LoraConfig(
    r=8,               # rank

    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],  # typical attention query/value

    lora_dropout=0.1,
    bias="none",
)

model = get_peft_model(model, lora_cfg)

# Train only the adapter parameters

model.print_trainable_parameters()

```

Only the adapter matrices (approximately `r × hidden_dim`) are stored and updated, while the base model remains frozen.

### Enabling Gradient Checkpointing in Transformers

Activation recomputation requires a single method call:

```python
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B")
model.gradient_checkpointing_enable()   # activates checkpointing

```

Activations are recomputed during the backward pass rather than stored, cutting peak activation memory significantly.

### 8-bit Quantized Training with BitsAndBytes

Quantization reduces storage for both weights and optimizer states:

```python
from transformers import BitsAndBytesConfig, AutoModelForCausalLM, Trainer, TrainingArguments

bnb_cfg = BitsAndBytesConfig(
    load_in_8bit=True,   # store weights in 8‑bit

    llm_int8_threshold=6.0,
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    quantization_config=bnb_cfg,
    device_map="auto"
)

training_args = TrainingArguments(
    output_dir="./lora_finetuned",
    per_device_train_batch_size=2,
    fp16=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=your_dataset,
)

trainer.train()

```

This configuration shrinks both model weights and optimizer buffers, directly addressing memory constraints identified in the book’s discussion of quantized training.

## Summary

- **Memory bottlenecks** in foundation model training arise from optimizer states (2-3x model size), activation storage, and full-model parameter updates.
- **Parameter-Efficient Fine-Tuning** via LoRA reduces trainable parameters by 99%+ while maintaining performance.
- **Gradient checkpointing** trades 20-30% additional computation for significant memory savings by recomputing activations.
- **Quantized training** with 8-bit or 4-bit precision reduces memory footprint for weights and optimizer states.
- **Activation offloading** and **model parallelism** extend available memory beyond single-GPU limitations.

## Frequently Asked Questions

### What consumes the most memory when training large language models?

The **optimizer state** typically consumes the most memory, as optimizers like Adam store two momentum tensors per parameter, effectively tripling the memory requirement beyond the model weights themselves. Additionally, activation storage for backpropagation scales with batch size and sequence length, often rivaling or exceeding optimizer memory in large-scale training runs.

### How does LoRA reduce memory usage compared to full fine-tuning?

LoRA freezes the base model weights and injects trainable low-rank decomposition matrices into specific layers, reducing trainable parameters by orders of magnitude. Since optimizer memory scales with the number of trainable parameters, LoRA dramatically reduces the memory required for both gradient storage and optimizer state while keeping the base model in lower-precision read-only memory.

### Does gradient checkpointing slow down training?

Yes, gradient checkpointing increases training time by approximately 20-30% because it recomputes activations during the backward pass rather than retrieving them from memory. However, this tradeoff enables training with significantly larger models or batch sizes that would otherwise exceed available GPU memory, making it essential for memory-constrained environments.

### Can quantization reduce memory usage during training, not just inference?

Yes, **quantized training** using libraries like BitsAndBytes can reduce memory usage during training by storing weights and optimizer states in 8-bit or 4-bit formats rather than 32-bit floats. However, as noted in the book’s analysis, quantization alone does not automatically reduce activation memory unless the entire training pipeline—including gradient computation—is adapted to use lower precision throughout.