Memory Planning
3 秒看懂
Memory Planning（内存规划）是 AI 编译器与运行时系统中在编译阶段为张量分配内存地址与生命周期的全局优化技术。它像一名城市规划师，在有限的 GPU/HBM 内存“土地”上，精确安排每一份中间数据的存放位置与起止时间，目标是以最少的内存资源支持最大的 AI 模型计算，从而降低训练与推理的硬件总拥有成本 (TCO)。
3 分钟产业解释
在 AI 计算中，最昂贵的资源不仅是算力 (FLOPs)，更包括高速内存（如 HBM）。一个大语言模型 (LLM) 训练或推理时，其激活值、权重、优化器状态等会迅速耗尽显存。
- 问题：传统方法为每个张量独立预分配固定内存，导致大量内存在中间结果完成后未能及时释放，呈现“闲置”与“碎片化”。
- 解决方案：Memory Planning 通过时间复用与空间复用，让多个张量的生命周期交错并共享同一内存区域。这是一个全局的、静态的优化过程，发生在模型编译阶段，而非运行时。
- 产业价值：它直接放大现有硬件的有效容量，使得同等 GPU 下能够训练更大参数量的模型或使用更大的批量大小 (Batch Size)，提升吞吐和训练效率，是 AI 编译器一项核心软件杠杆。
技术原理
Memory Planning 是将计算图中的张量生命周期映射为内存池分配时序的组合优化过程。
基本流程：
- 输入：经过算子融合、常量折叠等优化后的计算图（如 XLA HLO、MLIR、Torch FX Graph）。图中每个节点定义了输入/输出张量及依赖关系。
- 张量活跃区间分析：通过数据流分析确定每个张量从定义（首次写入）到最后一次被读取的指令跨度，将其标记为一段活跃区间。
- 冲突图构建：若两个张量的活跃区间在时间上有重叠，则它们之间产生冲突边，意味着不能占用同一块内存地址。
- 分配方案求解：以最小的峰值内存为目标，为每个张量分配（基地址，大小，生命周期）三元组，同时满足冲突约束。典型算法包括：
- 图着色 (Graph Coloring) 将张量视为顶点，冲突连边；用最少颜色数（内存块）着色，使相邻顶点颜色不同。
- 线性扫描 (Linear Scan) 按时间顺序遍历活跃区间，每当新张量需要分配时，优先复用已结束的生命周期所释放的内存。
- 整数线性规划 (ILP) 对小规模子图或关键路径精确求解，可获得理论最优方案，但计算代价高。
- 高级约束：
- 内存碎片控制：避免频繁小分配导致无法满足大块连续内存要求，需要在优化目标中加入碎片惩罚或使用大小分级的内存池。
- 流水线并行耦合：在计算与通信重叠的场景下，规划需要考虑数据传输的占用区间和暂存空间。
- 异构内存层：当模型超出 GPU HBM 容量时，规划需将张量在 HBM、CPU DRAM 和 NVMe SSD 之间动态调度，并结合预取与卸载策略。
与运行时协作：编译期产出一张静态分配表，运行时直接依表执行分配与释放，几乎无决策开销；对于动态形状模型，可辅以形状相关的分配方案模板或轻量级运行时调整。
关键参数
- 峰值内存占用 (Peak Memory Usage)：任意执行时刻所有活跃张量的内存使用量最大值。最核心优化目标，直接决定允许的最大模型规模或批量大小。
- 内存利用率 (Memory Utilization Ratio)：计算公式为 (总计算所需内存在不共享时的总和) ÷ (峰值内存占用)。该值越高，表示复用效率越好；经验上优化后可达 2×–5× 复用比（来源：根据编译文献中常见示例估算）。
- 规划时间 (Planning Time)：编译阶段求解分配方案所消耗的 CPU 时间。对于千亿参数模型，典型规划时间可能从数秒到数十分钟，需在方案质量与编译速度间折衷。
- 分配/释放操作数：运行时内存管理 API 调用次数，影响执行开销。静态全局规划可将此数降至近似常量。
- 外部碎片率：内存中因碎片无法满足大块请求的空间占总容量的比例。优质规划算法应将其控制在 5% 以内（行业经验值；公开资料未见统一标准）。
技术路线
| 路线 | 核心思想 | 优点 | 缺点 | 典型代表 |
|---|---|---|---|---|
| 静态编译期规划 | 在编译时求解完整分配方案，生成静态偏移表 | 运行时零开销，利于全局优化，峰值内存可控 | 无法适应动态形状，当图规模极大时求解器耗时 | XLA、TensorRT、TVM 静态图、PyTorch Inductor (部分) |
| 运行时动态规划 | 在算子执行过程中按需分配，利用生命周期信息即时复用 | 灵活处理动态图与控制流，实现简单 | 运行时有决策开销，可能非全局最优 | PyTorch 早期 eager 模式、PaddlePaddle 动态图 |
| 分层异构规划 | 联合建模多级内存（HBM、CPU RAM、NVMe），结合计算-通信重叠 | 可支持超出 GPU 物理容量的超大模型 | 系统复杂度高，数据搬移开销成为新瓶颈 | DeepSpeed ZeRO-Infinity、FlexFlow |
| 基于机器学习的规划 | 使用强化学习或图神经网络预测最优分配策略 | 可发现传统启发式算法难以捕获的模式 | 训练数据获取困难，泛化性待验证 | 学术前沿（见 MLSys 2022–2024 部分论文） |
路线收敛趋势：当前主流 AI 编译器（如 PyTorch 2.x 的 Inductor、OpenXLA）倾向于“静态为主、动态兜底”的混合方式：对静态子图采用编译期全局规划，对动态维度或控制流分支运行时补充轻量调整。
上游
Memory Planning 的上游输入包括：
- AI 模型与框架：开发者使用 PyTorch、JAX、TensorFlow 等框架编写模型代码，生成原始计算图（动态图由追踪或标图转换为中间表示）。
- 编译器前端优化：在进入内存规划之前，编译器通常会执行算子融合、死代码消除、公共子表达式消除、布局转换等 Pass，这些优化会显著改变张量生命周期和数量，直接影响规划难度与效果。
- 并行策略决策：在分布式场景下，张量并行、序列并行、流水线并行等策略会将张量切分并插入通信算子，重构后的计算图是内存规划的新输入。ZeRO 优化器的分片策略也会改变优化器状态的生命周期。
- 硬件内存规格：目标设备的内存层次结构、容量、带宽和对齐要求作为规划约束。
下游
Memory Planning 的输出与下游系统协作：
- 运行时执行引擎：根据分配表，调用底层 API（如 CUDA 的
cudaMalloc/cudaFree或统一内存池）管理实际内存。编译器生成的代码中嵌入分配指令序列。 - 定制内存分配器：为了减少碎片和开销，AI 运行时通常实现内存池或 Slab 分配器，Memory Planning 输出的偏移量直接映射到池内偏移，避免频繁系统调用。
- 硬件内存与缓存：规划结果影响 GPU L2 Cache 和 TLB 的命中率。良好的数据布局与生命周期控制可以提升访问局部性。
- 分析工具与可观测性：规划信息可导出为内存使用曲线和峰值报告，供开发者诊断 OOM 或进一步手工调优。
受益公司
Memory Planning 技术直接或间接惠及以下类型的企业：
- AI 编译器 / 框架提供商：Google (XLA)、Meta (PyTorch Inductor)、NVIDIA (TensorRT / cuDNN graph) 以及开源社区（MLIR、TVM）通过内置先进的 Memory Planning 提升了产品竞争力和生态壁垒。
- AI 芯片公司：如华为昇腾、寒武纪、壁仞科技等自研架构，必须提供自己的编译器栈，其中 Memory Planning 是发挥硬件极致效能的必要组件。2023 年后新发布的 AI 芯片普遍宣称支持 HBM 并配合高级编译器优化（来源：各公司官网）。
- 云计算与 AI 平台厂商：AWS、阿里云、微软 Azure 通过深度集成的编译器优化（包括内存规划）提升公共 AI 服务的规模化和性价比，降低自身运营成本。
- 大模型初创与垂直应用企业：模型训练与部署中应用优化过的编译器，能够节省数万至数十万 GPU·小时费用，是降本增效的直接受益者，如 Anthropic、Midjourney 等（根据公开技术报告提及使用定制优化）。
- 工具链初创：如 OctoML（Apache TVM 的商业支持公司）、d-Matrix（存内计算结合编译器），通过提供内存优化解决方案获取商业价值。
（注：此处仅为技术应用场景陈述，不代表任何投资建议。）
市场规模
- 直接市场：AI 编译器与系统优化软件市场仍处于早期，难以取得精确独立数据。但可参考 AI 训练与推理基础设施市场。据 Forture Business Insights 2024 年报告，全球 AI 基础设施市场 2023 年约 290 亿美元，预计 2030 年达约 1500 亿美元（口径含服务器、存储、软件），编译器优化作为效率提升的使能技术，其价值体现在减少所需硬件数量上。
- 间接经济效益：经 Memory Planning 等编译优化后，大模型训练任务可降低 20%–50% 的峰值显存占用（来源：行业案例综合评估，如 XLA 早期论文、PyTorch 博文），等效节约对应比例的硬件需求。以 Meta 公开的 LLaMA 2 训练使用了约 2.4 万 GPU 年为例，若优化使得需求减半，可节省上亿美元级别的硬件开支。
- 内存市场压力驱动：HBM 市场 2023 年约 23 亿美元，预计 2028 年达 127 亿美元（Yole Intelligence, 2024）。供给紧平衡导致硬件昂贵，软件内存优化因而具有更高经济价值。
（注：以上数字均为公开报告估计，不保证未来精确性。）
玩家对比
| 玩家 / 项目 | 规划方式 | 支持粒度 | 动态形状 | 异构内存 | 成熟度与适用场景 |
|---|---|---|---|---|---|
| XLA (Google) | 静态全局图着色 + 线性扫描混合 | 张量级 | 有限支持（动态维度边界） | 未原生支持，通常与手动交换配合 | 非常成熟，适用于 TPU / GPU 静态图 |
| PyTorch Inductor (Meta) | 基于 Triton 代码生成的静态偏移分配 | 张量级，部分融合内核内部复用 | 支持动态形状的复用策略 (by torch.export) | 与张量并行、ZeRO 协同 | 快速演进，是 PyTorch 2.x 核心后端 |
| TensorRT (NVIDIA) | 基于层的静态内存规划，支持跨层复用 | 层输出张量，规划结合构建时优化 | 部分支持，通过 shape binding | 不直接管理主机内存卸载 | 适用于推理，极致性能 |
| TVM | 静态规划，用户可自定义 pass | 张量级 | 借助 Relay VM 或自定义规划 | 字节跳动的项目中扩展至异构 | 灵活，学术与定制化场景多 |
| DeepSpeed ZeRO | 运行时数据分区与 offloading 调度 | 参数/梯度/优化器状态分区 | 适配动态 forward 流程 | 强力支持 NVMe offload | 分布式训练事实标准之一 |
| MLIR 生态 | 各 dialect 自定义，如 TF 的 memory-planning pass | 张量/缓冲区 | 通过 shape dialect 描述符号维度 | 可组合不同内存层级 | 基础设施，被多个框架集成 |
（对比信息综合自 2024 年各项目文档及社区会议演讲，版本差异可能导致细节变动。）
风险
- 算法复杂度与可扩展性：精确求解内存规划的 ILP 问题 NP 难，对于超大图（万亿参数）面临编译时间爆炸。启发式算法可能陷入局部最优，导致优化效果不及预期。
- 硬件快速演进：HBM4、存内计算 (PIM)、CXL 互联内存等新架构要求规划算法与内存访问模式重新适配，现有编译器团队需持续投入重写。
- 动态性与稀疏性冲击：MoE (混合专家) 和基于动态路由的模型使图结构极不固定，静态全局规划难以覆盖所有执行路径；过度依赖静态规划可能导致反效果。
- 生态碎片化：不同硬件、框架采用不同的内存规划方案，缺乏标准化接口，导致模型可移植性变差，也增加了用户迁移成本。
- 人才瓶颈：同时深刻理解编译器、AI 计算图、内存体系结构的工程师稀缺，制约技术迭代。
- 与更长链路优化的权衡：过度追求峰值内存最小化可能导致数据移动增加或缓存局部性恶化，甚至影响计算延迟。需要与算子调度、并行策略联合优化，系统复杂度陡增。
误读纠偏
- 误读一：“Memory Planning 只是一个好用的 malloc，做内存复用而已。”
事实：这是严重低估。常规内存分配器只管理单次分配/释放。Memory Planning 是系统级、跨张量的时空调度，基于全图依赖和生命周期信息进行协同优化，解的是 NP 难组合优化问题，其决策质量直接影响是否能在单设备上装载大模型。 - 误读二：“只对训练有意义，推理不需要。”
事实：推理场景下内存同样珍贵。高效规划能够：1) 在内存受限的边缘设备上运行更大模型；2) 通过内存复用提升并发批处理数量，直接提高服务吞吐；3) 避免因内存不足而 swap 或切分计算带来的延迟抖动。 - 误读三：“一旦模型在 GPU 上能够运行，内存规划的目标就达成了。”
事实：仅使模型“跑起来”是最低要求；更优的规划还可提升批量大小以加速训练，或通过更好的缓存行为降低推理延迟。这是性能持续挖潜的重要手段。 - 误读四：“动态形状模型无法做内存规划。”
事实：现代编译技术可通过符号维度分析为不同形状生成参数化的分配方案，或在运行时采用轻量动态修正，已可覆盖大部分动态形状需求（如 Torch 2.x 对 dynamic shapes 的支持）。
最新事件
- 2024 年 3 月 GTC：NVIDIA 发布 TensorRT-LLM 更新，内置增强的 KV cache 内存规划，支持更高效的 LLM 推理。
- 2024 年 10 月 PyTorch 2.5 发布：Inductor 后端增强了对于动态形状的 memory planning，引入基于启发式的复用池，使部分动态模型推理峰值内存降低达 30%（官方博客披露，实验环境）。
- 2024 年 MLSys 会议：多篇论文聚焦大模型训练的内存优化，如《MegaScale: Scaling Model Training to More than 10,000 GPUs》中详细描述了静态内存规划与并行策略的联合设计。
- 2025 年初 OpenXLA：社区提出统一 Memory Planning 接口草案，旨在让不同前端框架共享同一内存优化后端。
- 头部云厂商实践：阿里云 PAI 在 2024 年公开技术文章中介绍了其 Blade 编译器的内存复用优化，应用于千亿参数 LLM 训练，节省 40% 显存（口径：对比无规划 baseline）。
（以上事件均来源于公开博客、会议议程和官方发布说明。）
跟踪指标
- 优化后峰值显存占用 (GB) 及复用系数：对比关闭/开启 Memory Planning 的峰值内存，直接衡量收益。
- 编译时间增加比例：规划所导致的编译链路耗时增量，需监控以避免过度牺牲开发者迭代速度。
- 规划覆盖率：在所有产生中间张量的算子中，被纳入规划重复用的张量比例，过低的覆盖率表明仍有大片未优化空间。
- 动态形状场景的退化程度：对同一模型不同输入形状，规划后峰值内存的波动情况；理想状态下应接近最优。
- 下游硬件利用指标：规划前后，GPU SM 利用率、memory bandwidth utilization 的变化，确保内存优化未引入计算饥饿。
- 开源项目版本：关注的编译器（如 PyTorch, XLA, TVM）的 release notes 中与 memory、allocation 相关的 update 条目。
信源
- 系统会议与期刊：OSDI, SOSP, ASPLOS, MLSys, PLDI, SIGCOMM 等会议中关于 tensor lifetime analysis, AI compiler, memory-efficient training 的论文。
- 开源编译器文档与源码：
- XLA:
https://github.com/openxla/xla内部memory_space_assignment相关实现。 - PyTorch Inductor:
torch/_inductor/memory.py与官方博客 (pytorch.org/blog). - TVM:
tvm/src/te/schedule和tvm/src/relay/backend/memory_plan.cc. - MLIR:
mlir/lib/Dialect/MemRef/Transforms/内存相关 pass.
- XLA:
- 行业报告：Yole Intelligence “Memory for AI” (2024)；Fortune Business Insights “AI Infrastructure Market” (2024)；TrendForce 关于 HBM 市况的季度报告。
- 厂商技术博客：Google AI Blog (XLA 更新)；Meta Engineering Blog (PyTorch 编译优化)；阿里云 PAI 产品技术文章；NVIDIA Developer Blog (TensorRT).
- 综述与教程：Poeter et al. “A Survey of Memory Optimization Techniques for Deep Learning” (ACM Computing Surveys, 2023).
（所有数据均采自截至 2025 年 4 月的公开资料，统计口径见文内说明，不包含未公开的商业机密信息。）