Skip to content

显存优化(Memory Optimization)

专题类型:主学习路线 主服务目标:显存预算与资源取舍

页面导语

本专题研究训练和推理中的显存对象、生命周期与预算取舍,回答显存被什么占用、压力出现在哪个阶段、优化代价转移到哪里,以及当前方案是否值得采用。

训练侧关注参数、梯度、optimizer state、activation 和临时张量;推理侧关注权重、KV Cache、请求并发和临时 attention 空间。两者共享 dtype、内存层级和带宽基础,但训练与推理的项目证据分别记录。

本专题适合希望理解显存占用,并能在有限硬件上做资源取舍的学习者。学习从显存对象和账本开始,再根据问题进入训练、推理或分布式分支。

如何开始

  • 主学习路线: 从 Part 02 的 2.5 反向传播与显存优化 进入,再按下方 Task0–6 表格学习;需要补通用训练计算图时先看 Part 00 · 07 自动求导与反向传播
  • 快速上手: 训练显存不足时,先完成 Task0–2 的机制练习,再进入 73、76;推理 Cache、量化或多卡问题直接从路线表对应分支进入。
  • 按需回补: 需要 GPU 内存层级、LoRA / QLoRA 或模型架构基础时,从路线表的共享入口补充,不要求先完成所有扩展内容。

主学习路线与验证出口

主线先完成 Task0–2,建立显存对象、账本和单机策略;之后按问题进入训练侧 Task3,或进入推理侧 Task4–5。Task6 提供多卡与系统级扩展。每个 Task 都按“机制 → 策略 → 验证出口”组织,扩展内容不要求全部作为共同前置。

路线图用于查看 Task0–6 的学习顺序,以及训练、推理和分布式验证分支。 显存优化路线图:从显存账本到资源决策

知识地图补充显存对象、训练/推理策略和证据升级之间的关系;具体 Notebook 和项目入口以路线表为准。

显存优化知识地图:对象、策略与证据

Task1 的共同前置只保留 dtype → 参数规模 → 硬件条件 → 显存账本 这条核心链;Attention、混合精度、FlashAttention 和模型架构作为共享支撑或按需扩展,不要求在进入 Task2 前全部完成。

Task学习内容核心问题主学习线 / 项目入口学习顺序专题正文主入口
Task0显存对象与生命周期哪些状态会产生、驻留并在 backward 后释放?Part 00 · 07 自动求导与反向传播Part 02 · 18 激活与损失反向Part 02 · 17 注意力反向传播与自定义自动求导;CPU 检查局部梯度、saved tensors、梯度和 activation 生命周期计算图 → 局部 backward → Attention backward → 状态驻留与释放02 训练侧显存压力
Task1dtype、模型规模、硬件与显存账本当前显存压力来自哪个对象,理论容量和实际峰值应如何估算?核心: Part 01 · 01 数据格式与混合精度Part 01 · 02 参数量与算力推导Part 01 · 03 GPU 物理架构与内存层级Part 01 · 06 显存计算与 ZeRO 优化共享支撑: Part 02 · 04 多头注意力Part 01 · 12 Tensor Core 与混合精度Part 01 · 14 FlashAttention 显存模型架构扩展: Part 02 · 05 LLaMA3 Block 教程Part 02 · 06 MoE 路由器Part 02 · 07 MoE 负载均衡损失Part 02 · 08 架构技巧Part 02 · 61 架构验证核心: dtype → 参数规模 → 硬件条件 → 显存账本;共享支撑按需回补;架构扩展不作为共同前置01 显存账本与指标
Task2单机训练显存策略显存不够时,应该用更小的 micro-batch、更多重算,还是 CPU-GPU 搬运来换取空间?Part 02 · 12 梯度累积Part 02 · 19 激活检查点Part 02 · 42 激活卸载;CPU 检查逻辑、梯度对齐和状态变化micro-step → 重算 → CPU-GPU 搬运03 检查点与卸载
Task3训练侧测量与预算决策哪个训练策略在固定 workload、质量门槛和显存上限下值得采用?机制入口:Part 00 · 20 性能剖析与显存账本Part 01 · 13 性能分析与瓶颈定位;项目链:Part 02 · 73 训练性能分析Part 02 · 76 激活检查点与卸载对比Part 02 · 75 显存预算压缩Part 02 · 74 Profiling 驱动的显存优化;按需回补:Part 00 · 17 性能分析基础Part 00 · 18 显存分析与优化Part 00 · 19 调试与异常定位测量对象与指标 → 固定 workload → baseline → 策略比较 → 预算敏感性 → trace 解释06 基准测试与权衡决策
Task4推理侧 KV Cache 与容量上下文和并发增加时,KV Cache 为什么成为容量边界,如何组织、复用和验证?Part 01 · 11 KV Cache 与显存增长Part 02 · 22 vLLM 分页注意力Part 02 · 34 前缀缓存与分块预填充;项目 Part 02 · 66 推理性能对比实验Part 02 · 69 前缀缓存基准;架构扩展 Part 02 · 71 MLA 与 KV Cache 结构基准Part 02 · 24 SGLang 基数注意力增长 → 分页 → 复用 → 容量验证;扩展:RadixAttention / MLA04 推理 Cache 与显存预算
Task5量化与显存容量扩展压缩哪类对象、在什么时候压缩,才能真正换来更大的模型、上下文或并发?Part 01 · 21 量化理论与 INT4/INT8Part 02 · 25 W8A16 量化Part 02 · 40 GPTQ 与 AWQ 权重量化Part 02 · 41 FP8 与 KV Cache 量化 → 项目 Part 02 · 67 量化推理与部署对象与时机 → 权重格式 → 量化算法 → backend → 显存 / 质量验证05 量化作为显存工具
Task6分布式显存与系统级扩展单卡放不下时如何分摊状态,并解释通信、重算和搬运代价?分布式:Part 02 · 27 ZeRO 优化器模拟Part 02 · 28 Pipeline 并行微批次Part 02 · 29 Tensor 并行模拟Part 02 · 79 分布式并行基准 / Part 02 · 80 MoE 专家并行基准 / Part 02 · 81 分布式推理逻辑验证;Profiling 作为共享扩展,复用 Part 01 · 13 性能分析与瓶颈定位Part 02 · 74 Profiling 驱动的显存优化分布式切分 → 单卡显存分摊 → 通信代价 → 多卡证据;Profiling 不作为本 Task 的共同前置07 分布式显存与系统扩展

