Memory Planning

3 秒看懂

Memory Planning（記憶體規劃）是 AI 編譯器與執行時系統中在編譯階段為張量分配記憶體地址與生命週期的全域性最佳化技術。它像一名城市規劃師，在有限的 GPU/HBM 記憶體“土地”上，精確安排每一份中間資料的存放位置與起止時間，目標是以最少的記憶體資源支援最大的 AI 模型計算，從而降低訓練與推論的硬體總擁有成本 (TCO)。

3 分鐘產業解釋

在 AI 計算中，最昂貴的資源不僅是算力 (FLOPs)，更包括高速記憶體（如 HBM）。一個大語言模型 (LLM) 訓練或推論時，其啟用值、權重、最佳化器狀態等會迅速耗盡視訊記憶體。

• 問題：傳統方法為每個張量獨立預分配固定記憶體，導致大量記憶體在中間結果完成後未能及時釋放，呈現“閒置”與“碎片化”。
• 解決方案：Memory Planning 通過時間複用與空間複用，讓多個張量的生命週期交錯並共享同一記憶體區域。這是一個全域性的、靜態的最佳化過程，發生在模型編譯階段，而非執行時。
• 產業價值：它直接放大現有硬體的有效容量，使得同等 GPU 下能夠訓練更大引數量的模型或使用更大的批次大小 (Batch Size)，提升吞吐和訓練效率，是 AI 編譯器一項核心軟體槓桿。

技術原理

Memory Planning 是將計算圖中的張量生命週期對映為記憶體池分配時序的組合最佳化過程。

基本流程：

1. 輸入：經過運算元融合、常量摺疊等最佳化後的計算圖（如 XLA HLO、MLIR、Torch FX Graph）。圖中每個節點定義了輸入/輸出張量及依賴關係。
2. 張量活躍區間分析：通過資料流分析確定每個張量從定義（首次寫入）到最後一次被讀取的指令跨度，將其標記為一段活躍區間。
3. 衝突圖建置：若兩個張量的活躍區間在時間上有重疊，則它們之間產生衝突邊，意味著不能佔用同一塊記憶體地址。
4. 分配方案求解：以最小的峰值記憶體為目標，為每個張量分配（基地址，大小，生命週期）三元組，同時滿足衝突約束。典型演算法包括：
   • 圖著色 (Graph Coloring) 將張量視為頂點，衝突連邊；用最少顏色數（記憶體塊）著色，使相鄰頂點顏色不同。
   • 線性掃描 (Linear Scan) 按時間順序遍歷活躍區間，每當新張量需要分配時，優先複用已結束的生命週期所釋放的記憶體。
   • 整數線性規劃 (ILP) 對小規模子圖或關鍵路徑精確求解，可獲得理論最優方案，但計算代價高。
5. 高階約束：
   • 記憶體碎片控制：避免頻繁小分配導致無法滿足大塊連續記憶體要求，需要在最佳化目標中加入碎片懲罰或使用大小分級的記憶體池。
   • 流水線並行耦合：在計算與通訊重疊的場景下，規劃需要考慮資料傳輸的佔用區間和暫存空間。
   • 異構記憶體層：當模型超出 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.
• 行業報告：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 月的公開資料，統計口徑見文內說明，不包含未公開的商業機密資訊。）