Quantitative Analysis: PIM Chip Demands for LLAMA-7B inference

1 Architecture

如果将LLAMA-7B模子参数量化为4bit，则存储模子参数必要3.3GB。那么，至少PIM chip 的存储至少要4GB。

• AiM单个bank为32MB，单个die 512MB，至少必要8个die的芯片。
  
• 8个die集成在一个芯片上。
  
  
  
  • 提供8×16bank级别的访存带宽。
    
  
  
• 整个推理过程完全下放至PIM。
  
  
  
  • CPU把 prompt 传给 Controller
    
  • Controller 控制推理过程，将推理出的token返回给CPU
    
  
  
• Controller
  
  
  
  • ALUs
    
    
    
    • 处置惩罚softmax、Norm和向量乘等。
      
    
    
  • CRAM
    
    
    
    • cache
      
    
    
  • CMEM
    
    
    
    • 推理过程中，保存中间结果
      
    
    
  
  
• Die
  
  
  
  • DieRAM
    
    
    
    • 数据 buffer，Controller broadcast 数据时用到。
      
    
    
  
  
• Bank
  
  
  
  • MACs
    
    
    
    • Multiply-And-Accumulate
      
    • 用于GEMV and GEMM
      
    
    
  • BRAM
    
    
    
    • cache
      
    
    
  • BMEM
    
    
    
    • capacity: 32MB
      
    
    
  
  

2 Data partition

2.1 LLAMA-7B

2.2 Model parameter

在 batch_size=1 的环境下

• prefill 阶段，嵌入prompt，此时为GEMM算子
  
• decode 阶段，推理出一个个token，此时为GEMV算子
  

模子参数划分就是将上图中的矩阵划分至8×16个bank中。
2.2.1 一维划分

2.2.2 二维划分

• 在分布式盘算场景下的常用划分
  
• 优势：通信量小。但是，要求盘算节点间存在通信能力。
  
• 在PIM场景下，无法假设bank间存在通信能力，此时，Controller的reduction开销会很大。
  

2.3 KV cache

2.3.1 attention

• NUM_HEAD 个相互独立的 attention 操纵
  
• max sequence length = (8×16×32×1024×1024 - 6607077376/2)/(32×4096/2) = 15128
  

2.3.2 a bank for a head

• 一个head attention由一个bank执行
  
• 优势：Controller 与 banks 通信少
  
• 劣势：
  
  
  
  • NUM_HEAD < NUM_BANK，3/4的bank访存带宽和算力被浪费。
    
  • bank内不仅要支持MAC，还要支持softmax。
    
  
  

2.3.3 multiple banks for a head

• 一个head attention由多个bank执行
  
• 优势：所有的bank访存带宽和算力得到使用。
  
• 劣势：
  
  
  
  • Controller 与 banks 通信开销变大
    
  • Controller 必要举行softmax和reduction。
    
  
  

3 Demands

本节使用量化方法来分析PIM chip，希望能够复兴以下几个问题：

• CRAM、CMEM、BRAM做多大符合？
  
• Bank级并行带宽必要多少？并行算力必要多少？
  
• Controller 怎样与bank通信？通信带宽必要多少？
  
• Controller 必要提供多大的算力？
  

PIM chip面向端侧推理，一样平常来说，推理的batch size = 1。LLM在推理时，可以大致分为两个阶段：

• prefill
  
  
  
  • 在prefill阶段，模子嵌入prompt。
    
  • 假设嵌入prompt的长度为N，则在这个过程中模子参数会被复用N次。
    
    
    
    • 典范算子为GEMM
      
    
    
  
  
• decode
  
  
  
  • 在decode阶段，模子以自回归的方式推理出一个个token
    
  • 假设推理出S个token
    
    
    
    • 每推理出一个token，则在这个过程中都必须扫描一遍模子参数和kv cache。
      
    • 典范算子为GEMV
      
    
    
  
  

3.1 bandwidth

显然，LLM模子在decode阶段的瓶颈在于访存带宽。定量分析decode过程，也就可以分析出在给定访存带宽下，模子推理的速度。
3.1.1 hypothesis

• 假设推理的sequence的长度为L
  
• 在decode阶段，Controller和bank内的算力均可以吃满访存带宽。
  
  
  
  • 对于GEMV算子，Operational intensity = 2 Ops/weight byte
    
  • 这个假设完全公道。
    
  
  
• 并行intra-bank bandwidth 总带宽为BM
  
  
  
  • 对于AiM，BM = 8 × 512 GB/s
    
  
  
