LLM - 训练与推理过程中的 GPU 算力评估

其中浮点运算是指涉及小数点的计算，例如加减乘除。TFLOPs 比 FLOPs 单位更大，其用于衡量计算机或处理器的更高性能。以 L40S 为例，在精度为 FP32 的情况下，其算力可达 91.6 TFLOPs，而 A800 则只有 19.5 TFLOPs，RTX 3060 大约是 12.5 TFLOPS。如果是超级计算机其 TFLOPS 性能指标可达 1000 万亿次每秒。

三.训练阶段的 GPU 消耗

◆ 影响训练的因素

计算前先看下训练阶段有哪些因素会影响训练的进程。

- 训练数据大小

- 模型参数规模

- 训练 Epoch 次数

- 显卡算力

◆ GPT-3 训练统计

上面我们介绍了 FLOPs 和 TFLOPs，这里 ZettaFLOPs 代表每秒十万京 (10^21) 次的浮点运算，不同指标之间主要是单位的差异。以 GPT-3 为参照，其为 1750 亿参数的模型，45 TB 训练数据，训练一次大约需要 175 ZettaFLOPs 即 1.75 x 10^23 次浮点数运算。由于这里我们没有获取到 175 ZFLOPs 是一个 Epoch 还是完整训练，所以后续的计算我们先统一按完整训练所需为参照。

◆ 自定义训练 GPU 评估

这里假定我们训练的模型是 LLaMA-33B，参数量为 330 以，训练数据量为 50G，GPU 算力以 A800 下 FP32 计算即 19.5 TFLOPs，则训练该模型需要 GPU 算力为:

NeedFLOPs = (330 / 1750) * (0.048828125 / 45) * 1.75 * 10^23 FLOPs

我们使用设备为 A800，则其可提供的算力为:

CalcByA800 = 19.5 * 10^12 FLOPs

最后结合我们需要的天数，例如训练 5 天结束，则 5 天的 s 数为:

TrainTime = 86400 * 5

最终所需 GPU 数量为:

GPUCount = NeedFLOPs / (CalcByA800 * TrainTime)

为了方便计算我们直接改写为 Python 代码:

#!/usr/bin/python
# -*- coding: UTF-8 -*-

"""
    Returns the estimated number of GPUs

    Parameters
    ----------
    args_num : 当前模型参数数量，单位为 亿
    data_size : 当前训练数据大小，单位为 GB

    Returns
    -------
    count : 所需 GPU 数量, 未向上取整
"""


def calc_gpu_num(_args_num, _data_size, _train_days):
    need_flops = (_args_num / 1750) * (_data_size / 45) * 1.75 * 10 ** 23
    calc_by_a800 = 19.5 * 10 ** 12
    train_time = 86400 * _train_days
    gpu_count = need_flops / (calc_by_a800 * train_time)
    return gpu_count


if __name__ == '__main__':
    args_num = 330 # LLaMA-33B
    data_size = 0.048828125 # 50G
    train_days = 5 # 训练5天
    count = calc_gpu_num(args_num, data_size, train_days)
    print(count)

计算得到 count = 4.250628165558721，我们向上取整可得 5 天训练 50G 数据的 LLaMA-33B 模型需要 5 张 A800，不过这里并未考虑显存或者性能损耗等问题，只是粗略估算。

四.推理阶段的 GPU 消耗

◆ 影响推理的因素

- 输入输出数据

- 模型参数规模

- 显卡算力

推理阶段需要计算输入与输出的文本之和，实际计算中，需要将文本 tokenizer 到 token_ids 以便 transformers 计算，中文字符与 token 比例大概为 1:2 即 1 个中文字符对应 2 个 token_ids。GPU 推理计算时，其与输入输出文本和 L、模型维度 D 和模型层数 N 成正比。

◆ 自定义推理 GPU 评估

假设输入的 query 文本数为 100，输出的文本书为 1000，则 L = (100 + 1000) * 2 = 2200，按模型维度 D = 1280，层数 N = 96，则计算所需算力:

NeedFLOPs ≈ L * D * N ≈ 270336000 = 2.7 * 10^8

假设我们使用 A800 用 1s 时间完成本次推理请求:

CalcByA800 = 19.5 * 10^12 FLOPs

则需要 GPU 数量:

count = 270336000 / (19.5 * 10**12) = 1.3863384615384615e-05

反过来 CalcByA800 / NeedFLOPs 可以计算一张 A800 可以满足约 7.2 w 用户 1s 内获取回复，不过这只是理想情况：

people = (19.5 * 10**12) / 270336000 = 72132.45738636363

◆ 计算差异

推理时涉及到文字和计算，例如生成连贯的剧情、生成数学逻辑的推理。这些处理对算力的消耗几乎相同，这是 LLM 的本质是语言模型，只要在输入输出相同的情况下，其处理的都是 token_ids，推理得到的也都是 next_token 的分布，还是 token_ids，所以与具体的表现形式无关，即不同类型的任务，只要输入输出相近，则其对算力的消耗也相近。

五.基于 Token 的计算方式

上面的方式都是基于 FLOPs 为基准计算的，但是由于实际场景下很多信息我们获取的并不完整。例如 GPT-3 的训练 FLOPs 是否准确，训练的 FLOPs 是一个 epoch 还是一整个训练流程，因此实际场景下还可以采用另一种方式，即基于 Token 处理效率的计算，但是这种计算方式有一个前提，就是需要你事先测算出对应 GPU 设备每秒处理的 token 数量，即 Process Tokens / s。

以博主实际测算为例，P40 即 4090 的 token 处理能力约为 25 Token / s，假设我们一天需要请求 1000 次，每次输入输出共 1000 个汉字，共花费 10 h 处理完成，则需要 P40 的数量可通过下述公式计算：

def calc_gpu_num_by_token(post_token_num, post_time_all, token_process):
    return post_token_num / (post_time_all * token_process)


if __name__ == '__main__':
    token_num = 1000 * 1000 * 2  # 1000 次请求、1000 个汉字对应 2000 个 token
    time_cost = 10 * 3600  # 总处理时间
    token_process_num = 25  # 每s可处理 token 数目
    print(calc_gpu_num_by_token(token_num, time_cost, token_process_num))

算下来 10h 处理完这些需求需要 2.223 台 P40，则向上取整得到需要 3 台 P40。

六.总结

这里分享了几种 GPU 算力计算的方法，上面的计算都是基于理想状态，实际情况下还需要考虑多机多卡的 IO 延迟，机器的网络延迟，以及模型大小与 GPU 显存的供需关系等等。更多的数据还是需要自己实践出真知，有问题欢迎大家在评论区一直讨论，因为现在网上这方面的信息还是相对较少，可以一起分享一起进步！