【LLM】OpenELM: 一个高效的语言模型系列，具有开源训练和推理框架

一、结论写在前面

论文发布了OpenELM，一个基于Transformer的开源仅解码器语言模型。OpenELM使用了逐层缩放方法在Transformer模型内部实现高效的参数分配，从而相比现有模型提高了准确率。另外，论文已经开源了整个框架，包括训练日志、多个检查点、预训练配置和MLX推理代码。论文的源代码、预训练模型权重和训练配方可在https://github.com/apple/corenet获取。此外，OpenELM模型也可以在HuggingFace上找到:https://huggingface.co/apple/OpenELM。

二、论文的简单介绍

2.1 论文的背景

基于Transformer的大语言模型(LLM)正在革新自然语言处理领域。这些模型是各向同性的，意味着它们在每个Transformer层都具有相同的配置(例如，头数和前馈网络维数)。尽管这种各向同性模型很简单，但它们可能无法在模型内部高效分配参数。

在论文开发并发布了OpenELM，一个在公开可用数据集上进行预训练和微调的模型系列。OpenELM的核心是逐层缩放[30]，可以跨层更有效地分配参数。该方法在距离输入更近的Transformer层的注意力和前馈模块中使用更小的潜在维度，并随着接近输出逐渐扩大层的规模。

2.2 论文的方法

2.2.1 OpenELM架构

论文采用了基于Transformer的仅解码器架构。根据最先进的LLM，论文：(1)在任何全连接(也称为线性)层中都不使用可学习的偏置参数；(2)使用RMSNorm 应用预归一化，并且使用旋转位置嵌入(ROPE)对位置信息进行编码；(3)使用分组查询注意力(GQA)代替多头注意力(MHA)；(4)用SwiGLU FFN 代替前馈网络(FFN)；(5)使用闪电注意力计算缩放的点积注意力；(6)使用与LLama 相同的tokenizer。

现有的LLM在模型的每一层Transformer中使用相同的配置，导致参数在各层之间均匀分配。与这些模型不同，OpenELM中的每个Transformer层都有不同的配置(例如注意力头数和前馈网络维度)，从而导致模型的每一层参数数量不同。这使OpenELM能够更好地利用可用的参数预算来实现更高的准确性。论文使用逐层缩放(逐块缩放)来实现跨层的非均匀参数分配。    

逐层缩放：标准Transformer层由多头注意力(MHA)和前馈网络(FFN)组成。为实现Transformer层中参数的非均匀分配，论文调整每个Transformer层中的注意力头数和FFN乘数。假设具有均匀参数分配的标准Transformer模型有N个Transformer层，每层输入的维数为dmodel。

2.2.2. 预训练数据

对于预训练，论文使用公开数据集。具体来说，论文的预训练数据集包含RefinedWeb 、去重后的PILE 、RedPajama 的一个子集和Dolma v1.6的一个子集，总计约1.8万亿个tokens。这些细节也总结在表2中。

即时标记化和数据过滤：与利用预标记化数据[5，17]的之前方法不同，论文即时过滤和标记文本数据。这有助于轻松尝试各种标记器，从而显著简化了原型设计和研究工作。在论文的实验中，论文使用与LLama [46]相同的标记器。

为过滤掉低长度序列，论文应用两种过滤方法。第一种方法在字符级别进行操作，检查序列中字符数是否低于指定阈值。第二种方法在token级别进行操作，检查序列中的token数是否少于指定阈值。序列长度低于任一阈值时将被跳过。在论文的实验中，论文将字符级和token级过滤阈值设置为200个字符和256个token。

2.2.3. 训练细节

论文使用CoreNet(之前称为CVNets)训练OpenELM变体350k次迭代(或训练步骤)。论文使用AdamW 作为优化器。论文使用余弦学习率    

2.2.4. 评估细节

按照之前的工作，论文使用LM评估工具包在不同任务上评估性能:

•标准zero-shot任务。论文考虑了7个标准常识推理任务：ARC easy和challenge、BoolQ、HellaSwag、PIQA、SciQ和WinoGrande。

•OpenLLM排行榜任务。论文使用OpenLLM排行榜中的5个任务：ARC challenge、HellaSwag、MMLU、TruthfulQA和WinoGrande。

•LLM360排行榜任务。论文使用LLM360排行榜中的7个任务进行评估：ARC challenge、CrowS-Pairs(英语版)、HellaSwag、WinoGrande、MMLU、PIQA和RACE。

这些建立在LM评估工具包之上的评估框架，使论文能够全面评估OpenELM在推理(如ARC-c、HellaSwag和PIQA)、知识理解(如MMLU和RACE)以及虚假信息&偏差(如TruthfulQA和CrowS-Pairs)方面的表现。虽然这些框架中可能存在一些任务重叠，但它们主要区别在于少量示例设置，如表3所示。

2.3 论文的效果

2.3.1 实验结果

