AI 文摘

大模型训练需要花费多长时间:FLOPs的简单计算方法及calflop开源实现





作者: 老刘说NLP 来源: 老刘说NLP

​​​​​

今天是2023年8月28日,星期一,新周的第一天,北京,晴。‍‍‍

我们继续来看看一些大模型研发过程中的话题,在前面的文章中我们介绍了大模型部署推理以及训练过程中的一些显存占用计算方案,例如,现有的一下显卡配置:

不同精度下的字节耗费对应情况:‍‍‍

不同精度下每10亿参数需要的对应显存占用情况:‍‍‍‍‍‍

而进一步的,我们来看看在时间维度上的一些花费情况,本文主要介绍计算FLOPs的方法,从而可以推算出模型的训练时间,供大家一起参考。

一、FLOPs作用

根据知乎专栏——简单的算法笔记中的文章训练模型算力的单位:FLOPs、FLOPS、Macs 与 估算模型(FC, CNN, LSTM, Transformers&&LLM)的FLOPs中对FLOPs的解释:FLOPs(Floating Point Operations)本身指浮点运算次数,可以简单理解为评估计算量的单位,即FLOPs可以用来衡量一个模型/算法的总体复杂度(即所需要进行的运算量)。FLOPs通常在论文和技术报告中流行的单位是GFLOPs: 1GLOPs = 10^9 FLOPs,有时候也用MACs或者MAdd(Multiply–Accumulate Operations)代替FLOPs作为评估模型计算量的单位,1 MACs大约等价于 2FLOPs。

从定义中可以看出FLOPs 是一个与计算性能相关的指标,那么它的作用主要体现在当你在训练、部署模型时需要对模型的计算性能进行考虑。

  • 比如训练一个模型(LLM)时,通常通过计算模型训练全部FLOPs与使用GPU峰值的FLOPS以及GPU的利用率,来大致估算训练所需的时间和卡数。

  • 比如能够通过计算模型每个部分FLOPs得到哪个部分计算量最大,从而分析如何对模型计算性能进行优化。

  • 比如当几个模型的效果体验差不多,尽量选择FLOPs小的模型进行部署来满足服务的要求。

  • …..

二、FLOPs如何计算

既然FLOPs在考虑模型性能时有参考的价值,那我们如何得到一个模型的FLOPs?

有两种方式:

  • 根据计算公式和模型结构手动推算

  • 借助第三方工具:calflops、ptflops、thop、torchstat、torchsumary、fvcore

1、推公式

我们可以选择完全弄清楚一个模型的内部结构以及计算运行的整个过程,然后根据公式去手动计算。 这个方法可以加深你对于模型结构的理解,但缺点是估算可能有一定的误差。大部分时候只估算矩阵相乘之间计算,忽略掉激活函数,softamx之类的计算量。进行推导时,需要对模型不同层结构采用不同的计算方式,比如FC层,CNN层,LSTM层,Transformer层,具体计算方式可以看文章:训练模型算力的单位:FLOPs、FLOPS、Macs 与 估算模型(FC, CNN, LSTM, Transformers&&LLM)的FLOPs中的推导。

