CANN-ATB多卡推理-昇腾NPU上Llama70B怎么切到8张卡Llama2-70B 的权重 140GB单张 Atlas 800I A2 的 64GB 显存放不下。ATB 的多卡推理用 Tensor Parallel 把模型切到多张 NPU 上每张卡只存 1/N 的权重和 KV Cache。Tensor Parallel 的切法Llama2-70B 每层有两个并行维度可以切Attention 部分沿 head 维度切原始32 heads 在 1 张卡 TP216 heads 在卡016 heads 在卡1 TP48 heads 在每张卡 TP84 heads 在每张卡每个 head 独立计算 Attention不需要卡间通信。只在最后做 All-Reduce 把各 head 的输出合并。FFN 部分沿 ff_dim 维度切原始ff_dim28672 在 1 张卡 TP214336 在卡014336 在卡1 TP83584 在每张卡Gate/Up/Down 三个 Linear 各自切 ff_dim。Down Linear 的输出做 All-Reduce 合并。每层的通信模式Input → Attention (切head) → All-Reduce → FFN (切ff_dim) → All-Reduce → Output 不通信 合并 不通信 合并每层 2 次 All-Reduce。TP8 时每次 All-Reduce 通信量约 hidden_dim × 2 bytes 8KB延迟约 0.05ms。32 层 × 2 次 64 次通信总通信延迟约 3.2ms。ATB 的多卡配置fromatbimportLLM# 8卡 TPmodelLLM(meta-llama/Llama-2-70b-hf,devicenpu:0,1,2,3,4,5,6,7,# 8 张 NPUtensor_parallel_size8,dtypetorch.float16)outputmodel.generate(Hello,max_new_tokens100)ATB 自动把权重按 TP 维度切分每张卡加载自己那份权重。不需要手动切模型。权重加载HuggingFace 的权重是完整存储的。ATB 加载时按 TP rank 切分# ATB 内部逻辑简化版forname,paraminhf_model.named_parameters():ifq_projinnameork_projinnameorv_projinname:# Attention 权重按 head 切local_paramparam.chunk(tp_size,dim0)[tp_rank]elifgate_projinnameorup_projinnameordown_projinname:# FFN 权重按 ff_dim 切local_paramparam.chunk(tp_size,dim0)[tp_rank]else:# Embedding 和 LM Head 不切每张卡完整复制local_paramparam加载时间约 30-60 秒主要是磁盘 I/O。ATB 支持权重预加载到共享内存第二次启动只需 5 秒。KV Cache 的分摊TP 模式下 KV Cache 也按 head 维度切。每张卡只存自己负责的 kv_heads 的 KV CacheLlama2-70B (GQA, 8 kv_heads): 单卡 KV Cache 2 × 8 × 128 × seq_len × batch × 2 bytes TP8 每卡 KV Cache 2 × 1 × 128 × seq_len × batch × 2 bytes 单卡的 1/88 卡总 KV Cache 跟单卡一样大但分摊到 8 张卡上每卡显存压力小 8 倍。这就是为什么 TP 不省总显存但让大模型能放进多张小显存的卡。Pipeline Parallel vs Tensor ParallelATB 目前只支持 Tensor Parallel。Pipeline ParallelPP把不同层放到不同卡上层间串行通信。为什么 ATB 不支持 PP昇腾NPU的 HCCS 卡间带宽约 200GB/sAll-Reduce 延迟很低0.05ms。TP 的通信开销小不需要用 PP 来避免通信。PP 在带宽低的集群以太网互联才有优势。性能数据Atlas 800I A2 × 8Llama2-70BTP首 token 延迟 (ms)decode 速度 (tokens/s)单卡显存 (GB)1不支持显存不够--41201,200628852,40035TP8 的 decode 速度比 TP4 快 2 倍——因为每卡的 KV Cache 小了一半留给 batch 的显存更多并发更高。70B 模型用 8 卡 TP 是标配。ATB 的配置很简单指定tensor_parallel_size8就行。权重切分和通信都是自动的。仓库在这里https://atomgit.com/cann/ATB
