【新智元导读】DeepMind 最近被 ICML 2024 接收的一篇论文，完完全全暴露了他们背靠谷歌的「豪横」。一篇文章预估了这项研究所需的算力和成本，大概是 Llama 3 预训练的 15%，耗费资金可达 12.9M 美元。
发一篇顶会论文，需要多少实验预算？
最近，DeepMind 发表了一项研究，对 LLM 扩大规模时各种算法和架构细节，比如参数和优化器的选择，进行了广泛的实证调查。
这篇论文已被 ICML 2024 接收。
论文地址：https://arxiv.org/abs/2407.05872
63 页的论文涵盖了数以万计的模型，备选方案包括 3 种优化器、4 种参数化方案、几种对齐假设、十多个学习率，以及最高达 26.8B 的 14 种参数规模。
需要进行实验的 4 种参数化方案
仅仅听到这些数字，就不难知道，这项研究必定涉及海量的模型运行实验。
而有一位忠实读者，为了测试自己对论文内容的理解，统计了其中进行的所有实验，并估算出了复现论文的成本。
将所需算力全部加在一起，林林总总，居然达到了惊人的 1290 万美元。
考验基本功的时刻到了，假如你是研究团队的 leader，根据实验计划对所需算力和成本进行预估是一项必不可少的技能。
那就让我们跟着这篇博客文章盘一遍，这一千多万美元，究竟烧在哪里。
Transformer 架构信息
论文附录 C 提供了关于模型算法和架构的各种细节设置，比如使用 decoder-only 架构、层归一化、GeLU 激活函数、无 dropout、T5 分词器、批大小为 256、用 FSDP 并行等等。
实验模型的参数规模统计
通过架构方面的信息，我们可以大致估算出训练中每个 token 所需的 FLOPS，记为 M。
由于论文没有描述到任何 GQA / MQA 机制，所以就假设 Rkv=1，此外还有 lseq=512，Dhead=128，L=8（深度），V=32101（分词器词汇量）。
模型总参数量可以表示为：
因此，就可以得到 M 的计算公式：
默认情况下，每次实验处理的 token 数（tokens per experiment, TPE）为 5k（训练步数）256（批大小）512（lseq），约为 6.5536e9。
defM(d:int,L=8,l_seq=512,V=32101)->int:return6*d*(L*(12*d+l_seq)+V)TPE=50000*256*512
对齐实验
假设对齐实验中，直接使用了后面的学习率扫描得出的最优结果，并没有单独进行学习率扫描，因此这一步的成本计算比较简单：
defalignment()-int:return4*TPE*sum(M(d)fordin[1024,2048,4096])#>>f'{alignment()3E}'#'3.733E+20'#>>cost_of_run(alignment())[0]#888.81395400704
如果 H100 每运行 1 小时的花费以 3 美元计算，对齐实验的成本大致为 888 美元。
学习率
子问题：最佳评估损失（eval loss）实验
论文的表 E1 记录了 6 种模型规模下，所有可能的优化器 参数化方案 模型大小 实验设置的组合，分别进行基础学习率扫描，以获得最佳评估损失。
总共包括如下几个实验变量：
-
模型维度 D∈3072,4096,6144,8192,12288,16384
-
4 种参数化方案
-
3 种优化器，其中 SGD 仅有 5 个实验设置，Adam 和 Adam+Param Scaling 有 7 个实验设置
假设这里的实验都是单独进行，没有从其他地方复制结果，因此如果全部运行一遍，有成本上限预估：
H=[1,2,4,6,8,12,16,20,24,32,48,64,96,128]D=[h*128forhinH]deftable_e1()-int:sets_x_optims=5+7+7return4*sets_x_optims*TPE*sum(M(d)fordinD[-6])#>>f'{table_e1()3E}'cost_of_run(table_e1())#'1.634E+23'#(388955.999106498616206.499962770775)
这部分的成本就接近 40 万美元，虽然仍属于可接受范围内，但对于大多数学术预算来说，已经算是非常昂贵了。
表 E1 给出了最佳评估损失，但没有描述 LR 的扫描策略，每张图上的点数也不尽相同。
由于没有得到论文作者的答复，我们也无法确定具体机制，因此假设每个最佳评估损失都经过了 15 次实验（目测发现，每条线的点数约为 10～15）。
参数
根据论文 4.2 节内容，学习率还涉及到两个超参数的选择： 和 。
如果仅有 参数，则被称为「LR+default」设置：
这部分包括 3 优化器，4 参数化，加上全局和单层（GlobalLR、Perlayer-fullalign）分别进行实验，以及未知的 LR 扫描数量：
def_only()->int:return3*4*2*PpL*TPE*sum(M(d)fordinD)#7.988E+23(1902022.3291813303,79250.93038255542)
从公式就可以看出，成本和下文的 epsilon 实验类似，都是 200 万美元。
参数
相比 参数的实验，这部分有两个细节差异。
首先，除了 GlobalLR、Perlayer-fullalign 两种设置外，还需要加上 Perlayer-noalign 设置。
其次，仅针对 d=1024=b，进行 3D 超参数搜索 (_1,_h,_L+1)，因此有额外的 800 次运行。
两者结合后的计算公式为：
这部分的预估成本与 Adam 的 epsilon 热力图实验接近，约为 320 万美元。
defgamma_expts()->int:return36*TPE*(800*M(1024)+PpL*sum(M(d)fordinD))#gamma_expts1.354E+24(3224397.534237257,134349.8972598857)
Adam 优化器的 Epsilon 参数
论文 4.3 节所述的 Epsilon 参数实验是计算量的大头。
根据上面的推断，每次找到最佳评估损失时都尝试过 15 个不同的学习率（points per line），那么图 6 所示的 epsilon 参数变化图耗费的计算量为：
计算结果透露出一种简洁的昂贵，也就是 200 万美元的账单而已。
PpL=15#unprincipledestimatedefeps_variants()->int:return4*6*PpL*TPE*sum(M(d)fordinD)'''>>>f'{eps_variants():.3E}';cost_of_run(eps_variants())'7.988blog.byteway.netE+23'(1902022.3291813303,79250.93038255542)'''
除了图 6 左侧的折线图，还有附录 F 热力图的结果。
假设每个方块值都是经过 13 次学习率扫描后得到的结果，这部分计算量则为：
结果发现，仅仅要得到这 8 张热力图，成本就是 320 万美元。而且，由于我们将 LR 扫描数量建模为常数 13，这个数字可能低于实际成本。
defeps_heatmaps()-int:#eps-type*eps-val*parameterizations*LRrange*...return2*6*4*13*TPE*sum(M(d)fordinD[-6])'''>>f'{eps_heatmaps()3E}'cost_of_run(eps_heatmaps())'1.341E+24'(3193533.466348094133063.89443117057)'''
权重衰减
权重衰减实验（附录 G）比较好理解，对 4 参数化方案以及所有参数进行一次基本的 LR 扫描：
比 epsilon 实验便宜不少，也就是湾区工程师一年的工资 ——31.7 万美元。
defweight_decay()->int:return4*PpL*TPE*sum(M(d)fordinD)'''>>>f'{weight_decay():.3E}';cost_of_run(weight_Ligthing Newsdecay())'1.331E+23'(317003.7215302217,13208.488397092571)'''
Adafactor 优化器
这部分实验在附录 C3 中有详细描述，是为了检验 Adafactor 和 Adam+parameter scaling 是否有相似的宽度缩放机制。
共有 24 张图，其中每个优化器收集 11 个数据点，因此计算公式为：
账单上再加 18.8 万美元。
defadafactor()->int:return2*2*4*PpL*TPE*sum(M(d)fordinD[:11])'''>>>f'{adafactor():.3E}';cost_of_run(adafactor())'7.918E+22'(188532.80765144504,7855.533652143543)'''
计算最优化
论文尝试改变注意力头 H 的数量，希望找到计算最优化的设置，但其中涉及步长和数据集的改变，因此这部分不使用公式描述，计算代码如下：
defP(d:int,L=8,V=32101)->int:return2*d*(6*L*d+V)defcompute_optimal():indices_50k=(14,14,12)return4*PpL*sum([TPE*sum(sum(M(d)fordinD[:i])foriinindices_50k),20*sum(P(d)*M(d)fordinD[:11])*3,])#compute_optim7.518E+23(1790104.1799513847,74587.67416464102)
总结
将以上各部分实验的算力和成本汇总在一起：
alignment3.733E+20(888.81395400704,37.033914750293334)table_e11.634E+23(388955.9991064986,16206.499962770775)eps_variants7.988E+23(1902022.3291813303,79250.93038255542)eps_heatmaps1.341E+24(3193533.466348094,133063.89443117057)_onlLigthing Newsy7.988E+23(1902022.3291813303,79250.93038255542)gamma_expts1.354E+24(3224397.534237257,134349.8972598857)weight_decay1.331E+23(317003.7215302217,13208.488397092571)adafactor7.918E+22(188532.80765144504,7855.533652143543)compute_optim7.518E+23(1790104.1799513847,74587.67416464102)
结果发现，整篇论文的运算量为 5.42e24 FLOPS。
这个数字仅仅是 Llama 3 训练计算量的 15%，如果在 10 万卡 H100 集群上运行，只需要 2 天时间即可完成所有实验。
total_flops=5.421E+24rentalprice:US$12.9Mh100nodemonthsrequired:746.9595590938408(sanitycheck)D=[128,256,512,768,1024,1536,2048,2560,3072,4096,6144,8192,12288,16384](sanitycheck)modelsizes:['0.00979B','0.0227B','0.058B','0.106B','0.166B','0.325B','0.534B','0.794B','1.1B','1.87B','4.02B','6.97B','15.3B','26.8B'](sanitycheck)M/6P:['63.4%','68.5%','75.3%','79.7%','82.8%','86.8%','89.3%','91.0%','92.2%','93.9%','95.7%','96.7%','97.7%','98.3%']
然而，如果不从 LLM 预训练的标准来衡量，仅把 DeepMind 的这篇论文看做一篇学术研究，这个计算量就显得相当奢侈了。
如果实验室仅有 10 张 H100，就根本不可能进行这个量级的研究。
有 100 张 H100 的大型实验室，或许能用几年时间跑完以上所有实验。
参考资料：
-
https://152334h.github.io/blog/scaling-exponents/
-
https://news.ycombinator.com/item?id=41107721
-
https://arxiv.org/abs/2407.05872
本文来自微信公众号：微信公众号（ID：null），作者：新智元
广告声明：文内含有的对外跳转链接（包括不限于超链接、二维码、口令等形式），用于传递更多信息，节省甄选时间，结果仅供参考，IT之家所有文章均包含本声明。
Related suggestion: 下一个大事件： 新华健康｜北京市民线上购药可用支付宝进行医保结算
总结具体而言，参保市民在京东App、美团外卖App平台买药时，需先登录两家平台，选择带有“医保”蓝色标识的药品，提交订单后，选择“北京医保个人账户支付”即可跳转至支付宝App进行医保电子凭证的激活和认证，进而支付结算。 家住北京市朝阳区的张女士说：“以前家里老…