24年6月来自Meta AI论文“Meta Large Language Model Compiler: Foundation Models of Compiler Optimization”。
大语言模型 (LLM) 在各种软件工程和编码任务中都表现出了卓越的能力。然而,它们在代码和编译器优化领域的应用仍未得到充分探索。训练 LLM 需要大量资源,需要大量的 GPU 时间和大量的数据收集,这可能会让人望而却步。为了解决这一差距,Meta推出 大语言模型编译器 (LLM Compiler),这是一套强大的、公开可用的、预训练的模型,专为代码优化任务而设计。LLM Compiler 建立在 Code Llama 的基础上,增强了对编译器中间表示 (IR)、汇编语言和优化技术的理解。该模型已经在 5460 亿个 LLVM-IR 和汇编代码的庞大语料库上进行了训练,并经过指令微调以解释编译器行为。LLM Compiler 是根据定制的商业许可发布的,允许广泛重复使用,有两种大小:70 亿个和 130 亿个参数。还提供了该模型的微调版本,展示了其在优化代码大小和从 x86_64 和 ARM 汇编反汇编回 LLVM-IR 方面的增强功能。这些实现了自动调优搜索的 77% 优化潜力和 反汇编来回的45%(14% 精确匹配)。此版旨在为学术研究人员和行业从业者在编译器优化方面的进一步研究和开发提供可扩展、经济高效的基础。
LLM Compiler 模型以编译器优化为目标。它们有两种模型大小:7B 和 13B 参数。LLM Compiler 模型使用相应大小的 Code Llama 模型权重进行初始化,并在额外的 5460 亿 tokens 数据上进行训练,这些数据主要包括编译器中间表示和汇编代码。然后,用额外的 1640 亿 tokens 数据进一步训练 LLM Compiler FTD 模型,完成两个下游编译任务:标志调优(flag tuning)和反汇编。在训练的所有阶段,都会使用来自先前阶段的少量代码和自然语言数据来帮助保留基本 Code Llama 模型的功能。
如图所示:LLM Compiler 模型是 Code Llama 的特殊版本,通过两个阶段对 5460 亿个以编译器为中心的数据进行训练而来。第一阶段,模型主要在未标记的编译器 IR 和汇编代码上进行训练。在下一阶段,模型将进行指令微调,预测优化的输出和效果。然后,LLM Compiler FTD 模型将进一步在 1640 亿个下游标志调优和反汇编任务数据集上进行微调,总共 7100 亿个训练 tokens。在四个训练阶段中,每个阶段都会保留 15% 的先前任务数据。
用于训练编码 LLM 的数据通常主要由 Python 等高级源语言组成。
汇编代码在这些数据集中所占比例微不足道,而编译器 IR 所占比例则更小。
为了构建能够很好地理解这些语言的 LLM,用 Code Llama 的权重初始化 LLM Compiler 模型,然后在以编译器为中心的数据集上训练 4010 亿个tokens,该数据集主要由汇编代码和编译器 IR 组成,如表所示。
LLM Compiler 主要在 LLVM(Lattner & Adve,2004)版本 17.0.6 生成的编译器中间表示和汇编代码上进行训练。
这些代码来自用于训练 Code Llama 的同一公开代码数据集。
在如下表中总结了该数据集。
与 Code Llama 一样,也从自然语言数据集中获取了训练批次的一小部分。
为了理解代码优化的机制,指令微调 LLM Compiler 模型来模拟编译器优化,如图所示。
其思想是通过对这些程序应用随机生成的编译器优化序列,从一组有限的未优化种子程序中生成大量示例。
然后,训练模型来预测优化生成的代码。
还训练模型来预测应用优化后的代码大小。
众所周知,操纵Compiler标志,会对运行时性能和代码大小产生相当大的影响(Fursin,2005)。
在下游任务上训练 LLM Compiler FTD 模型,即为 LLVM 的 IR 优化工具 opt 选择标志的下游任务,产生最小的代码大小。
以前机器学习方法对标志的调优已经显示出良好的结果,但很难在不同程序之间泛化(Cummins,2022)。
以前的工作通常需要编译新程序数十或数百次来尝试不同的配置并找出性能最佳的选项。
通过预测标志来最小化未见过程序的代码大小,在该任务的零样本版本上训练和评估 LLM Compiler FTD 模型。
这个方法与所选的编译器和优化指标无关,在未来将以运行时性能为目标。
目前,针对代码大小进行优化简化了训练数据的收集。
如图说明了用于生成训练数据的过程以及如何使用模型进行推理。评估时,只需要所生成的传递列表。从模型输出中提取传递列表并使用给定的参数运行工具 opt。然后可以评估模型预测二进制大小和优化输出代码的准确性,但这些都是辅助学习任务,不需要使用的。