1. GLake：大模型时代显存+传 输 管理与优化 蚂蚁集团-基础智能-AI Infra 赵军平 2024.5

2.

3. 提纲 • 背景与技术挑战 • 显存墙、传输墙 • GLake介绍：显存与传输优化 • 训练场景 • 推理场景 • 其他：混部、serverless • 总结：开源

4. 自我介绍 • 赵军平，蚂蚁异构计算与推理引擎负责人 • CCF HPC和存储专委委员，~200 中/ 美技术专利 • 异构加速，虚拟化，K8S，推理优化，文件系统，企业级存储-保护等 • “ 数据密集型应用系统设计” 译者 • 2007~2019：EMC、Dell EMC，Start-up

5. LLM技术挑战1：显存墙 • 显存容量、访存带宽（特别是推理小batch场景） 模型参数量 vs. 单卡显存容量 单卡算力 vs. 访存带宽

6. Llama2-70B推理为例 • 权重与KV cache的显存消耗分析 （假定FP16 未优化） 2机 A100 4机 L40/20 2机 A100 2机 L40/20 A100 *8 2机 L40/20 A10 *8 L40/20 *4 A10 *8 A100 *3 L40/20 *4 A100 *2 L40/20 *6 A100 *4 KV cache显存占用 = H# * H_dim * L# * 2 * 2B * seq# * b#

7. LLM技术挑战2：传输墙 • 访存、卡-卡传输、PCIe传输等发展 << 算力发展

8. 业界方案 思路 精度 性能 范围 模型/算法/实现改造 • DeepSeek • YOCO • MoD, MoE Case by case Case by case case by case 梯度累计 depends 中~高 训练过程适配 重计算 无影响 中 通用 计算换显存 offloading/swapping 无影响 低(同步) ~ 高(异步) 通用 包括UMA 混合精度、小型化 可能影响 - 无损压缩不影响 中-高 Dependents QAT, AMP, INT8 多卡并行:PP, TP, EP, SP - - Depends 模型改造 以上组合

9. 观察 • 显存碎片化 • 训练阶段管理开销，特别是recompute、offloading、并行等叠加使用 • 推理阶段LLM KV cache: 动态生成 • 计算-传输-显存需一体优化 • 多通道、细粒度融合、overlapping • 从单任务 -> 多任务、弹性服务（例如共享混部、ser verless） •  系统层、全局的管理与优化；模型透明  GLake

10. 蚂蚁集团-GLake总体架构 • 显存、传输一体系统优化 • 全局池化： 跨卡、跨进程、 跨容器 • 模型基本透明 • 开源、开放、共建

11. 训练场景：显存碎片问题 • 例子 400MB

12. 显存碎片原因分析  CUDA不支持释放部分 (空闲)显存区域  访存特征的动态变化，LLM更加复杂  数据集长短不一、训练 - >评估、cu DNN workspace等  大模型： 原始pattern 激活recompute后  Recompute  多卡/ stream并发  FSDP/ Deep Speed offloading  Lo RA 不同卡数 Memory Utilization = 不同优化策略 GB

13. 如何显存碎片优化 挑战 思路  减少碎片产生 释放部分使用的block?   ”碎片整理”的性能影响？   CUDA没有直接提供该能力API 搬移-拷贝数据需同步计算，影响性能 复杂性？   Tensor及时释放-> 工具  BFC(first-best-match-sz) -> least frag impact  结合Tensor类型: 同类型(权重/中间值/ws)临近分配  提取tensor访问特征：生命周期相近的临近分配 对用户/模型透明，尽量无感知 优化解决碎片  基于细粒度chunk分配和释放，  异步分配、释放：减少性能影响  无数据搬移: 对应用透明, 降低复杂度

14. 基本思路 2层指针与动态remapping （基于CUDA VMM） - 虚拟：对用户可见，保证连续 - 物理：CHUNK_SIZE（2MB）粒度分配，不要求连续 - Remapping：动态回收chunk与重映射 虚拟内存addr 对齐到chunk (2MB*N) addr handles[ ] 98MB 98MB

15. Remap例子 data Sub-block 物理地址 新分block B1 H1 B3 H3 B1 (驱动内部对物理 显存进行整理) B3 H4 对齐到chunk， B1-1 split block 释放空闲chunks: H1, 3~5 H1 B3-2 H3 H4

16. GLake：优化大模型训练显存碎 片 • PyTorch caching allocator, pluggable; 对模型透明 • 重点：映射元数据管理（无数据copy），策略控制 GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitching, ASPLOS24, 蚂蚁集团、上海交大等

17. Case Study

18. GMLake分配策略设计

19. 效果评测 硬件环境 16 × NVIDIA A100 SXM （80GB） 训练框架 FSDP, DeepSpeed ZeRO, Colossal-AI Gemini 模型 8 LLMs including OPT, Vicuna, GPT-2, GPT-NeoX and so on 不同对比 Batch-size, Strategies, Distributed Training Devices and Different Training Engines 负载 76 fine-tune workloads

20. 其他 • 与PyTorch ExpandSegment对比 • 相同：都借助VMM接口 • 不同： GLake的分配、粘合策略不同。内部实测效果优于 Expand Segment。二者可互补（ 实现中） • 扩展成为PyTorch pluggable allocator，方便lib替换方式使用 • 扩展支持跨stream：机会复用 • 开源、复用策略可灵活配置

21. GLake: 推理场景 • case1: 单模型多进程实例（python） • 挑战：跨进程的权重、中间结果的显存如何复用 • case2：LLM推理KV cache显存管理 • 长短不一动态生成，reser ve低效；动态分配容易碎片 • vLLM：PagedAttention，额外做特殊管理，修改kernel • GLake思路: 坚持系统层面全局优化入手，对模型、 kernel透明

