Model Architecture Mamba 基准测试

# Mamba 性能基准测试 ## 推理速度对比 ### 吞吐量 (tokens/sec) **Mamba-1.4B 与 Transformer-1.3B** 在单张 A100 80GB 上的对比： | 序列长度 | Mamba-1.4B | Transformer-1.3B | 加速比 | |----------------|------------|------------------|---------| | 512 | 8,300 | 6,200 | 1.3× | | 1024 | 7,800 | 4,100 | 1.9× | | 2048 | 7,200 | 2,300 | 3.1× | | 4096 | 6,800 | 1,200 | 5.7× | | 8192 | 6,400 | 600 | **10.7×** | | 16384 | 6,100 | OOM | ∞ | **核心见解**：加速比随序列长度增加而增长 (Mamba 为 O(n)，Transformer 为 O(n²)) ### 延迟 (ms per token) **生成延迟** (batch size 1, 自回归生成)： | 模型 | 首个 Token | 每个 Token | 总计 100 Tokens | |-------|-------------|-----------|------------------| | Mamba-130M | 3 ms | 0.8 ms | 83 ms | | Transformer-130M | 5 ms | 1.2 ms | 125 ms | | Mamba-1.4B | 12 ms | 3.2 ms | 332 ms | | Transformer-1.3B | 18 ms | 8.5 ms | 868 ms | | Mamba-2.8B | 20 ms | 6.1 ms | 631 ms | | Transformer-2.7B | 35 ms | 18.2 ms | 1855 ms | **Mamba 优势**：无论上下文长度如何，每个 token 的延迟保持恒定 ## 内存占用 ### 训练内存 (BF16, 单 GPU) **Mamba-1.4B** 训练内存分解： | 序列长度 | 激活值 | 梯度 | 优化器 | 总计 | 相比 Transformer | |----------------|-------------|-----------|-----------|-------|----------------| | 512 | 2.1 GB | 3.2 GB | 11.2 GB | 16.5 GB | 0.9× | | 1024 | 3.8 GB | 3.2 GB | 11.2 GB | 18.2 GB | 0.6× | | 2048 | 7.2 GB | 3.2 GB | 11.2 GB | 21.6 GB | 0.4× | | 4096 | 14.1 GB | 3.2 GB | 11.2 GB | 28.5 GB | 0.25× | | 8192 | 28.0 GB | 3.2 GB | 11.2 GB | 42.4 GB | 0.15× | **注意**：Transformer 在 40GB A100 上于 8K 序列长度时会发生 OOM (内存溢出) ### 推理内存 (FP16, batch size 1) | 模型 | KV Cache (8K 上下文) | 状态 (Mamba) | 比率 | |-------|------------------|---------------|-------| | 130M | 2.1 GB | 0 MB | ∞ | | 370M | 5.2 GB | 0 MB | ∞ | | 1.4B | 19.7 GB | 0 MB | ∞ | | 2.8B | 38.4 GB | 0 MB | ∞ | **Mamba 不存储 KV Cache** - 每个 token 的内存占用恒定！ 实际 Mamba 状态大小： - 130M: ~3 MB (d_model × d_state × n_layers = 768 × 16 × 24) - 2.8B: ~13 MB (2560 × 16 × 64) ## 语言建模基准 ### 常用数据集上的困惑度 (Perplexity) **在 The Pile (300B tokens) 上训练的模型**： | 模型 | 参数量 | Pile (验证集) | WikiText-103 | C4 | Lambada | |-------|--------|------------|--------------|-----|---------| | Pythia | 160M | 29.6 | 28.4 | 23.1 | 51.2 | | **Mamba** | **130M** | **28.1** | **26.7** | **21.8** | **48.3** | | Pythia | 41