显存优化(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 一页: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 的逻辑和梯度对齐 | 显存节省、重算 / 搬运代价、吞吐和 OOM | 03 检查点与卸载 |
| 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 证据或多卡通信,分别转到推理优化、量化与压缩、性能分析或通信与并行。
