Glake：效透明的大模型显存管理和优化-赵军平.pdf

1、GLake：大模型时代显：大模型时代显存存+传传输输 管理与优化管理与优化蚂蚁集团-基础智能-AI Infra赵军平2024.5提纲提纲 背景与技术挑战 显存墙、传输墙 GLake介绍：显存与传输优化 训练场景 推理场景 其他：混部、serverless 总结：开源自我介绍自我介绍 赵军平，赵军平，蚂蚁异异构构计计算算与与推推理理引引擎擎负责人 CCF HPC和存储专委委员，200 中/美技术专利 异构加速，虚拟化，K8S，推理优化，文件系统，企业级存储-保护等“数据密集型应用系统设计”译者 20072019：EMC、Dell EMC，Start-upLLM技技术术挑挑战战1：显存显存墙墙 显

2、存容量容量、访存带带宽宽（特别是推理小batch场景）模型参数量 vs.单卡显存容量单卡算力 vs.访存带宽Llama2-70B推推理为理为例例 权重与KV cache的显存消耗分析（假定FP16 未优化）A10*8L40/20*4 A100*2A100*32机机 A1004机机 L40/202机机 A1002机机 L40/20A100*82机机 L40/20L40/20*6A100*4A10*8L40/20*4KV cache显存占用=H#*H_dim*L#*2*2B*seq#*b#LLM技技术术挑挑战战2：传输传输墙墙 访存、卡-卡传输、PCIe传输等发展 多任务、弹性服务（例如共享混部、

3、serverless）系统层、全局的管理与优化；模型透明 GLake蚂蚁集蚂蚁集团团-GLake总体总体架架构构 显存、传输一体系统优化 全局池化：跨卡、跨进程、跨容器 模型基本透明 开源、开放、共建训练场景：训练场景：显存显存碎碎片问片问题题 例子400MB显存碎片原显存碎片原因分因分析析CUDA不支持释放部分(空闲)显存区域访存特征的动态变化，LLM更加复杂 数据集长短不一、训练-评估、cu DNN workspace等 大模型：Recompute 多卡/stream并发 FSDP/Deep Speed offloading Lo RA不同优化策略 GB不同卡数Memory Utiliza

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

5、 无数据搬移无数据搬移:对应用透明,降低复杂度思路基本思路基本思路2层指针与动态remapping（基于CUDA VMM）-虚拟虚拟：对用户可见，保证连续-物理物理：CHUNK_SIZE（2MB）粒度分配，不要求连续-Remapping：动态回收chunk与重映射虚拟内存addr对齐到chunk(2MB*N)addrhandles 98MB98MBRemap例子例子Sub-blockB1B3H3H4dataH1B1-1B3-2对齐到chunk，split blockB1B3H1H3H4释放空闲chunks:H1,35新分block(驱动内部对物理 显存进行整理)物理地址GLake：优化优化大大

6、模模型训型训练练显存显存碎碎片片 PyTorch caching allocator,pluggable;对模型透明 重点：映射元数据管理（无数据copy），策略控制GMLake:Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training withVirtual Memory Stitching,ASPLOS24,蚂蚁集团、上海交大等Case StudyGMLake分分配配策策略设略设计计效果评测效果评测硬件环境硬件环境16 NVIDIA A100 SXM（80GB）训练框架训练框架FSDP