• Controler-bank bandwidth 总带宽为CBM
  
  
  
  • 一样平常来说，CBM << BM
    
  • 对于AiM，CBM = 8 × 32 GB/s
    
  • 这儿应该举行更加精细的讨论
    
    
    
    • 每次 Controler-bank 通信的基础开销（时延）设置为                                                            λ                                                  \lambda                                 λ
      
    
    
    
    
  
  
• intra-Controler bandwidth 带宽为CM
  
  
  
  • 能够到达类似CPU 100GB/s的访存带宽？
    
  
  
• 在decode过程，推理出一个token的时延 = bank内并行访存（GEMV）的时延 + Controler-bank 通信的时延 + Controler 内访存（softmax、Norm 和 reduction 等）的时延
  

3.1.2 intra-bank bandwidth

bank内并行访存（GEMV）的时延包罗两部门：模子参数相关的GEMV的时延和kv cache相关的GEMV的时延。

• 模子参数相关的GEMV的时延
  
  
  
  • 这个时延非常好算，其实就是 模子总参数量/BM。
    
  • 对于AiM，这个时延为 (6607077376 / 2 / 1024 / 1024 / 1024)/(BM) = 7.5 × 10^-4 s
    
  
  
• kv cache相关的GEMV的时延
  
  
  
  • kv cache 大小
    
    
    
    • L×H×NUM_LAYER×2
      
    
    
  • 与kv cache相关的attention算子的时延盘算就比较复杂，因为其有两者盘算方案。
    
  • a bank for a head
    
    
    
    • NUM_HEAD 的数量为32，bank数量为8×16
      
    • bank并行访存带宽的使用率为25%
      
    • 时延 = ((L × H + NUM_HEAD × L + L × H) × NUM_LAYER / 2 / 1024 / 1024 / 1024)/(BM / 4)
      
    • 如果L=4096，时延 = 4.9 × 10^-4 s
      
    
    
  • multiple banks for a head
    
    
    
    • 4 banks for a head
      
    • 时延 = ((L × H + L × H) × NUM_LAYER / 2 / 1024 / 1024 / 1024)/(8 × 512)
      
    • 如果L=4096，时延 = 1.2 × 10^-4 s
      
    
    
  
  

3.1.3 Controler-bank bandwidth

• 模子参数相关的通信时延
  
  
  
  • 通信次数
    
    
    
    • 运算一个GEMV算子，必要两次通信，向bank发送向量，取回结果。
      
    • 通信次数 2 × (1 + NUM_LAYER × 7)
      
    
    
  • 通信量
    
    
    
    • (V + H) + NUM_LAYER × (4 × 2 × H + 3 × (H + DIM_MLP))
      
    
    
  • 通信时延
    
    
    
    • (2 × (1 + NUM_LAYER × 7))                                                            λ                                                  \lambda                                 λ + ((V + H) + NUM_LAYER × (4 × 2 × H + 3 × (H + DIM_MLP))) /2 / 1024 / 1024 / 1024 / CBM
      
    • UPMEM中                                                            λ                                                  \lambda                                 λ=0.0001s，AiM中CBM=8 × 32 GB/s，此时，时延 = 450                                                            λ                                                  \lambda                                 λ + 4.61×10^-6 s
      
    
    
  
  
• kv cache相关的通信时延
  
  
  
  • a bank for a head
    
    
    
    • 通信次数
      2 × NUM_LAYER
      
    • 通信量
      NUM_LAYER × (2H)
      
    • 通信时延
      (2 × NUM_LAYER)                                                                 λ                                                      \lambda                                    λ + NUM_LAYER × (2H) /2 / 1024 / 1024 / 1024 / (CBM/4)
      64                                                                 λ                                                      \lambda                                    λ + 1.91 × 10^-6 s
      
    
    
  • multiple (4) banks for a head
    
    
    
    • 通信次数
      2 × NUM_LAYER × 2
      
    • 通信量
      NUM_LAYER × (4H + NUM_HEAD × L + NUM_HEAD × L + 4H)
      
    • 通信时延
      (2 × NUM_LAYER × 2)                                                                 λ                                                      \lambda                                    λ + NUM_LAYER × (4H + NUM_HEAD × L + NUM_HEAD × L + 4H) /2 / 1024 / 1024 / 1024 / CBM
      192                                                                 λ                                                      \lambda                                    λ + 1.72 × 10^-5 s
      
    
    
  
  

3.1.4 intra-Controler bandwidth

