设置
  • 日夜间
    随系统
    浅色
    深色
  • 主题色

DeepMind 研究成本大起底,一篇 ICML 论文烧掉 1290 万美元

新智元 2024/8/3 13:18:50 责编:清源

【新智元导读】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。

def M(d: int, L=8, l_seq=512, V=32101) -> int:
    return 6*d * (L*(12*d + l_seq) + V)
TPE = 50000 * 256 * 512

对齐实验

假设对齐实验中,直接使用了后面的学习率扫描得出的最优结果,并没有单独进行学习率扫描,因此这一步的成本计算比较简单:

def alignment() - int:
    return 4 * TPE * sum(M(d) for d in [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 * 128 for h in H]
def table_e1() - int:
  sets_x_optims = 5 + 7 + 7
  return 4 * sets_x_optims * TPE * sum(M(d) for d in D[-6])
# >> f'{table_e1()3E}'cost_of_run(table_e1())
# '1.634E+23'
# (388955.9991064986 16206.499962770775)

这部分的成本就接近 40 万美元,虽然仍属于可接受范围内,但对于大多数学术预算来说,已经算是非常昂贵了。

表 E1 给出了最佳评估损失,但没有描述 LR 的扫描策略,每张图上的点数也不尽相同。

由于没有得到论文作者的答复,我们也无法确定具体机制,因此假设每个最佳评估损失都经过了 15 次实验(目测发现,每条线的点数约为 10~15)。

β 参数

根据论文 4.2 节内容,学习率还涉及到两个超参数的选择:β 和 γ。

如果仅有 β 参数,则被称为「LR+default」设置:

这部分包括 3× 优化器,4× 参数化,加上全局和单层(GlobalLR、Perlayer-fullalign)分别进行实验,以及未知的 LR 扫描数量:

def _only() -> int:
  return 3*4*2*PpL * TPE * sum(M(d) for d in D)
# 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 万美元。

def gamma_expts() -> int:
  return 36*TPE * (800*M(1024) + PpL*sum(M(d) for d in D))
# gamma_expts 1.354E+24 (3224397.534237257, 134349.8972598857)

Adam 优化器的 Epsilon 参数

论文 4.3 节所述的 Epsilon 参数实验是计算量的大头。

根据上面的推断,每次找到最佳评估损失时都尝试过 15 个不同的学习率(points per line),那么图 6 所示的 epsilon 参数变化图耗费的计算量为:

计算结果透露出一种简洁的昂贵,也就是 200 万美元的账单而已。

PpL = 15 # unprincipled estimate
def eps_variants() -> int:
  return 4 * 6 * PpL * TPE * sum(M(d) for d in D)
'''
>>> f'{eps_variants():.3E}';cost_of_run(eps_variants())
'7.988E+23'
(1902022.3291813303, 79250.93038255542)
'''

除了图 6 左侧的折线图,还有附录 F 热力图的结果。

假设每个方块值都是经过 13 次学习率扫描后得到的结果,这部分计算量则为:

结果发现,仅仅要得到这 8 张热力图,成本就是 320 万美元。而且,由于我们将 LR 扫描数量建模为常数 13,这个数字可能低于实际成本。

def eps_heatmaps() - int:
  # eps-type * eps-val * parameterizations * LR range * ...
  return 2 * 6 * 4 * 13 * TPE * sum(M(d) for d in D[-6])
'''
>> f'{eps_heatmaps()3E}'cost_of_run(eps_heatmaps())
'1.341E+24'
(3193533.466348094 133063.89443117057)
'''

权重衰减

权重衰减实验(附录 G)比较好理解,对 4× 参数化方案以及所有参数进行一次基本的 LR 扫描:

比 epsilon 实验便宜不少,也就是湾区工程师一年的工资 ——31.7 万美元。

def weight_decay() -> int:
  return 4 * PpL * TPE * sum(M(d) for d in D)
'''
>>> f'{weight_decay():.3E}'; cost_of_run(weight_decay())
'1.331E+23'
(317003.7215302217, 13208.488397092571)
'''

Adafactor 优化器

这部分实验在附录 C3 中有详细描述,是为了检验 Adafactor 和 Adam+parameter scaling 是否有相似的宽度缩放机制。

共有 2×4 张图,其中每个优化器收集 11 个数据点,因此计算公式为:

账单上再加 18.8 万美元。

def adafactor() -> int:
  return 2*2*4*PpL*TPE*sum(M(d) for d in D[:11])
'''
>>> f'{adafactor():.3E}'; cost_of_run(adafactor())
'7.918E+22'
(188532.80765144504, 7855.533652143543)
'''

计算最优化

论文尝试改变注意力头 H 的数量,希望找到计算最优化的设置,但其中涉及步长和数据集的改变,因此这部分不使用公式描述,计算代码如下:

def P(d: int, L=8, V=32101) -> int:
    return 2 * d * (6*L*d + V)

def compute_optimal():
  indices_50k = (14, 14, 12)
  return 4*PpL*sum([
    TPE * sum(sum( M(d) for d in D[:i] ) for i in indices_50k),
        20  * sum(P(d)*M(d) for d in D[:11]) *3,
  ])
# compute_optim 7.518E+23 (1790104.1799513847, 74587.67416464102)

总结

将以上各部分实验的算力和成本汇总在一起:

alignment       3.733E+20 (888.81395400704, 37.033914750293334)
table_e1        1.634E+23 (388955.9991064986, 16206.499962770775)
eps_variants    7.988E+23 (1902022.3291813303, 79250.93038255542)
eps_heatmaps    1.341E+24 (3193533.466348094, 133063.89443117057)
_only       7.988E+23 (1902022.3291813303, 79250.93038255542)
gamma_expts     1.354E+24 (3224397.534237257, 134349.8972598857)
weight_decay    1.331E+23 (317003.7215302217, 13208.488397092571)
adafactor       7.918E+22 (188532.80765144504, 7855.533652143543)
compute_optim   7.518E+23 (1790104.1799513847, 74587.67416464102)

结果发现,整篇论文的运算量为 5.42e24 FLOPS。

这个数字仅仅是 Llama 3 训练计算量的 15%,如果在 10 万卡 H100 集群上运行,只需要 2 天时间即可完成所有实验。

total_flops=5.421E+24
rental price: US$12.9M
h100 node months required: 746.9595590938408

(sanity check) D=[128, 256, 512, 768, 1024, 1536, 2048, 2560, 3072, 4096, 6144, 8192, 12288, 16384]
(sanity check) model sizes: ['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']
(sanity check) 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之家所有文章均包含本声明。

相关文章

关键词:deepmind谷歌

软媒旗下网站: IT之家 最会买 - 返利返现优惠券 iPhone之家 Win7之家 Win10之家 Win11之家

软媒旗下软件: 软媒手机APP应用 魔方 最会买 要知