预训练结果：论文评估了OpenELM在zero-shot和少量示例设置下的表现(表3)。论文将OpenELM与公开可用的LLM进行了比较，即PyThia、Cerebras-GPT、TinyLlama、OpenLM、MobiLlama和OLMo。与论文工作最相关的是MobiLlama和OLMo。这些模型是在可比的数据集组合上训练的，具有相似或更大的预训练token数量。

在图1中，论文绘制了OpenELM在7个标准zero-shot任务上的准确率与训练迭代次数的关系。论文观察到，在大多数任务中，随着训练时间的延长，准确率总体上有所提高。另外，通过平均最后五个检查点(间隔5000次迭代收集)获得的检查点，其准确率与350k次迭代后获得的最终检查点相当或略有改善。这种改进可能是由于权重平均导致噪声减少。因此，论文在表4的主要评估、表5的指令调优实验和表6的参数高效调优实验中使用了平均检查点。

表4中的结果跨越了各种评估框架，并突出了OpenELM相较于现有方法的有效性。例如，具有11亿参数的OpenELM变体比12亿参数的OLMo的准确率分别高出1.28%(表4a)、2.36%(表4b)和1.72%(表4c)。值得注意的是，OpenELM在使用2倍少的预训练数据的情况下实现了这一准确率水平。    

指令调优结果：论文使用UltraFeedback清理后的变体数据集进行指令调优，该数据集包含60k个提示。论文使用Alignment Handbook库进行指令调优。在优化过程中，论文使用统计rejection sampling方法或直接preference optimization方法。表5显示，指令调优可以使OpenELM在不同评估框架下的平均准确率持续提高1-2%。

参数高效微调(PEFT)结果：论文使用CommonSense推理训练和评估设置。该设置为PEFT研究提供了8个多重选择数据集的170k训练样本，可用于不同方法，包括LoRA和DoRA。论文将OpenELM与这些方法集成，并使用8个NVIDIA H100 GPU对生成的模型进行三个epoch的微调。表6显示，PEFT方法可应用于OpenELM。在给定的CommonSense推理数据集上，LoRA和DoRA在平均准确率方面表现相似。    

指令调优结果：论文使用经过清理的UltraFeedback 数据集变体进行指令调优，该数据集包含60k个提示。论文使用Alignment Handbook库进行指令调优。在优化过程中，论文使用统计rejection sampling方法或直接preference optimization方法。表5显示，指令调优可以使OpenELM在不同评估框架下的平均准确率持续提高1-2%。

参数高效微调(PEFT)结果：论文使用CommonSense推理训练和评估设置。该设置为PEFT研究提供了8个多重选择数据集的170k训练样本，可用于不同方法，包括LoRA和DoRA。论文将OpenELM与这些方法集成，并使用8个NVIDIA H100 GPU对生成的模型进行三个epoch的微调。表6显示，PEFT方法可应用于OpenELM。在给定的CommonSense推理数据集上，LoRA和DoRA在平均准确率方面表现相似。    

2.3.2 基准测试

评估：论文分别为token吞吐量(以每秒处理的token数衡量)提供了两种独立的测量：(1)prompt处理(预填充)和(2)token生成。此外，论文还报告了总的组合吞吐量。论文按顺序对所有模型进行基准测试，并为第一个模型执行一次完整的"干运行"，生成1024个token，因为论文发现这可以显著提高后续模型的生成吞吐量。在测量每个单独模型之前，论文通过执行单次前向传递来预热模型，以允许框架执行进一步的自动优化(如有)。在所有实验中，论文都使用了键值缓存，并在所有测试中生成了1024个token外加prompt token。只要支持，论文就使用静态键值缓存。所有运行都使用了相同的prompt，导致prompt长度为35-36个token(取决于tokenizer)。    

结果：表7a和7b分别显示了GPU和MacBook Pro上的基准测试结果。尽管在相似参数量下OpenELM的准确率更高，但论文观察到它比OLMo慢。虽然本研究的主要重点是可重复性而非推理性能，但论文还是进行了全面的分析以了解瓶颈所在。论文的分析发现，OpenELM的处理时间中有相当一部分可归因于论文对RMSNorm的朴素实现(表8)。

具体来说，朴素的RMSNorm实现导致了许多单独的内核启动，每个内核处理一小部分输入，而不是像LayerNorm那样启动单个融合内核。通过用Apex的RMSNorm替换朴素的RMSNorm，论文观察到OpenELM的吞吐量显著提高。然而，与使用优化LayerNorm的模型相比，仍然存在相当大的性能差距，部分原因是(1)OpenELM有113个RMSNorm层，而OLMo只有33个LayerNorm层；(2)Apex的RMSNorm对小输入未做优化。为进一步说明RMSNorm导致的性能下降，论文用RMSNorm替换了OLMo中的LayerNorm，观察到生成吞吐量显著下降。在未来的工作中，论文计划探索优化策略以进一步提高OpenELM的推理效率。