• 模子参数相关的GEMV的时延
  
  
  
  • 主要是做处置惩罚softmax、Norm和向量乘等。
    
  • 处置惩罚数据量
    
    
    
    • V + NUM_LAYER × (H + H + H + 2 × DIM_MLP + H)
      
    
    
  • 处置惩罚时延
    
    
    
    • (V + NUM_LAYER × (H + H + H + 2 × DIM_MLP + H)) /2 / 1024 / 1024 / 1024 / CM
      
    • 5.87 × 10^-6 s
      
    
    
  
  
• kv cache相关的GEMV的时延
  
  
  
  • a bank for a head
    
    
    
    • 无
      
    
    
  • multiple (4) banks for a head
    
    
    
    • softmax 和 reduction
      
    • 处置惩罚数据量
      NUM_LAYER × (L + 4 × H)
      
    • 处置惩罚时延
      NUM_LAYER × (L + 4 × H) /2 / 1024 / 1024 / 1024 / CM
      3.05 × 10^-6 s
      
    
    
  
  

3.1.5 summary

• sequence len = 4096，推理出一个token的总时延
  
  
  
  • a bank for a head
         intra-bank 模子参数intra-bank qkvControler-bank 模子参数Controler-bank qkvintra-Controler 模子参数intra-Controler qkvtotal7.5 × 10^-4 s4.9 × 10^-4 s450                                                                      λ                                                          \lambda                                       λ + 4.61×10^-6 s64                                                                      λ                                                          \lambda                                       λ + 1.91 × 10^-6 s5.87 × 10^-6 s0514                                                                      λ                                                          \lambda                                       λ+1.25 × 10^-3 s
    
  • multiple (4) banks for a head
         intra-bank 模子参数intra-bank qkvControler-bank 模子参数Controler-bank qkvintra-Controler 模子参数intra-Controler qkvtotal7.5 × 10^-4 s1.2 × 10^-4 s450                                                                      λ                                                          \lambda                                       λ + 4.61×10^-6 s192                                                                      λ                                                          \lambda                                       λ + 1.72 × 10^-5 s5.87 × 10^-6 s3.05 × 10^-6 s642                                                                      λ                                                          \lambda                                       λ+1.27 × 10^-3 s
    
  • 如果可以将Controler-bank通信时延优化掉，1s 可以推理出500+ token。这时间，性能是冗余的。
    
  • Controler要放在PIM chip上，否则Controler-bank通信基础开销会成为系统瓶颈。
    
  • 减少成本，将上述三种带宽均减小一个数量级，系统吞吐 50+ token/s，性能也能满足需求。
    
  
  

3.2 Computing

LLM模子在prefill阶段的瓶颈在于硬件算力。
3.2.1 Controler

• Controler 部门负担softmax、Norm 和 reduction 访存麋集性算子
  
• Controler 的算力只要能吃满 Controler 的访存带宽就可以。
  
• CRAM容量
  
  
  
  • L最长16K，能够放入L长度的向量举行softmax就可以
    
  • CRAM容量最小 8KB，设置64KB是公道的。
    
  • 固然，如果追求更大的并行度，设置 NUM_HEAD × 8KB = 256 KB
    
  
  
• CMEM容量
  
  
  
  • 在prefill阶段，要能够放入L×L的矩阵，最小128MB
    
  • 设置256MB是公道的
    
  
  

3.2.2 bank

• bank 负担GEMM算子
  
  
  
  • 嵌入prompt的长度为N，模子参数加载进入cache可以重用N次
    
  • 焦点还是MAC
    
    
    
    
  
  
• 可以将 N 以 batch_size (比方16) 为粒度举行切分，以tile的方式举行GEMM
  
  
  
  • 可以设置BRAM的大小为 16 KB ~ 64 KB，更大的BRAM可以允许更大的tile。
    
  
  
• 必要算力：BM×2×batch_size ops
  

3.3 conclusion

如果decode阶段必要50+ token/s的推理速度：

• CRAM、CMEM、BRAM做多大符合？
     CRAMCMEMBRAM256 KB256MB16 KB ~ 64 KB
  
• Bank级并行带宽必要多少？并行算力必要多少？
  
  
  
  • 512 GB/s
    
  • 512×2 ~ 512×2×16 Gops
    
    
    
    • prefill 阶段的embedding速度和并行算力相关。
      
    • 512×2 Gops 对应50+ token/s的embedding速度
      
    • 512×2×16 Gops 对应50×16+ token/s的embedding速度
      
    
    
  
  
• Controller 怎样与bank通信？通信带宽必要多少？
  
  
  
  • Controller 必须在片内，降低通信时延
    
  • 4 ~ 32GB/s
    
  
  
• Controller 必要提供多大的访存带宽和算力？
  
  
  
  • 16 GB/s
    
  • 32 Gops
    
  
  

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。