Управление GPU памятью для LLM: vRAM оптимизация
opensourceaillmgpumemoryvramoptimizationit
Введение: GPU Memory
Проблема
GPU memory is the bottleneck for LLM deployment:
- Llama-3-70B (FP16) needs ~140 GB vRAM
- A100-80GB can only fit one model
- Multiple models = multiple GPUs
- Memory fragmentation wastes resources
Understanding memory breakdown helps optimize:
- Model weights
- KV cache
- Activations
- Overhead
Факт: A single Llama-3-70B (FP16) needs 140 GB — two A100-80GB GPUs.
Memory Breakdown
Где живёт память?
GPU Memory Usage for LLM:
┌─────────────────────────────────────────┐
│ GPU Memory (A100-80GB) │
├─────────────────────────────────────────┤
│ Model Weights: 40% (FP16) │
│ KV Cache: 30% (context=4K) │
│ Activations: 15% (batch=32) │
│ Overhead: 15% │
├─────────────────────────────────────────┤
│ Total: 100% │
└─────────────────────────────────────────┘
Example: Llama-3-70B, FP16, batch=32, context=4K
Weights: 140 GB → needs 2× A100-80GB
KV Cache: ~20 GB (depends on context)
Activations: ~10 GB
Model Weights
Memory per precision:
Model Size FP32 FP16 INT8 INT4
────────────────────────────────────────────────
7B 28 GB 14 GB 7 GB 3.5 GB
13B 52 GB 26 GB 13 GB 6.5 GB
30B 120 GB 60 GB 30 GB 15 GB
65B 260 GB 130 GB 65 GB 32.5 GB
70B 280 GB 140 GB 70 GB 35 GB
KV Cache
KV Cache memory formula:
memory = 2 × num_layers × hidden_size × num_tokens × num_heads × head_size × dtype_size
For Llama-3-70B:
num_layers: 80
hidden_size: 8192
num_heads: 64
head_size: 128
Context=4K, batch=32, FP16:
KV cache = 2 × 80 × 8192 × 4096 × 64 × 128 × 2 bytes
≈ 21 GB
Context=8K, batch=32, FP16:
KV cache ≈ 42 GB
Activations
Activation memory:
activations = 2 × batch_size × max_context × hidden_size × dtype_size
For Llama-3-70B, batch=32, context=4K:
activations = 2 × 32 × 4096 × 8192 × 2
≈ 8.6 GB
For Llama-3-8B, batch=32, context=4K:
activations = 2 × 32 × 4096 × 4096 × 2
≈ 2.1 GB
Quantization
Снижение памяти
Quantization reduces memory:
Method Bits Quality Speed Memory
────────────────────────────────────────────────
FP32 32 100% 1x 1x
FP16 16 100% 2x 0.5x
BF16 16 100% 2x 0.5x
INT8 8 ~99% 4x 0.25x
INT4 4 ~95% 8x 0.125x
INT2 2 ~85% 16x 0.0625x
GGUF Quantization
GGUF quantization levels:
Q2_K: ~2.5 bits/weight - Large quality loss
Q3_K: ~3.5 bits/weight - Noticeable loss
Q4_0: ~4.5 bits/weight - Good balance (default)
Q4_K_M: ~5.0 bits/weight - Best quality
Q5_0: ~5.5 bits/weight - Near-FP16 quality
Q5_K_M: ~5.8 bits/weight - Excellent
Q6_K: ~6.5 bits/weight - Almost perfect
Q8_0: ~8.0 bits/weight - Near-FP16, larger
Quantization Example
# FP16 vs INT4 memory comparison
import torch
model_fp16 = load_model("llama-3-70b", dtype=torch.float16)
model_int4 = load_model("llama-3-70b", quantize="int4")
# Memory comparison
fp16_memory = sum(p.numel() * 2 for p in model_fp16.parameters())
int4_memory = sum(p.numel() * 0.5 for p in model_int4.parameters())
print(f"FP16: {fp16_memory / 1e9:.1f} GB") # 140 GB
print(f"INT4: {int4_memory / 1e9:.1f} GB") # 35 GB
PagedAttention
vLLM подход
PagedAttention solves memory fragmentation:
Traditional approach:
- Pre-allocate KV cache per request
- Wastes memory (over-allocation)
- Fragmentation
PagedAttention:
- Virtual KV cache blocks
- Dynamic allocation
- No fragmentation
- 24x throughput improvement
Block Manager
Block Manager architecture:
Physical Blocks (GPU):
┌────┬────┬────┬────┬────┬────┐
│ B0 │ B1 │ B2 │ B3 │ B4 │ B5 │
└────┴────┴────┴────┴────┴────┘
Virtual Blocks (per request):
Request 1: [V0] → [B0]
Request 2: [V0] → [B1], [V1] → [B2]
Request 3: [V0] → [B3]
Benefits:
- Blocks shared for prefix caching
- Dynamic allocation
- No over-allocation
Memory Calculation
# vLLM memory calculation
def calculate_memory(
num_gpus: int,
gpu_memory: int, # bytes
num_layers: int,
hidden_size: int,
num_heads: int,
head_size: int,
max_context: int,
batch_size: int,
):
# Block size (16 tokens typical)
block_size = 16
# Bytes per block
bytes_per_block = 2 * num_layers * 2 * batch_size * \
num_heads * head_size * 2 # FP16
# Number of blocks
gpu_blocks = (gpu_memory * num_gpus * 0.9) // bytes_per_block
# Max concurrent requests
max_requests = gpu_blocks * block_size // max_context
return {
"gpu_blocks": gpu_blocks,
"max_requests": max_requests,
"bytes_per_block": bytes_per_block,
}
Tensor Parallelism
Распределение модели
Tensor Parallelism splits model across GPUs:
Layer split (TP=2):
GPU 0: GPU 1:
┌──────────────┐ ┌──────────────┐
│ Q_proj (half) │ │ Q_proj (half) │
│ K_proj │ │ K_proj │
│ V_proj │ │ V_proj │
│ Attention │ │ Attention │
│ O_proj (half) │ │ O_proj (half) │
│ W1 (FFN) │ │ W1 (FFN) │
│ W2 (half) │ │ W2 (half) │
└──────────────┘ └──────────────┘
Communication: All-Reduce after attention, All-Gather after FFN
Memory per GPU
With tensor parallelism:
Model: Llama-3-70B (140 GB FP16)
TP=1 (1 GPU): 140 GB → doesn't fit
TP=2 (2 GPU): 70 GB per GPU → fits A100-80GB
TP=4 (4 GPU): 35 GB per GPU → fits T4
TP=8 (8 GPU): 17.5 GB per GPU → fits any GPU
Trade-off:
More GPUs = less memory per GPU
But more communication overhead
PyTorch FSDP
# Fully Sharded Data Parallel
from torch.distributed.fsdp import FSDP, FullStateDictConfig
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-3-70B")
# Wrap with FSDP
fsdp_model = FSDP(
model,
sharding_strategy=ShardingStrategy.FULL_SHARD,
cpu_offload=CPUOffload(offload_params=True),
)
# Memory per GPU:
# Without FSDP: 140 GB (doesn't fit on single GPU)
# With FSDP: ~50 GB (weights + gradients + optimizer states)
Context Length Optimization
KV Cache optimization
KV Cache grows with context length:
Llama-3-70B, batch=32, FP16:
Context KV Cache Total Memory
──────────────────────────────────────
1K 5 GB 155 GB
2K 10 GB 160 GB
4K 20 GB 170 GB
8K 40 GB 190 GB
16K 80 GB 230 GB
32K 160 GB 310 GB
With PagedAttention:
- 24x throughput improvement
- Better memory utilization
- Prefix caching for common prompts
RoPE Scaling
RoPE (Rotary Positional Embeddings) scaling:
Methods:
- Linear: works up to training context
- NTK-aware: extrapolation support
- YARN: focused extrapolation
- LongRoPE: extended context
For longer contexts:
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3-70B",
rope_scaling={
"type": "ntk",
"factor": 2.0, # extends to 8K from 4K
},
)
Monitoring GPU Memory
nvidia-smi
# Monitor GPU usage
watch -n 1 nvidia-smi
# Key metrics:
# GPU-Util: kernel utilization
# Mem.Use: allocated memory
# Perf: P0 (max), P8 (idle)
# Temp: thermal throttling > 85°C
PyTorch Memory Tracking
import torch
# Track memory usage
torch.cuda.reset_peak_memory_stats()
# Run inference
with torch.no_grad():
output = model(input_ids)
# Check memory
peak_memory = torch.cuda.max_memory_allocated()
reserved = torch.cuda.memory_reserved()
print(f"Peak allocated: {peak_memory / 1e9:.2f} GB")
print(f"Reserved: {reserved / 1e9:.2f} GB")
# Memory leaks detection
torch.cuda.memory_summary()
Memory Profiling
# PyTorch memory profiler
from torch.profiler import profile
with profile(
activities=[torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.CUDA],
record_shapes=True,
) as prof:
model(input_ids)
# Export and analyze
prof.export_chrome_trace("trace.json")
# Key insights:
# - Which layers use most memory
# - Kernel execution times
# - Memory allocation patterns
Optimization Strategies
Memory optimization checklist
1. Quantization:
- Use INT4 for deployment
- FP16 for training
- Consider Q4_K_M for best quality
2. Tensor Parallelism:
- Split across GPUs for large models
- Use NVLink for fast inter-GPU
3. Batch Size:
- Reduce batch_size if OOM
- Use gradient accumulation
4. Context Length:
- Use shortest context needed
- Implement context sliding window
5. KV Cache:
- Use PagedAttention (vLLM)
- Enable prefix caching
6. Offloading:
- CPU offload for small batches
- Disk offload for inference
Practical examples
# Memory-efficient inference
from vllm import LLM
llm = LLM(
model="meta-llama/Llama-3-70B",
tensor_parallel_size=2, # 2 GPUs
gpu_memory_utilization=0.9, # 90% of GPU memory
max_model_len=4096, # Limit context
max_num_batched_tokens=4096, # Limit batch tokens
max_num_seqs=256, # Limit concurrent requests
dtype="float16", # Use FP16
quantization="awq", # Use AWQ quantization
)
# This fits on 2× A100-80GB
Заключение
GPU memory management is critical for LLM deployment:
Key concepts:
- Model weights, KV cache, activations
- Quantization reduces memory 2-32x
- PagedAttention eliminates fragmentation
- Tensor parallelism distributes across GPUs
Best practices:
- Use INT4 for production deployment
- Enable PagedAttention (vLLM)
- Monitor memory usage with nvidia-smi
- Optimize batch size and context length
- Use tensor parallelism for 70B+ models