手动推导FLOPs原则:

  • 手动推导模型的FLOPs时,优先推导整个过程计算量占大头部分,通常忽略激活函数、layer normalize,softmax等等部分计算量。

  • 手动推导模型的FLOPs时只推导前向传播,大部分情况默认模型后向传播的计算量是前向传播的2倍, 总共FLOPs是前向的3倍。(结论出自——https://epochai.org/blog/backward-forward-FLOP-ratio)

  • 由于LLM模型参数过大,占用显存过多,有时候为了降低显存在训练采用将中间参数保留在内存里——激活重计算。因此推导LLM训练时FLOPs如果考虑到中间参数的激活重计算的过程,需要计算整体FLOPs需要再加一份前向计算量,即1(前向) + 2(反向)+ 1(激活重计算)= 4 倍 计算量。 (结论出自——https://arxiv.org/pdf/2205.05198.pdf)

2、使用工具-torchstat,thop,fvcore,ptflops, calflops

目前在github上计算模型FLOPs以及Params数目的相关工具有calflops、ptflops、thop、torchstat、 fvcore等等。本文将这些工具尝试了一遍,先从模型计算FLOPs结果的准确性上发现这里面存在将FLOPs与MACs混用甚至弄反了的情况(torchstat), 没有区分FLOPs与MACs、输出是MACs(fvcore, thop, ptflops)的情况 。其中只有calflops 计算的结果是严格区分FLOPs和MACs。

以计算模型resnet50 为例:

from torchvision.models import resnet50  
model = resnet50()  
  
# torchstat 不推荐  
from torchstat import stat  
stat(model, (3, 224, 224))  
# 在终端或者jupyter中显示结果,MACS与FLOPs结果弄反了,MACS:8.22GACs, FLOPs:4.12GFLOPs,params: 25,557,032, 不支持AdaptiveAvgPool2d操作。  
  
# thop  
from thop import profile  
input = torch.randn(1, 3, 224, 224)  
macs, params = profile(model, inputs=(input, ))  
print(macs, params)  
# 4133742592.0 25557032.0, 只返回MACs的值  
  
# fvcore   
from fvcore.nn import FlopCountAnalysis, parameter_count_table  
tensor = (torch.rand(1, 3, 224, 224),)  
flops = FlopCountAnalysis(model, tensor)  
print(flops.total())   
# 4144854528, 虽然FlopCounter, 返回的是MACs的值  
  
# ptflops  
from ptflops import get_model_complexity_info  
flops, params = get_model_complexity_info(model, (3, 224, 224))  
  
print(flops, params)  
# 4121925096 25557032  返回也是MACs  
  
# calflops  
from calflops import calculate_flops  
flops, macs, params = calculate_flops(model, input_shape=(1, 3, 224, 224))  
# 8211108352, 4089184256, 25557032  返回的结果区分flops, macs,   

基于上面测试,首先排除torchstat和fvcore,前者甚至将FLOPs与MACs弄反了,后者将MACs与FLOPs弄混了。 接着再测试一下对预训练的语言模型、参数更大的LLM来说thop,ptflops, calflops是否计算准确、使用是否方便。

3、测试thop,ptflops, calflops对LLM计算 FLOPs

由于基于Transformer预训练模型的输入需要经过其tokenizer将文本构造成对应Token id,再转化成Tensor后才能输入到模型中。 在包thop和ptflops中,都需要将构造transformer模型输入的过程先处理好,然后再借助它们A计算FLOPs的API提供的参数进行传入。例如参考文章在thop中需要实现构造transformers模型的输入,具体如下代码所示:

def _input_constructor(input_shape, tokenizer):  
    max_length = input_shape[1]  
      
    model_input_ids = []  
    model_attention_mask = []  
    model_token_type_ids = []  
    for _ in range(input_shape[0]):  
        inp_seq = ""  
        inputs = tokenizer.encode_plus(  
            inp_seq,  
            add_special_tokens=True,  
            truncation_strategy='longest_first',  
        )  
  
        input_ids, token_type_ids = inputs["input_ids"], inputs["token_type_ids"]  
        attention_mask = [1] * len(input_ids)  
        padding_length = max_length - len(input_ids)  
        pad_token = tokenizer.pad_token_id  
        pad_token_segment_id = 0  
        input_ids = input_ids + ([pad_token] * padding_length)  
        attention_mask = attention_mask + ([0] * padding_length)  
        token_type_ids = token_type_ids + ([pad_token_segment_id] * padding_length)  
        assert len(input_ids) == max_length  
        assert len(attention_mask) == max_length  
        assert len(token_type_ids) == max_length  
        model_input_ids.append(input_ids)  
        model_attention_mask.append(attention_mask)  
        model_token_type_ids.append(token_type_ids)  
  
    labels = torch.tensor([1] * input_shape[0])  
    # Batch size input_shape[0], sequence length input_shape[128]  
    inputs = {  
        "input_ids": torch.tensor(model_input_ids),  
        "token_type_ids": torch.tensor(model_token_type_ids),  
        "attention_mask": torch.tensor(model_attention_mask),  
    }  
    return inputs  
  
inputs = _input_constructor  
tuples = (inputs["input_ids"],   
          inputs["token_type_ids"],  
          inputs["attention_mask"])  
total_ops, total_params = profile(model, inputs=tuples,)  

包ptflops在计算LLM模型与上面类似,而在包calflops中,它只需要将模型对应的tokenizer作为参数transformer_tokenizer 的赋值传入即可,calflops会自动构造指定大小的模型输入,使用更加简便:

flops, macs, params = calculate_flops(model=model,  
                                      input_shape=(batch_size, seq_length),  
                                      transformer_tokenizer=tokenizer)  

从计算FLOPs和Parameterde的结果上,thop、ptflops、calflops 对于LLM的参数是一致,FLOPs有差异是因为统计上对有些算子进行忽略。总体来说,三者使用一致情况下,都可以作为参考工具。

从具体实现细节上,相比ptflops,calflops参考deepspeed profile实现了支持torch.nn.function.*操作统计FLOPs的计算量,从而支持更细粒度的实现方式模型(包括自定义模型)统计FLOPs,因此计算准确度更高。

4、推荐工具——calflops

基于上面测评,综合从计算FLOPs的准确性和对语言模型使用更简易上,计算模型FLOPs最推荐工具是calflops ,其开源地址:https://github.com/MrYxJ/calculate-flops.pytorch,使用文档非常详细。

安装方式:

pip install calflops  

5、计算部分LLM的FLOPs与Parameter

calflops对于LLM计算FLOPs提供很方便的使用方式,目前作者已自测部分开源LLM的FLOPs在其文档中:

github: https://github.com/MrYxJ/calculate-flops.pytorch

鉴于目前LLM依然在层出不求,欢迎更多人使用calflops测试新出的LLM的FLOPs继续更新在该项目的文档中,毕竟LLM时代的模型性能直接关系训练和推理的体验,祝大家训练LLM速度++,loss降降。

6、从计算FLOPs到预估LLM训练时间

从本文上面的手动推算FLOPs所讨论的原则中模型训练有:

我们利用工具calflops等工具计算出模型FLOPs只是模型选定batch_size下一次step的FLOPs,等价于公式中 模型单次。

由于step受到batch_size的影响,实际训练时不是一个定量。在LLM中,预训练语料Token数量是一个用来评估训练预料规模的定量。还是根据文章中关于LLM模型FLOPs的计算推导有:

  • 如果开启中间参数激活重计算:

则LLM 最终训练

总结

本文主要介绍了大模型训练过程中TFLOPS的计算方式以及pytorch实现,可以用于大模型训练时间花费的估计。‍

当然,本文中给出了几个进一步延伸阅读的工作,感兴趣的可以进一步查看,以增进理解。‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍

本文由Converge供稿,是个不错的工作,欢迎去看看。‍‍‍‍‍‍‍‍‍

参考文献

1、https://www.zhihu.com/column/c_1615427893089206272

2、https://arxiv.org/pdf/2205.05198.pdf

3、https://github.com/MrYxJ/calculate-flops.pytorch

关于我们

老刘,刘焕勇,NLP开源爱好者与践行者,主页:https://liuhuanyong.github.io。

老刘说NLP,将定期发布语言资源、工程实践、技术总结等内容,欢迎关注。

对于想加入更优质的知识图谱、事件图谱、大模型AIGC实践、相关分享的,可关注公众号,在后台菜单栏中点击会员社区->会员入群加入。

​​​​​

更多AI工具,参考Github-AiBard123国内AiBard123

可关注我们的公众号:每天AI新工具