0. 这篇文章干了啥?
ChatGPT等大语言模型在日常、科研、工作上的作用不必多说,但是大模型的训练代码一般是不公开的,只会开放专用接口或者推理代码。
但最近OLMo开源了!这是一个最先进的、真正开放的语言模型和框架,完整提供训练数据、训练和评估代码、中间模型检查点和训练日志。
对于训练和建模,OLMo包括完整的模型权重、训练代码、训练日志、切除信息、以Weights & Biases日志形式的训练指标以及推理代码。OLMo也是首次发布包括7B规模的语言模型的四个变体,对应不同的架构、优化器和训练硬件,以及一个1B规模的模型,都至少训练了2T Token。还在HuggingFace上发布了数百个中间检查点作为修订版本。
对于数据集构建和分析,OLMo包括用于这些模型的完整训练数据,包括产生训练数据的代码。
对于评估,OLMo包括AI2的Catwalk用于下游评估和Paloma用于perplexity评估。
对于微调,OLMo发布了Open Instruct,并且将很快发布正在使用它制作一个调整过的版本的OLMo。并且所有代码和权重都在Apache 2.0许可下发布。
下面一起来阅读一下这项工作~
1. 论文信息
标题:OLMo: Accelerating the Science of Language Models
作者:Dirk Groeneveld,et al
机构:艾伦人工智能研究所、华盛顿大学、耶鲁大学、纽约大学、卡内基梅隆大学
原文链接:https://arxiv.org/abs/2402.00838
代码链接:https://github.com/allenai/OLMo
权重链接:https://huggingface.co/allenai/OLMo-7B
数据链接:https://huggingface.co/datasets/allenai/dolma
评估链接:https://github.com/allenai/OLMo-Eval
微调链接:https://github.com/allenai/open-instruct
2. 摘要
语言模型(LMs)已经在自然语言处理研究和商业产品中变得无处不在。随着它们在商业上的重要性不断增加,最强大的模型变得封闭起来,被封锁在专有接口背后,其训练数据、架构和开发的重要细节未公开。考虑到在科学研究这些模型时,包括它们的偏见和潜在风险在内的这些细节的重要性,我们认为研究社区有必要能够访问功能强大、真正开放的LMs。为此,本技术报告详细介绍了OLMo的首次发布,这是一款领先的、真正开放的语言模型,以及用于构建和研究语言建模科学的框架。与大多数先前仅发布模型权重和推理代码的努力不同,我们发布了OLMo和整个框架,包括训练数据、训练和评估代码。我们希望这一发布能够赋予和加强开放研究社区,并激发新一波创新。
3. 发布清单
整个OLMo发布以下内容:
(1)训练和建模代码:https://github.com/allenai/OLMo
(2)7B模型:https://huggingface.co/allenai/OLMo-7B
7B-twin-2T模型:https://huggingface.co/allenai/OLMo-7B-Twin-2T
1B模型:https://huggingface.co/allenai/OLMo-1B的训练模型权重。
对于所有模型,不仅发布最终的模型权重,还发布间隔为1000步的500多个Checkpoint。
(3)训练数据Dolma:https://huggingface.co/datasets/allenai/dolma
(4)Dolma的工具包,用于构建新的数据集:https://github.com/allenai/dolma,以及用于数据集分析的WIMBD:https://github.com/allenai/wimbd。
(5)用于下游评估的评估代码:https://github.com/allenai/OLMo-Eval,使用Catwalk(https://github.com/allenai/catwalk)进行和用于perplexity评估的Paloma(https://paloma.allen.ai/)。
作者计划很快发布另一个版本,其中包括以下内容:
训练日志、消融和发现。
训练运行的Weights & Biases日志。
带有指导调整和RLHF的调整OLMo,包括其训练和评估代码以及使用Open Instruct库的数据:https://github.com/allenai/open-instruct。
4. OLMo框架
4.4 OLMo 模型和架构
OLMo采用了解码器-专用 Transformer 架构,并提供了 1B 和 7B 两个变体,65B 版本即将推出。
大语言模型通常通过在硬件上优化训练吞吐量并最小化丢失峰值和缓慢发散的风险来选择超参数。下表将OLMo的设计选择与最新的开放语言模型进行了比较。OLMo对原始 Transformer 架构的主要更改可以总结如下:
(1)无偏置。与 LLaMA、PaLM 等人一样,我们在架构中排除了所有偏置项,以提高训练稳定性。
(2)非参数层归一化。使用层归一化的非参数形式,在该归一化中,没有范围内的仿射变换,即没有"自适应增益"(或偏置)。作者认为这是最安全的选择,而且与其他变体相比速度更快:参数层归一化和 RMSNorm。
(3)SwiGLU 激活函数。与 LLaMA、PaLM 等人一样,使用 SwiGLU 激活函数而不是 ReLU,且根据 LLaMA,激活隐藏大小约为 8/3d,但增加到最接近 128 的倍数(例如,对于7B 模型为 11008)以提高吞吐量。
(4)旋转位置嵌入(RoPE)。与 LLaMA、PaLM 等人一样,用旋转位置嵌入替换绝对位置嵌入。
(5)词汇表。使用 GPT-NeoX-20B 的基于 BPE 的分词器的修改版本,其中还包括用于屏蔽个人可识别信息 (PII) 的额外标记。最终的词汇表大小为 50,280。但为了最大化训练吞吐量,将模型中相应嵌入矩阵的大小增加到 50,304,以使其成为 128 的倍数。
4.2 预训练数据:Dolma
尽管在获取模型参数方面取得了进展,但预训练数据集仍然不够开放。预训练数据集通常不会与开放模型(更不用说闭源模型)一起发布,关于这些数据的文档通常缺乏复制或完全理解工作所需的详细信息。这使得支持语言模型研究的某些方向变得困难,比如理解训练数据如何影响模型的能力和限制。为促进语言模型预训练的开放研究,OLMo构建并发布了预训练数据集 Dolma,这是一个包含来自 7 个不同数据源的 5B 文档中的 3T Token的多样化、多源语料库,这些数据源(1)在大规模语言模型预训练中经常见到,并且(2)对一般公众可访问。
Dolma 是通过(1)语言过滤,(2)质量过滤,(3)内容过滤,(4)去重,(5)多源混合和(6)分词的流程构建的,并提供在 Dolma 的中间状态上训练语言模型的额外分析和实验结果。作者在整个整理过程中以及最终发布时都将每个来源的文档保持分开,还开源了高性能数据整理工具;这个工具包可以用于在 Dolma 上进一步实验、重现工作,并实现对预训练语料库的快速且轻松的整理。最后,还开源了WIMBD 工具以帮助进行数据集分析。
4.3 评估
在两作者个阶段进行模型评估:在线评估用于制定模型设计决策,离线评估用于评估模型检查点。
对于离线评估,使用 Catwalk 框架,这是公开提供的评估工具,具有广泛的数据集和任务格式访问权限。使用 Catwalk进行下游评估以及对新 perplexity 基准 Paloma进行内在语言建模评估。对于下游和 perplexity 评估,使用作者的固定评估管道,与几个公开可用的模型进行比较。
循环内训练消融: 在整个模型训练过程中,进行下游评估以做出关于模型架构、初始化、优化器、学习速率调度和数据混合的决策。作者称之为在线评估,因为它每隔 1000 次训练步骤(或约 4B 训练令牌)运行一次,提供关于正在训练的模型质量的早期和持续信号。
下游评估:在很多先前的工作的基础上,报告了zero-shot性能在一组下游任务上。作者的评估套件包括 9 个核心任务。鉴于正在评估的模型的规模,这些任务是在模型开发初期选择的,因为它们的自然性(例如,都可以被制定为文本完成评分任务)和在训练过程中提供有意义信号的能力。
内在语言建模评估:为了衡量 OLMo-7B 如何适应语言分布超出留出的训练数据,使用 Paloma,这是一个包含 585 种不同领域文本的新 perplexity 基准。这些领域从 nytimes.com 到 Reddit 上的 r/depression 等,都来自于 18 个不同的数据源,这允许更均等地包含在它们的源语料库中代表性不足的文本领域。
作者的目标不仅是将 OLMo-7B 与其他模型进行最佳性能比较,而且是展示它如何实现更充分、更受控制的科学评估。OLMo-7B 是具有明确去污染的 perplexity 评估的最大语言模型。按照 Paloma 描述的方法,作者删除了任何在 Paloma 评估数据泄漏的段落的预训练文档。没有去污染,其他模型可能低估 perplexity(即高估模型的样本外适应性)。作者还发布了中间检查点,允许与另外两个发布检查点的模型进行更丰富的比较。
5. 训练OLMo
5.1 分布式训练框架
OLMo使用ZeRO优化器策略,通过PyTorch的FSDP框架进行模型训练,该框架通过在GPU上分片模型权重及其对应的优化器状态来减少内存消耗。在7B规模下,可以在硬件上每个GPU上使用4096个令牌的微批量大小进行训练。对于OLMo-1B和-7B模型,使用约4M令牌(2048个实例,每个实例的序列长度为2048个令牌)的常量全局批量大小。对于OLMo-65B模型(目前正在训练),使用一个批量大小的预热,从约2M令牌(1024个实例)开始,然后每100B令牌增加一倍,直到达到约16M令牌(8192个实例)。为提高吞吐量,使用混合精度训练,通过FSDP的内置设置和PyTorch的amp模块实现。后者确保像softmax这样的某些操作始终以完全精度运行,以提高稳定性,而所有其他操作则以半精度bfloat16格式运行。
5.2 优化器
OLMo使用AdamW优化器。对于所有模型大小,在5000步(∼21B令牌)内调整学习率,然后线性衰减至剩余训练过程中峰值学习率的十分之一。在预热期后,剪切梯度,使参数梯度的总l2-范数不超过1.0。
5.3 数据
训练数据集是从作者的开放数据集Dolma中提取的2T令牌样本构建的,每个文档的令牌在附加特殊的EOS令牌到达每个文档的末尾后连接在一起,然后将连续的2048个令牌的块组合成训练实例。每次训练运行,训练实例都以完全相同的方式进行洗牌。所有发布的模型都至少经过2T令牌的训练。
5.4 硬件
为了验证代码库在NVIDIA和AMD GPU上都能够在性能上无损使用,在两个不同的集群上训练了模型:
• LUMI:由LUMI超级计算机提供,在该集群上使用了最多256个节点,每个节点由4个AMD MI250X GPU组成,内存为128GB,互连速度为800Gbps。
• MosaicML:在该集群上使用了27个节点,每个节点由8个NVIDIA A100 GPU组成,内存为40GB,互连速度为800Gbps。
尽管为了优化训练吞吐量而对批量大小进行了轻微调整,但两次运行在2T令牌时的评估套件中表现出几乎相同的性能。
6 结果
6.1 下游评估
核心下游评估套件包括:arc(arc easy和arc challenge),boolq,openbookqa,sciq,hellaswag,piqa,copa和winogrande。作者总结了对OLMo-7B的零次评估结果,并将其与6个规模相当的其他公开可用模型进行了比较。OLMo-7B检查点在2个端任务上表现优异,并在评估套件的8/9端任务中保持前三。总体而言,OLMo-7B对比其他6个公开可用的模型检查点是具有竞争力的。
9个核心端任务的准确性分数进展情况。除了OBQA之外的所有任务在OLMo-7B训练更多令牌时都呈上升趋势。在最后一步和倒数第二步之间,许多任务的准确性急剧上升,显示了在线性降低LR至0的最后1000训练步骤中的收益。
6.2 内在语言建模评估
对于内在评估,Paloma分别检查每个领域的性能到对Paloma的18个源组合的更概括的结果。在图2的Sources Combined子图中,展示了OLMo-7B在Paloma的11个数据源组合上与6个相似大小的语言模型的性能。总体而言,OLMo具有竞争力,特别是考虑到其训练数据明确经过Paloma的净化。
7. 总结 & 未来工作
作者开源了首个版本的OLMo,这是一款先进的、真正开放的语言模型及其用于构建和研究语言建模科学的框架。与大多数之前仅发布模型权重和推理代码的尝试不同,作者发布了OLMo和整个框架,包括训练数据、训练和评估代码。很快,还将发布训练日志、消融实验、研究结果以及Weights & Biases日志。还正在探索通过指导调整和不同变种的RLHF对OLMo进行适应。最终将发布适应的模型以及所有的模型适应代码和数据。
作者打算持续支持和扩展OLMo及其框架,并继续推动开放式语言模型的边界,以增强开放研究社区的能力。为此,会将不同的模型大小、模态、数据集、安全措施和评估引入到OLMo家族中。
对更多实验结果和文章细节感兴趣的读者,可以阅读一下论文原文~