22. case1：多进程推理“显存dedup” • 权重(RO)：全局、细粒度、透明共享，类似存储 “重删” • 全局：支持跨线程、跨进程、跨容器共享 进程1 vP1 • 透明：无需用户参与、无需配置 不同于 Triton Server/ONNXRT: - 粒度：模型所有W[]完全相同 vs. 显存块级（通常比tensor粒度更小） - 范围：单进程 vs. 跨进程、跨容器 - 使用：手工配置 vs. 自动发现，透明使用 进程3 vP3 权重显存块 • 细粒度：层或者block级共享，基于hash和字节比对 • 中间结果：多模型（DAG）、跨进程中间结果的时分复用 进程2 vP2 M1 M2 多模型DAG M3

23. 设计：跨进程、跨模型 • 核心： • 虚拟- 显存指针的动态管理与映射 • 跨进程expor t、引用计数等

24. 效果评测 • 6个进程总显存13.2->4.6GB (- 65%) - 权重优化: - 3.3GB, - 中间结果优化: - 5.3GB,     精度 不影响 性能 不下降 模型 不感知 用户 不用配 25% 40%

25. Case2：LLM KV cache管理 • 问题：动态生成，***GB：与模型Hd#、Layer#以及运行时Batch#、Seq#成正比 • 现有方案：PagedAttention(vLLM): 显式管理block index，显著降低碎片&提高batch • 不足1 ：现有attention kernel需要改造，比如FlashAtt, FlashDecoding • 不足2 ： ~10%~20+% kernel性能开销 • GLake：继续系统层优化思路。模型看到大而连续的虚地址，物理显存动态映射 • 好处1 ：所有atten kernel几乎不需改造， 虚拟地址连续 • 好处2 ： 特定场景下， 20%~350% 的kernel性能优势（特别是GQA decoder） • GLake+FlashAtten/ FlashDecoding/ FlashInfer • 其他：业界最新类似的思路，vAttention（ 未开源） vs. v LLM Paged Atten kernel

26. VMM不足 • 部分API调用耗时波动严重 • cuMemSetAccess：10~100+X • 已反馈至NVIDIA：确认问题并初步分析了原因（ lock attention） • 部分硬件上kernel有影响 • A100：~10% • 其他如A10、L40s、L20、H**等可持平 • 其他：API优化建议反馈至NV

27. 其他场景：混部、serverless 流量波动、性能SLA需求差异 容器/进程保活下，显存自动保存和恢复 显存加载 在 线 批 量 研 发 显存卸载 GPU 多卡 显存+DRAM 多任务共享、混部 - - 显存动态切分 算力动态压制 GPU 多卡 显存+DRAM Serverless - 显存(persistent)自动卸载 按需加载 显存+传输的全局、动态管理  提高GPU利用率、降低成本

28. 总结 • 显存-传输的巨大挑战 • 硬件、互联演进：HBM4, NVLinkSwitch, CXL … • 软件的深层优化 • GLake：系统层全局、动态管理与优化，对模型、用户尽量透明 • 进行中： • 混部+Ser verless显存动态管理 • 推理显存全局统一管理、精细分配 • KV cache Smar t Rebuild • L2 cache, DSM, …

29. 参考 • [ 1 ] GMLake : Efficie nt and Transparent GPU Memor y D e f r a g m e n t a t i o n for Large - scale DNN Training with Vir tual Memor y Stitching, https://a r x i v. o rg / a b s / 2 4 0 1 . 0 8 1 5 6 • [ 2 ] Efficie nt memor y management for large language model ser ving with pag ed attention , https://ar x i v. o rg / a b s / 2 3 0 9 . 0 6 1 8 0 • [ 3 ] v Attention : Dynamic Memor y M a na ge m e nt for Ser ving LLMs without Paged Atte ntion, https://ar x i v. o rg / a b s / 2 4 0 5 . 0 4 4 3 7 • [ 4 ] LLM Inference Unveiled : Sur vey and Roofline Model Insights, https://ar x i v. o rg / a b s / 2 4 0 2 . 1 6 3 6 3 • [ 5 ] SKVQ : Sliding- window Key and Value Cache Q ua ntiza tion for Large Language Models, https://ar x i v. o rg / a b s / 2 4 0 5 . 0 6 2 1 9 • [ 6 ] Attention Store : Cost- effective Attention Reuse across Multi - turn Conver sations in Large Language Model Ser v ing, http s://ar x i v. o rg / a b s / 2 4 0 3 . 1 9 7 0 8 • [ 7 ] You Only Cache Once : Decoder- Decoder A rchitectures for Language Models, http s://ar x i v. o rg / a b s / 2 4 0 5 . 0 5 2 5 4 • [ 8 ] Deep Seek- V2 : A Strong, E conom ical, and E fficient Mixture - of- Exper ts Language Model, http s://ar x i v. o rg / a b s / 2 4 0 5 . 0 4 4 3 4 • [ 9 ] Loong Ser ve : E fficiently Ser v ing Long - context Large Language Models with Elastic Sequence Parallelism , http s://ar x i v. o rg / a b s / 2 4 0 4 . 0 9 5 2 6 • [ 10 ] QSer ve : W 4 A 8 KV 4 Q uantiz ation and System Co - design for E fficient LLM Ser v ing, http s://ar xiv. o rg / a b s / 2 4 0 5 . 04532

30.

31. 谢谢！ 开源： https://github.com/intelligent-machine-learning/glake 欢迎star、欢迎交流、欢迎共建