读完这篇，你能做什么
能列出预算假设、核验 Adam 状态占用的字节数，并在内存占用随微批次大小单调变化的模型下，计算能容纳的最大微批次。
学习理由、准备要求与知识点
为什么学它
参数能装进设备，不代表训练也能运行。分清各部分占用、单位和峰值，才能决定是减小批次、分片还是调整训练目标。
先读这些内容
知识点清单
- 参数量与共享参数去重
- 字节与
GiB单位 - 权重梯度和优化器状态
- 激活与临时峰值
- 6ND计算近似
- 吞吐与利用率
- 算力最优的数据参数分配
- 预算模型二分搜索
运行前先确认Java17、PyTorch2.9.1 CPU。字节估算针对无AMSGrad、全参数FP32 AdamW；CPU张量统计不是GPU峰值测量，6ND仅为稠密模型粗估。
01先弄懂它是什么
估算训练需要多少资源，先把“多少参数”“占多少内存”“算多少次”分开。这里 N 是可训练参数个数，D 是整个训练实际处理的词元数，C 是累计浮点运算次数，B 是一次前向计算的微批大小，T 是序列长度。内存用字节计，FLOPs 表示运算总数，FLOP/s 表示每秒运算速度，单位不同不能直接相加。预算表还要注明：GB 是十亿字节，GiB 是二的三十次方字节。
一个参数不只占权重本身的空间。以普通全参数 FP32 AdamW 为例，粗略计算时，每个参数的权重占四字节、梯度占四字节、两个动量状态合计八字节，主体一共十六字节。这里没有计算 AMSGrad 的最大二阶矩，也没有额外主权重副本。即使用混合精度，也不能把总数直接减半，还要逐项检查实际数据类型、主权重、分片方式，以及哪些参数被冻结。
激活是前向中为反向保存的中间张量，依赖批次、长度、层数和内核。传统注意力分数形状含T乘T，长度翻倍可能让这一部分增长四倍；高效注意力可改变中间存储，却没有取消所有长度成本。模型加载成功只证明权重分配成功，首次反向和优化器更新往往才出现新峰值。
文中这几个词，先认一认
- 静态状态
- 随参数数量增长的权重、梯度和优化器张量。它们也可能在训练到某个阶段时才创建，并非一开始就全部占满。
- 峰值内存
- 运行过程中同一时刻需要占用的最大内存，包含当时仍在使用的对象和运行时缓冲；它不是所有文件大小相加。
- Scaling Law
- 在给定数据、模型和训练范围内，通过实验总结规模与损失之间的关系。它不能无条件推广到任何训练设置。
02看看它怎么工作
对以大型矩阵乘法为主的稠密语言模型，常用C约等于6ND进行量级估算。它把一次乘加计作两次操作，并用前向与反向的相对成本作近似；词表头、注意力随长度的开销、重计算与稀疏路由都可能改变结果。估算耗时还要用实测有效吞吐，而非直接除以设备宣传峰值。CPU小模型只能验证算术，不给出GPU训练时长承诺。
算力预算固定时，增加参数会减少可处理数据量。Chinchilla研究通过多组训练比较这种取舍，说明只增大模型而不调整数据量可能偏离其条件下的有效配置。具体最优比例依赖数据质量、重复、结构、推理成本和日程；不能把一个经验比例当成所有中文领域模型的强制处方，更不能从一次玩具实验重新证明规模规律。
实验先构建一个线性层，核对参数总数。反向之后梯度才出现，第一次AdamW更新之后动量才出现；逐张量相乘numel与element_size核验字节。Java算法采用独立的预算模型M(B)=固定状态+每样本激活乘B，用二分寻找最大可行B。模型中的每样本激活值需要由目标实现测量或估计，搜索算法正确不意味着估计系数准确。
先跑通一个最小例子
运行 experiment.py，在CPU统计一个线性层首次AdamW更新前后的张量字节。没有GPU峰值或大模型训练测量。
import torch
torch.set_num_threads(1)
torch.manual_seed(27)
model = torch.nn.Linear(3, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.01, foreach=False)
params = list(model.parameters())
n = sum(p.numel() for p in params)
size = lambda tensors: sum(t.numel() * t.element_size() for t in tensors)
assert n == 8
assert size(params) == 4*n
# 状态延迟初始化，不能在首次step之前估计完整优化器内存
assert len(optimizer.state) == 0
x, target = torch.randn(4, 3), torch.randn(4, 2)
loss = (model(x) - target).square().mean()
loss.backward()
assert size([p.grad for p in params]) == 4*n
optimizer.step()
moments = [optimizer.state[p][key] for p in params
for key in ("exp_avg", "exp_avg_sq")]
assert size(moments) == 8*n
assert size(params) + size([p.grad for p in params]) + size(moments) == 16*n
# step等标量与容器开销另计，16N不是进程内存峰值
assert all("step" in optimizer.state[p] for p in params)
n_big, tokens = 10**9, 20*10**9
assert 6*n_big*tokens == 120_000_000_000_000_000_000
print("8 compute and tensor-state byte checks passed")
核对运行结果参数8个，权重32字节、梯度32字节、两个动量64字节；主体128字节，另有step等对象。
正在加载流程图…
查看流程图源码（Mermaid）
flowchart TD
A["参数与dtype清单"] --> B["状态字节预算"]
C["批次长度和内核"] --> D["激活与临时预算"]
B --> E["加安全余量的容量约束"]
D --> E
E --> F["最大微批次估计"]
G["参数量和训练Token"] --> H["计算量粗估"]
H --> I["实测吞吐校准时间"]按顺序看，每一步做了什么
- 明确N、D、B、T与每个张量的dtype。
- 分别列出权重、梯度、优化器、激活和临时缓冲。
- 把内存模型转换为单调可行性判定并搜索批次。
- 以真实实现的阶段测量校准，再报告假设与误差。
03对照实现代码
PyTorch v2.9.1 · Adam._init_group；Adam.step；state['exp_avg']；state['exp_avg_sq']
下例为原创状态预算模型，省略真实优化器更新、AMSGrad、foreach及fused分支；不是完整Adam实现。
def state_budget(parameter_count, weight_bytes=4, grad_bytes=4, moment_bytes=4):
# 数量与单元素字节分开，才能对应真实dtype
if parameter_count < 0 or min(weight_bytes, grad_bytes, moment_bytes) <= 0:
raise ValueError("invalid budget")
# 普通Adam有一阶和二阶动量，主权重或AMSGrad需额外列项
moments = parameter_count * 2 * moment_bytes
# 只返回三类主体，不包含激活、标量和运行时缓冲
return parameter_count * (weight_bytes + grad_bytes) + moments
Adam._init_group遍历具有梯度的参数，在空状态时初始化step和动量，因此优化器构造后立即统计会低估。
一阶与二阶状态使用对应参数形状；默认zeros_like路径说明应检查真实参数dtype，不能仅凭开启autocast推断状态字节。
AMSGrad会增加max_exp_avg_sq，特定实现还可能分配额外临时量；教学十六字节口径有明确开关前提。
AdamW在该版本复用Adam机制并启用解耦权重衰减，状态预算与权重衰减形式是不同问题。
代码先调用谁，再调用谁
AdamW构造基础Adam状态配置- step收集梯度并调用
_init_group - 首次更新初始化动量后执行更新内核
04动手写一个例子
单调预算下搜索最大微批次
输入非负固定字节、正的每样本激活字节、总预算和候选上限，求固定+批次乘单样本不超过预算的最大批次。固定成本已超预算返回-1。
- 输入 1
fixed=100, per=30, budget=200, limit=10- 输出 1
3
内存模型随批次单调增加，所以一旦某个批次超预算，更大批次也必定不可行。
二分维护可行下界与待判定上界，使用偏向上侧的中点防止只有两个候选时停滞。
乘法与总量使用BigInteger，避免模型规模过大造成long回绕后把超预算误判为可行。
为什么这样写是对的固定成本可行时，lo始终可行，真正最大可行批次位于闭区间[lo,hi]内。
U为批次上限；k为输入与中间乘积位数。大整数运算不是严格常数，本例乘以long范围批次，操作成本随数字位数增长。
保存为 Main.java，运行 javac --release 17 -encoding UTF-8 Main.java，再运行 java Main。程序包含固定输入自测，断言失败时会抛出错误。
展开收起代码 · Main.java · 完整可运行程序
长行可在代码区域内左右滚动，按 Tab 可把焦点移入代码区。
import java.math.BigInteger;
public class Main {
static BigInteger bi(long x){return BigInteger.valueOf(x);}
static long maxBatch(BigInteger fixed, BigInteger per, BigInteger budget, long limit){
if(fixed.signum()<0 || per.signum()<=0 || budget.signum()<0 || limit<0)
throw new IllegalArgumentException();
// 先区分连零批次都放不下的情况
if(fixed.compareTo(budget)>0)return -1;
long lo=0,hi=limit;
while(lo<hi){
long distance=hi-lo;
// 上中点写法避免lo+hi溢出
long mid=lo+distance/2+distance%2;
BigInteger bytes=fixed.add(per.multiply(bi(mid)));
if(bytes.compareTo(budget)<=0)lo=mid;else hi=mid-1;
}
return lo;
}
static BigInteger compute(BigInteger n, BigInteger tokens){
// 保留整数精度，只表示6ND近似中的算术部分
if(n.signum()<0 || tokens.signum()<0)throw new IllegalArgumentException();
return n.multiply(tokens).multiply(bi(6));
}
static void check(boolean ok){if(!ok)throw new AssertionError();}
public static void main(String[] args){
check(maxBatch(bi(100),bi(30),bi(200),10)==3);
check(maxBatch(bi(100),bi(30),bi(100),10)==0);
check(maxBatch(bi(100),bi(30),bi(99),10)==-1);
check(maxBatch(bi(0),bi(1),bi(Long.MAX_VALUE),Long.MAX_VALUE)==Long.MAX_VALUE);
check(maxBatch(bi(0),bi(1),bi(1000),4)==4);
check(compute(bi(1_000_000_000L),bi(20_000_000_000L))
.equals(new BigInteger("120000000000000000000")));
try{maxBatch(bi(0),bi(0),bi(2),4);throw new AssertionError();}
catch(IllegalArgumentException expected){}
System.out.println("7 memory budget checks passed");
}
}
这些特殊情况也要试一试
- 固定成本大于预算返回-1，与仅允许零批次区分。
- 上限为零和
Long.MAX_VALUE都不使中点溢出。 - 每样本成本必须为正；真实非单调内核选择不符合本模型。
进阶选读 · 按序列长度建立分段预算模型
场景不同长度选择不同内核，单个线性系数无法覆盖整个输入域。
做法对固定长度区间测量若干批次，加入安全余量，逐区间建立单调上界；输出预测与实际峰值误差，超出校准范围时重新测量。
复杂度每个区间O(log U)次预测；实际测量成本由模型前后向决定。
好处与代价保守上界减少容量失败但降低利用率，不能把预算搜索结果当作设备实测。
05项目里怎么用
预算报告应附训练模式、序列长度、优化器选项和参数冻结范围。共享Embedding与输出层可能引用同一权重，参数统计应避免重复计数；序列并行、分片和量化又改变各项持有方式。先把最简单可运行配置测准，再增加策略，才能知道每一项节省发生在哪里。
测性能时，先经过预热，再分别记录每轮计算、等待数据、设备通信和保存训练进度各花多久。每秒处理多少样本，与每秒处理多少有效 token 不是同一个口径；设备看起来很忙，也可能在处理大量填充或其他不计入目标的工作。少存一些中间值、到反向传播时再重算，会增加运算次数，却可能降低内存峰值，让一次能处理更大的微批。最后应比较完成同样有效训练目标需要多少时间和预算，不能只挑利用率或某一项吞吐数字。
规模规律适合提出实验假设。先用多个参数量与数据量做小规模对照，保持日程和评测一致，报告拟合范围与外推不确定性。模型最终用于高频推理时，可以接受更长训练来降低部署模型大小；只考虑训练算力的最优点不必等于产品全生命周期最优点。
容易踩的坑
- 用下载权重文件大小预测训练显存。
- 开启混合精度后直接把十六字节减半。
- 第一次
optimizer.step之前读取空状态并宣布占用很小。 - 用理论设备峰值计算精确训练完成时间。
几种做法，怎么选
表格有四列，窄屏可左右滑动查看代价和适用条件。
| 方案 | 优势 | 代价与限制 | 适用条件 |
|---|---|---|---|
| 减小微批次并累积 | 降低部分激活峰值 | 增加调用次数且状态内存仍在 | 激活主导的容量问题 |
| 减少参数或冻结部分层 | 降低可训练状态 | 改变优化空间与能力假设 | 状态主导的适配任务 |
| 分片和重计算 | 在容量限制内运行更大配置 | 引入通信或额外计算 | 已测量瓶颈的大规模训练 |
06检验是否读懂了
先自己想、自己写，再展开答案核对。场景题是假设的练习情境；遇到还没学过的内容，可以做完基础后再回来。
先用自己的话解释
01十亿FP32参数权重占多少字节？
展开参考答案与推导
四十亿字节，约3.73GiB；这只包括权重。
02为什么第一次更新可能才发生OOM？
展开参考答案与推导
优化器延迟创建动量，还可能出现更新临时张量，加载和前向阶段未覆盖这些峰值。
036ND能否直接用于任意MoE模型总参数？
展开参考答案与推导
不能；激活专家、路由和共享模块决定实际运算，需用适配的计数模型。
换个情况，还会处理吗
01固定状态已经超过设备预算，减小B是否足够？
展开参考答案与推导
不够；需要减少或分片状态、冻结参数等改变固定项的办法，单纯减少激活无解。
02长序列从2K到4K占用远超两倍怎么办？
展开参考答案与推导
检查注意力二次中间项、内核切换和缓存，分别测量各阶段，不能只用线性外推。
03小实验拟合比例能直接推到百倍模型吗？
展开参考答案与推导
应当报告外推不确定性，增加中间尺度验证，并考虑架构、数据与优化日程变化。
写代码试一试
01算普通FP32 Adam主体。
- 输入 / 条件
N=1000，无AMSGrad- 输出 / 预期
16000字节
提示 权重、梯度、两个动量。
展开参考答案与推导
每个参数四项各四字节，总量16N。
02求预算最大批次。
- 输入 / 条件
固定100，每样本30，预算200- 输出 / 预期
3
提示 检验3与4两侧。
展开参考答案与推导
批次3用190，批次4用220，因此最大3。
03计算全训练词元数。
- 输入 / 条件
每步有效2048词元，1000步，无变化- 输出 / 预期
2048000
提示 明确是累计处理量。
展开参考答案与推导
有效词元每步乘步数；如果重复读同一文档，这仍是处理量，不是独立语料大小。
读完之后，独立完成
课后作业：自己做一遍
本篇共 12 道作业，先完成 9 道基础与综合任务，再按前置知识挑战进阶题。其中包含 2 道思维拓展，试着换个条件，看看原来的办法是否还管用。不附参考答案。
能列出预算假设、核验 Adam 状态占用的字节数，并在内存占用随微批次大小单调变化的模型下，计算能容纳的最大微批次。
本篇重点加练
内存峰值、计算近似与预算分配互相制约，需要避免用单一参数量代替整体成本。
- 激活与临时峰值
- 6ND计算近似
- 算力最优的数据参数分配
重点内容会多练几次。先说清楚怎么做，再试试空数据、出错或条件改变时会发生什么。
先自己做，卡住时记下问题，再回来查。问答题写清“为什么、什么时候适用”；动手题留下输入、输出和运行步骤，方便下次重做。
基础必做
- 作业 01代码推演
模型由两个共享同一权重的线性层组成，权重形状 [128,64]。列出参数对象、引用次数及不同状态精度下的字节预算。
- 提交内容
- 去重参数表与预算
- 怎样算完成
- 共享对象不重复计数
- 字节与
GiB标注一致
练到的知识点参数量与共享参数去重 · 字节与GiB单位 · 权重梯度和优化器状态
- 作业 02原理问答
画出一次前向、反向和优化更新中权重、激活、梯度及优化器状态的存活区间，指出静态状态表无法描述哪些峰值。
- 提交内容
- 内存生命周期图
- 怎样算完成
- 状态类型完整
- 稳态与临时峰值分开
练到的知识点权重梯度和优化器状态 · 激活与临时峰值
- 作业 03原理问答
两个实现参数量相同，一个保存全部中间激活，一个重计算部分激活。列出计算近似与峰值预算中会改变和不会改变的项，并声明模型假设。
- 提交内容
- 成本项变化表
- 怎样算完成
- 计算和存储口径分开
- 重计算范围明确
练到的知识点激活与临时峰值 · 6ND计算近似
- 作业 04代码推演
给三个候选参数量和总计算预算，自行定义采用的近似公式及每种配置可训练 token 数，审阅“最大模型必定最佳”的选择依据。
- 提交内容
- 候选预算表及缺失证据
- 怎样算完成
- 近似适用条件明确
- 可训练量不直接等同质量
练到的知识点6ND计算近似 · 算力最优的数据参数分配
综合应用
- 作业 05编程实践
编写参数化内存预算求解器，输入固定开销、批次相关开销及预算上限，输出可行批次范围；测试无可行解、临界等号和计数溢出。
- 提交内容
- CPU 程序与边界测试
- 怎样算完成
- 预算模型明确
- 整数与舍入策略可复核
练到的知识点激活与临时峰值 · 预算模型二分搜索
- 作业 06面试场景
训练计划采用 N=5×10^7、D=2×10^8 的近似预算。审阅计算量、理论耗时与实测吞吐之间的推断链，列出必须测量的条件。
- 提交内容
- 预算审查与测量方案
- 怎样算完成
- 近似假设明确
- 理论峰值不直接当实测速度
练到的知识点6ND计算近似 · 吞吐与利用率
- 作业 07排错实验
在 CPU 小型网络上改变批大小和序列长度，记录时间及可获取的进程内存指标；明确这些指标与 GPU 峰值显存的差别。
- 提交内容
- 实验表与测量局限
- 怎样算完成
- 测量区间一致
- 设备外推受到限制
练到的知识点激活与临时峰值 · 吞吐与利用率
- 作业 08编程实践
编写 CPU 内存生命周期账本，输入不超过十个对象的分配/释放事件及字节数，计算随时间占用；测试共享引用和同一时刻事件顺序。
- 提交内容
- 模拟器、事件契约与曲线
- 怎样算完成
- 模拟对象身份唯一
- 逻辑载荷不冒充进程实测
练到的知识点激活与临时峰值
- 作业 09设计评审
预算同时受内存和截止时间限制，某配置单步更快但允许批量更小。自拟三配置测量表，设计可行性与质量比较流程。
- 提交内容
- 双约束预算与比较规则
- 怎样算完成
- 两个约束都被检查
- 质量指标在独立数据上定义
练到的知识点激活与临时峰值 · 算力最优的数据参数分配
进阶挑战 · 掌握前置后选做
- 作业 10设计评审
为固定训练预算设计三组模型参数量与数据量组合，预先规定质量评价和预算容差；说明小规模实验如何支持或限制后续扩展。
- 提交内容
- 组合预算与研究协议
- 怎样算完成
- 参数和数据预算共同计算
- 不把近似比例当普适定律
练到的知识点6ND计算近似 · 算力最优的数据参数分配 · 预算模型二分搜索
思维拓展
先写出你的猜想，再用例子或实验验证。答案可以不同，但要说清依据、适用情况，以及还有什么没弄懂。
- 作业 11思维拓展 · 条件变化
训练数据含大量重复时，固定 6ND 近似预算还能保证“有效训练信息量”相同吗？设计不超过一千样本的小实验区分计算消费与学习收益。
- 提交内容
- 重复率实验与信息边界说明
- 怎样算完成
- 预算记录真实样本访问
- 有效信息不随意定义成唯一标量
练到的知识点6ND计算近似 · 算力最优的数据参数分配
- 作业 12思维拓展 · 跨域迁移
将训练容量规划迁移到流水线生产：平均占用低但峰值超过容量的情况如何出现？构造三阶段生命周期例子，比较总工作量与最大同时占用。
- 提交内容
- 生命周期图及容量反例
- 怎样算完成
- 阶段重叠关系明确
- 工作总量和峰值维度不同
练到的知识点激活与临时峰值 · 6ND计算近似
完成后，再检验一次
给每题标记“独立完成 / 查阅后完成 / 尚未完成”。查过资料或没做出的题，记下卡在哪里。隔一天换一组输入再试，能解释结果和出错原因，再把这个知识点标为掌握。
继续阅读
先看能帮你完成当前练习的部分。例子跑通后，再读更深入的解释。
- 必读参考文档英文
PyTorch2.9.1 Adam 状态源码 （新标签页）↗
确认延迟初始化和动量形状。
- 深入参考文档英文
Chinchilla 原论文 （新标签页）↗
阅读算力约束下的数据参数取舍与附录F运算计数。
- 深入参考文档英文
Scaling Laws 原论文 （新标签页）↗
理解6N每训练Token近似的定义和适用模型。
- 深入参考文档英文
PyTorch2.9.1 AdamW 源码 （新标签页）↗
核对
AdamW与Adam基础实现的关系。