开篇概述
对于一个 LLM 来讲,大家现在最经常干的事就是 Prompt 优化、然后是 RAG、如果这两个都不成,会考虑用一些 FineTune 来搞一下,大部分人的工作也就止步于此了。是大家不想继续深入下去吗?是大家不想做完整的 From Scratch Training 吗?其实是因为大部分团队没有这样的技术框架基础与数据基础。
现在 Olmo 给了所有人一个一窥全貌的机会。Allen AI发布了OLMo的报告, 开源一整套大模型构建的代码, 包括:完整的训练代码、模型权重、中间检查点、训练数据、评估代码。 此外,还发布了训练日志、训练指标、训练日志、以及用于指令调整的Open Instruct库。所有这些工件都在Apache 2.0许可证下发布,即可商用。下面, 让我们仔细分析一下这篇报告了。
论文基本信息
- 论文标题:OLMo: Accelerating the Science of Language Models
- 发表机构:Allen Institute for Artificial Intelligence
- 论文链接:arxiv.org/pdf/2402.00…
- 相关项目地址:
- OLMo模型:huggingface.co/allenai/OLM…
- 项目代码:github.com/allenai/OLM…
- 预训练数据集Dolma:huggingface.co/datasets/al…
- 评估工具OLMo-Eval:github.com/allenai/OLM…
OLMo模型框架全景
OLMo框架由多个组成部分构成,包括模型、预训练数据集Dolma和评估工具。
1. OLMo模型和架构
OLMo采用基于Vaswani等人的decoder-only Transformer架构[1],并提供了1B和7B两种规模的变体。该架构在标准Transformer的基础上进行了多项改进,包括:
1. 不使用偏置项: 与LLaMA、PaLM等模型一样,OLMo在其架构中排除了所有偏置项。
2. 非参数化层归一化: OLMo采用了非参数化的层归一化公式,其中没有层内的仿射变换,即没有“自适应增益”(或偏置)。这被认为是最安全的选项,并且在考虑的其他变体中速度最快。
3. SwiGLU激活函数: 与LLaMA和PaLM等模型一样,OLMo使用SwiGLU激活函数代替ReLU,并且激活隐藏层的大小大约是8d,但为了提高吞吐量,增加到最接近128的倍数。
4. 旋转位置嵌入(RoPE): 与LLaMA、PaLM等模型一样,OLMo使用旋转位置嵌入代替绝对位置嵌入。
5. 词汇表: OLMo使用了GPT-NeoX-20B的基于BPE的分词器的修改版本,并添加了额外的标记以掩盖个人可识别信息(PII)。最终的词汇表大小为50280,但为了最大化训练吞吐量,模型中对应的嵌入矩阵大小增加到50,304,以确保它是128的倍数。
简而言之,和其他大模型相比,OLMo本身并没有特别创意的设计,模型架构是中规中矩的。
2. 预训练数据:Dolma
OLMo开源了名为Dolma的语料库,包含7个不同数据源的3T个token,覆盖了5B个文档。Dolma的构建包括语言过滤、质量过滤、内容过滤、去重、多源混合和分词等步骤。
数据源的来源和大小都算是中规中矩,但是Dolma提供了language detection, dedup之类的工具。其中比较有趣的是用Rust做dedup。Allen AI还开源了个What's In My Big Data来分析文本的分布, 算是个挺有用的工具。
3. 评估工具
OpenAI的Greg Brockman说Eval is all you need,对于LLM的eval本身是一个很有难度的事情,事实上现在进行eval本身也是一个小的研究方向。设计好的benchmark, 并且能快速计算出来对分析模型非常重要。
OLMo模型的评估分为在线评估和离线评估两个阶段。在线评估在模型训练过程中进行,以便对模型架构、初始化、优化器、学习率计划和数据混合等做出决策。每1000个训练步骤(约4B个训练标记)评估一次。
OLMo的评估工具包括Catwalk用于下游评估,Paloma用于基于困惑度的评估。这些工具既可以用于在线评估也可用于离线评估。 促进对语言模型预训练的开放研究,帮助研究者理解训练数据如何影响模型的能力和局限性。
训练过程
OLMo模型的分布式训练框架、优化器设置、数据准备和硬件配置
OLMo模型的训练采用了分布式训练框架,具体使用了ZeRO优化器策略,通过PyTorch的FSDP(Fully Sharded Data Parallel)框架实现。在7B模型规模下,允许在每个GPU上以4096个token的batch大小进行训练。
在优化器方面,OLMo模型使用了AdamW优化器,并在训练的前5000步(约21B个token)内对学习率进行预热,然后在训练的剩余部分中将其线性衰减至峰值学习率的十分之一。
为了提高吞吐量,模型采用了混合精度训练,通过FSDP的内置设置和PyTorch的amp模块确保某些操作如softmax始终以全精度运行以提高稳定性,而其他操作则以半精度的bfloat16格式运行。
硬件配置方面,在两个不同的集群上训练了模型。一个是LUMI超级计算机提供的集群,使用了最多256个节点,每个节点包含4个AMD MI250X GPU,每个GPU显存128GB。另一个是MosaicML(Databricks)提供的集群,使用了27个节点,每个节点包含8个NVIDIA A100 GPU,每个GPU拥有40GB显存。
这里比较有意思的是OLMo提供了对于AMD GPU的支持,这下AMD Yes了。
实验结果与分析
引用
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.