Hung-yi Lee ML 2025 TA Session: Multi-GPU Training for LLMs

Notes from the TA session of Professor Hung-yi Lee’s (李宏毅) course “Machine Learning in the Era of Generative AI (2025)” at National Taiwan University.

Why Multi-GPU Training?

Large language models are… large. An 8B parameter model requires:

• ~32GB just for model weights (FP32)
• Additional memory for gradients, optimizer states, and activations
• Total: ~128GB+ for training

A single GPU (even an 80GB H100) often can’t fit everything needed for training.

Memory Components in LLM Training

1. Model Parameters: The weights themselves
2. Gradients: Same size as parameters
3. Optimizer States: For Adam optimizer, 2x parameter size (momentum + variance)
4. Activations: Intermediate outputs from each layer (can be huge for long contexts)

Part 1: Optimizing Parameters, Gradients, and Optimizer States

DeepSpeed ZeRO (Zero Redundancy Optimizer)

Microsoft’s DeepSpeed provides three levels of memory optimization:

ZeRO-1: Partition optimizer states across GPUs

• Each GPU stores 1/N of optimizer states
• Minimal communication overhead

ZeRO-2: Partition optimizer states + gradients

• Further memory reduction
• Slightly more communication

ZeRO-3: Partition everything (optimizer + gradients + parameters)

• Maximum memory savings
• Most communication overhead

CPU Offloading

When GPU memory is still insufficient:

• Move optimizer states to CPU RAM
• Move gradients to CPU
• Trade-off: Slower training due to CPU-GPU transfer

Gradient Accumulation

When batch size is limited by memory:

1. Process mini-batches
2. Accumulate gradients without updating
3. Update model after N mini-batches
4. Effective batch size = mini-batch × N

Part 2: Optimizing Activations

The Activation Problem

For long context training (32K+ tokens), activations dominate memory usage:

• Each layer stores intermediate outputs
• Attention computation is O(N²) in sequence length

Flash Attention

A revolutionary kernel that:

• Rewrites attention computation for GPU efficiency
• Reduces memory from O(N²) to O(N)
• Significantly faster than naive implementation

Key insight: Standard attention is memory-bound, not compute-bound. Flash Attention optimizes memory access patterns.

Liger Kernel

An open-source project providing fused kernels for common LLM operations:

• Combines multiple operations into single GPU kernels
• Reduces memory transfers
• Easy to use: just add a decorator

Activation Checkpointing (Gradient Checkpointing)

Trade compute for memory:

• Don’t store all activations during forward pass
• Recompute activations during backward pass as needed
• Significant memory savings with modest compute overhead

Part 3: Quantization

What is Quantization?

Reduce precision of stored values:

• FP32 (32-bit) → FP16 (16-bit) → INT8 (8-bit) → INT4 (4-bit)
• Lossy compression but often acceptable accuracy loss

Mixed Precision Training

• Store master weights in FP32
• Compute forward/backward in FP16
• Best of both worlds: memory savings + numerical stability

Quantization for Inference

Even more aggressive quantization possible:

• 8-bit and 4-bit models common for deployment
• Significant memory reduction with minimal quality loss

Practical Recommendations

Memory Issues?

1. First: Try DeepSpeed ZeRO-1 or ZeRO-2
2. Still OOM: Enable gradient checkpointing
3. Still OOM: Try ZeRO-3
4. Last resort: CPU offloading (will be slow)

Speed Issues?

1. Use Flash Attention (almost always beneficial)
2. Try Liger Kernel for fused operations
3. Use torch.compile for graph optimization

Tools Summary

Tool	Purpose	Difficulty
DeepSpeed	Multi-GPU, memory optimization	Medium
Flash Attention	Fast attention, less memory	Easy
Liger Kernel	Fused kernels	Easy
Quantization	Reduce precision	Easy

Key Takeaways

1. Know your bottleneck: Is it parameters, activations, or optimizer states?
2. Start simple: Don’t over-engineer; add complexity only when needed
3. Flash Attention is essential: Use it for any serious LLM work
4. DeepSpeed ZeRO scales: From single GPU to hundreds