Transformer-based模型到底要训练多久
- -
嘿，你知道一年多前我干过什么蠢事吗？说出来你别笑！
我手里攥着实验室仅有的两张Titan XP，师兄们天天在跑Bert的各种变体，我一边羡慕一边焦虑：这些Transformer模型到底要训练多久？我这破卡能干嘛？
然后我给自己找了个课题——研究FLOPs和显卡算力的关系，写出一个通用的训练时间估算公式，最好能投个会议。为此我泡了一星期Arxiv和NVIDIA文档，把Transformer里每个矩阵乘法都拆开算了一遍。
结果呢？你猜对了。论文鸽了——考试、赶项目、再鸽。就像大三时鸽了微软老师的强化学习程序一样，这事儿成了我心里一个疙瘩。
但那段研究真没白费！至少现在我看到任何开源模型，第一反应不是“卧槽牛逼”，而是“让我算算得跑多久”。你看，傻事也能长本事。
- -
说到Transformer刚出来那会儿，可不是现在这待遇。
《Attention Is All You Need》确实是神仙论文，但在当时的Google Brain，它只是每年几千篇里的之一。研究员们的KPI是发顶会刷SOTA，Transformer在机器翻译上效果很棒，但谁也没料到它会统治整个NLP——更别说后来连CV、语音也通吃了。
前Google Brain的高级研究科学家Ashish Vaswani说过一句话，我记到现在：“意义是事物之间的关系所产生的结果，而自注意力是一种学习关系的通用方法。”这话现在看，简直预言家！但2017年你敢说这个架构会让Google搜索重构一遍？估计会被当疯子。
而且你绝对想不到，早期的Transformer训练成本有多温柔。原始论文里base模型在8张P100上跑了3.5天，big模型也是3.5天——因为有warmup和checkpoint average。放在今天，这简直是神仙日子！那时候的我们哪知道，后面会跟着上千张卡、训练几个月的巨兽。
- -
说回训练时间估算这事儿。
我当初研究FLOPs的时候，发现网上大部分资料都在绕弯子。核心就一句话：Transformer的训练运算量主要来自矩阵乘法，不考虑Embedding查表和dropout之类的小操作，一次前向+反向的总FLOPs大约是 6 × 模型参数量 × 训练token数。
怎么来的？你看啊，拿一个解码层举例：self-attention里Q、K、V三个线性层，加上attention输出映射，加上两个FFN层，每个参数在前向和反向都要被计算两遍——前
读者评论 2