正文页按机制聚合,不是一 Task 一页:02 是 Task0–2 的训练侧机制入口,06 是 Task3 的训练侧决策入口,07 是 Task6 的分布式扩展入口;Task1、Task2、Task4、Task5 分别有独立的主正文页。任务表中的“专题正文主入口”表示第一次阅读应进入哪里,其他正文可以作为支撑阅读。

Task1 建立显存账本,Task2 比较单机训练策略;完成前置机制后,Task3 可以沿项目链执行训练侧验证:先由 73 建立 baseline,再由 76 比较策略、75 进行预算决策,最后由 74 用 trace 解释结果。需要补充时间测量、显存对象或异常排查时,分别回补 Part 00 · 17、18、19;它们是按需支撑,不要求全部完成后才能进入项目链。Task4–5 分别处理推理缓存和量化分支;Task6 提供分布式扩展,Profiling 作为跨分支的证据方法。Part 02 · 61 架构验证是架构扩展项目,Part 02 · 71 MLA 与 KV Cache 结构基准属于推理显存分支;Part 02 · 08 架构技巧、Part 02 · 06 MoE 路由器、Part 02 · 07 MoE 负载均衡损失和 LoRA / QLoRA 也不属于共同前置。

证据边界与项目出口

路线表负责选择入口,正文负责解释机制;需要按现象分流时进入显存优化判断手册,需要沿“发现问题 → 建账本 → 做实验 → 下结论”连续阅读时进入显存优化深入阅读

内容层级CPU 可以确认GPU、backend 或多卡才能确认主要验证出口
Task0–1 机制与账本生命周期、shape、dtype、参数、梯度、optimizer state 和 activation 的理论关系实际峰值、allocator reserved、带宽和 OOM 边界01 显存账本
Task2 单机策略accumulation、checkpoint、offload 的逻辑和梯度对齐显存节省、重算 / 搬运代价、吞吐和 OOM03 检查点与卸载
Task3 训练项目workload、指标、报告和预算决策逻辑73 baseline、76 策略比较、75 预算敏感性、74 trace 解释73 → 76 → 75 → 74
Task4–5 推理与量化KV Cache shape、容量估算、量化误差和决策逻辑backend 命中、TTFT / TPOT、格式、kernel、真实显存和质量04、05 与 66–71
Task6 分布式扩展ZeRO、pipeline、tensor、expert parallel 的切分模拟多卡显存分摊、通信时间、拓扑影响和 profiler 归因07 与 79–81

CPU 运行可以使用 GPU 机器,但 device='cpu' 的结果仍属于 CPU 证据。73、76、75 使用匹配的模型、dtype、batch、seq_len 和 workload;74 是跨项目的 profiling 收口,不把不同条件下的数字直接横向比较。FP32 长序列 OOM 是容量边界,应与 BF16、LoRA / QLoRA 或 activation-only workload 分开记录。

同一 Notebook 在不同路线中只切换观察目标:显存路线关注对象账本、峰值和容量,推理路线关注 KV Cache、TTFT / TPOT 和并发,算子与编译路线关注 kernel、访存和融合,训练微调路线关注 loss、梯度和稳定性。Notebook 保留一份权威实现;不同模型、设备、dtype 和 workload 的结果不能直接合并。不要把“代码运行成功”写成“显存优化成功”:稳定决策至少需要固定 workload、baseline / candidate、质量门槛和报告文件。

环境与验证

基础机制可以 CPU-first;真实训练、显存峰值和策略对比需要 NVIDIA GPU。运行前确认 PyTorch CUDA 可用,并按 Notebook 输出保存 JSON。73–76 的运行顺序、GPU 检查、结果文件和 74 profiling 要求见项目验证清单。如果问题首先表现为请求链路速度、低比特压缩、profiler 证据或多卡通信,分别转到推理优化量化与压缩性能分析通信与并行

Released under the MIT License.