86633004ca
Add comprehensive memory analysis for 64k inference on Llama 3.1 8B: New documentation: - docs/64k_memory_analysis.md: GPU-only vs offload memory analysis, OOM root cause (memory fragmentation), RTX 3090 limitations, theoretical vs actual memory usage breakdown Test configuration updates: - tests/test_ruler.py: Add --num-kv-buffers parameter for ring buffer size tuning (default 4, can reduce to 1 for lower memory) - Update default data_dir to ruler_64k - Update default max_model_len to 65664 for 64k support CLAUDE.md updates: - Add 64k_memory_analysis.md to documentation index - Document num_kv_buffers parameter in Configuration section - Add 64k hardware requirements note to Model Limits Key findings: 64k inference requires ~26GB (GPU-only) or ~23GB (offload) due to memory fragmentation on 24GB GPUs, making A100 (40GB+) the recommended hardware for 64k workloads. Co-Authored-By: Claude <noreply@anthropic.com>
132 lines
4.0 KiB
Markdown
132 lines
4.0 KiB
Markdown
# 64k 推理内存分析
|
||
|
||
本文档分析 Llama 3.1 8B 模型在 64k 长度推理时的内存占用，以及 RTX 3090 (24GB) 上的 OOM 问题。
|
||
|
||
## 模型配置
|
||
|
||
```python
|
||
hidden_size = 4096
|
||
intermediate_size = 14336
|
||
num_layers = 32
|
||
num_heads = 32
|
||
num_kv_heads = 8
|
||
head_dim = 128
|
||
seq_len = 65536
|
||
dtype = bfloat16 (2 bytes)
|
||
```
|
||
|
||
## 理论内存占用
|
||
|
||
### GPU Only 模式
|
||
|
||
| 组件 | 计算公式 | 内存占用 |
|
||
|------|----------|----------|
|
||
| 模型权重 | 8.03B × 2 bytes | **16.06 GB** |
|
||
| KV Cache | 32 × 65536 × 8 × 128 × 2 × 2 | **8.19 GB** |
|
||
| Prefill 激活值峰值 | max(QKV, MLP) | **~2 GB** |
|
||
| **总计** | | **~26 GB** |
|
||
|
||
**结论**：GPU only 模式需要 ~26 GB，**RTX 3090 (24GB) 无法运行**。
|
||
|
||
### CPU Offload 模式
|
||
|
||
| 组件 | 计算公式 | 内存占用 |
|
||
|------|----------|----------|
|
||
| 模型权重 | 8.03B × 2 bytes | **16.06 GB** |
|
||
| Ring buffer | num_kv_buffers × seq_len × 128 KB/token | 258-1034 MB |
|
||
| GPU KV blocks | num_gpu_blocks × block_size × 128 KB/token | 256 MB (2 blocks) |
|
||
| Per-layer decode buffer | 32 layers × 缓冲 | 128 MB |
|
||
| 激活值峰值 (chunked) | chunk_size × hidden_size × 2 | ~50 MB |
|
||
| PyTorch 开销 | CUDA 上下文 + 碎片 | ~5-6 GB |
|
||
| **理论小计** | | **~17.5 GB** |
|
||
| **实际需求** | | **~23 GB** |
|
||
|
||
**配置参数**：
|
||
- `num_kv_buffers`: Ring buffer 大小 (1-4)，默认 4
|
||
- `num_gpu_blocks`: GPU 上的 KV cache block 数量
|
||
- `block_size`: 每个 block 的 token 数
|
||
|
||
## OOM 问题分析
|
||
|
||
### 实际观测（RTX 3090, num_kv_buffers=1）
|
||
|
||
```
|
||
PyTorch allocated: 22.49 GB
|
||
PyTorch reserved: 429 MB
|
||
Free: 306 MB
|
||
Total available: 735 MB
|
||
Failed to allocate: 508 MB (torch.cat)
|
||
```
|
||
|
||
### 内存碎片来源
|
||
|
||
| 来源 | 说明 | 影响 |
|
||
|------|------|------|
|
||
| Binned 分配器 | PyTorch 使用固定大小的内存池 | 中等 |
|
||
| torch.compile 缓存 | 编译后的 kernel 代码和常量 | 高 (~2-3 GB) |
|
||
| 频繁分配/释放 | chunked 处理中每个 chunk 的创建销毁 | 高 |
|
||
| 不同大小张量 | (128,4096), (65536,6144) 等 | 中等 |
|
||
|
||
### torch.cat 内存需求
|
||
|
||
Chunked MLP 处理（chunk_size=128）：
|
||
```
|
||
65536 / 128 = 512 chunks
|
||
每个 chunk 输出: (128, 4096) × 2 bytes = 1 MB
|
||
torch.cat 拼接需要: (65536, 4096) × 2 bytes = 508 MB (连续)
|
||
```
|
||
|
||
## 已尝试的优化
|
||
|
||
| 优化项 | 效果 |
|
||
|--------|------|
|
||
| 移除 `@torch.compile` | PyTorch: 23.13 → 22.80 GB (-300 MB) |
|
||
| 减少 `num_kv_buffers` (4→1) | Ring buffer: 1034 → 258 MB (-776 MB) |
|
||
| Chunked QKV/MLP/LayerNorm | 峰值激活: ~2 GB → ~50 MB |
|
||
| 降低 GPU 利用率 (0.9→0.75) | 无明显效果 |
|
||
| 减小 chunk_size (4096→128) | 峰值降低，但 torch.cat 需要连续内存 |
|
||
|
||
### 最终状态
|
||
|
||
```
|
||
理论需求: ~17.5 GB
|
||
实际分配: 22.49 GB
|
||
剩余空间: 735 MB (306 MB + 429 MB reserved)
|
||
分配失败: 508 MB (torch.cat 需要连续内存)
|
||
```
|
||
|
||
## 结论
|
||
|
||
### 根本原因
|
||
|
||
**不是绝对内存不足，而是内存碎片导致的分配失败**。
|
||
|
||
理论需求 17.5 GB < 24 GB，但由于：
|
||
- PyTorch 开销（CUDA 上下文、碎片）：~5-6 GB
|
||
- torch.compile 缓存：~2-3 GB（已移除）
|
||
- 内存碎片导致无法分配 508 MB 连续块
|
||
|
||
### 硬件限制
|
||
|
||
| GPU | 显存 | 64k GPU Only | 64k Offload |
|
||
|-----|------|--------------|--------------|
|
||
| RTX 3090 | 24 GB | ❌ | ⚠️ 碎片问题 |
|
||
| RTX 4090 | 24 GB | ❌ | ⚠️ 碎片问题 |
|
||
| A100 | 40 GB | ✅ | ✅ |
|
||
| A100 | 80 GB | ✅ | ✅ |
|
||
|
||
### 建议
|
||
|
||
1. **64k 推理建议使用 40GB+ 显存的 GPU**
|
||
2. RTX 3090/4090 适合 32k 或更短的场景
|
||
3. 如必须在 24GB GPU 上运行 64k：
|
||
- 使用 RAPIDS RMM 分配器
|
||
- 预分配 torch.cat 需要的内存
|
||
- 或使用流式处理避免 torch.cat
|
||
|
||
## 参考
|
||
|
||
- [PyTorch 内存管理文档](https://docs.pytorch.org/docs/stable/generated/torch.cuda.memory.memory_stats.html)
|
||
- [PyTorch 内存碎片讨论](https://discuss.pytorch.org/t/how-to-reduce-memory-fragmentation-when-enable-expandable-segments/221805)
|
||
- [STWeaver - 减少 79% 内存碎片](https://arxiv.org/html/2507.16274v1)
|