Files

Zijie Tian 86633004ca 📝 docs: add 64k memory analysis and test configuration updates

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>

2026-01-14 07:02:09 +08:00

4.0 KiB

Raw Blame History

64k 推理内存分析

本文档分析 Llama 3.1 8B 模型在 64k 长度推理时的内存占用，以及 RTX 3090 (24GB) 上的 OOM 问题。

模型配置

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	✅	✅

建议

64k 推理建议使用 40GB+ 显存的 GPU
RTX 3090/4090 适合 32k 或更短的场景
如必须在 24GB GPU 上运行 64k：
- 使用 RAPIDS RMM 分配器
- 预分配 torch.cat 需要的内存
- 或使用流式处理避免 torch.cat

4.0 KiB Raw Blame History Unescape Escape