生成式人工智能/大语言模型

应用具有推理能力的专用大语言模型(LLM)加速电池研究

人工评估材料通常会拖慢电池创新等复杂领域的科学研究,限制了每天只有数十个候选材料的进展。在这篇博客文章中,我们探讨了通过推理能力增强的领域适应型大语言模型 (LLM) 如何改变科学研究,尤其是在电池创新等高风险、复杂领域。我们深入探讨了 SES AI 的 Molecular Universe LLM,这是一个具有 70B 参数的科学 LLM,展示了这种方法的实际应用。

您将了解使用 NVIDIA NeMo Curator NVIDIA NeMo Framework NVIDIA DGX Cloud NVIDIA NIM 构建的训练和推理工作流,以及如何结合领域自适应、指令调整和推理对齐等技术来加速科学发现,同时提高专家工作效率。

简介

Flowchart depicting the Molecular Universe LLM training pipeline.
图 1。The Molecular Universe LLM 训练管线

LLM 在推进科学研究方面展现出巨大潜力,助力完成论文总结、综合复杂见解和生成新颖假设等任务。然而,由于预训练期间对专业术语和上下文知识的接触有限,通用型 LLM 通常无法满足特定领域的任务需求。

为弥合这一差距,适应领域的 LLM 提供了一种更加可行的解决方案。域自适应预训练 (DAPT) 通过精心策划自定义和与领域相关的语料库来扩展现有基础模型 (例如 LLaMA) 的功能,而不是从零开始产生高昂的训练成本和计算需求。

这种方法显著提高了科学等专业领域的性能,同时保留了原始模型的广泛语言能力。此外,我们还对模型进行了微调,以增强其响应常规查询和特定任务查询的能力。虽然领域适应和指令调整可提高任务性能,但它们无法为模型配备推理能力。

为弥补这一差距,我们引入了推理对齐,使模型能够在逻辑上导航假设生成、思维链推理和自校正等过程。这些功能对于解决多步骤问题和推动材料探索至关重要。

SES AI 是一家专门从事电池创新的公司,他们构建了一个专有模型 — Molecular Universe LLM,这是一个具有 70B 参数的大型自定义推理模型。它基于 Llama 3.1 70B,为特定领域的科学任务设定了新的基准,优于同类别的其他模型。

它展示了一种计算高效的训练和对齐策略,可将基础模型转换为高性能、领域自适应的模型,同时强调了将 DAPT、instruction tuning 和 reasoning-based fine-tuning 结合起来用于特定领域任务的有效性。

Molecular Universe LLM 是一个由 AI 驱动的电池研究 LLM,利用高级推理对潜在的电解质溶剂和合成添加剂进行排序。以前,他们的科学家根据专业领域的专业知识对电解质溶剂和添加剂进行人工排名,这项工作仅限于每天评估几十个候选化合物。

通过整合长语境理解、结构化推理和专家级决策,这种方法凸显了领域适应性推理模型如何加速科学创新的突破,显著提高电池专家的工作效率。

Molecular Universe LLM 在 NVIDIA DGX Cloud 上使用 NVIDIA NeMo 框架 通过三步流程进行训练:

这种方法可确保模型提供特定领域、上下文相关的高质量响应。通过将 Molecular Universe LLM 与 NVIDIA ALCHEMI GPU 加速模拟数据和 NVIDIA cuML 生成的分子图集成, SES AI 将数十年的电池研究时间缩短到了几个月。

我们来深入探讨构建此模型所涉及的步骤。

A detailed workflow diagram illustrating the data pipeline and model training process for a scientific language model based on Llama 3.1.
图 2。用于训练 Molecular Universe 推理模型的 End-to-end 工作流程

基础架构设置

Molecular Universe LLM 在 NVIDIA DGX 云 上的 128 个 NVIDIA H100 GPU 上进行训练,NVIDIA DGX 云是一个与领先云提供商共同设计的全托管 AI 训练平台。DGX 云包含 NVIDIA 管理的 Kubernetes 和 NVIDIA Run:ai ,用于工作负载优化、作业调度和编排。开发者可以立即开始在专用集群上运行分布式训练,无需复杂的集群启动或管理底层基础设施。

NVIDIA NeMo 框架 作为基于 NVIDIA DGX 云的 AI 开发平台,可提供无缝的加速体验,以高效地大规模构建、定制和部署生成式 AI 模型。它支持先进的模型和算法,同时通过 4D 并行和其他优化确保数千个 GPU 的高训练吞吐量和可扩展性。

借助 NVIDIA Run:ai,管理员能够通过“项目”和“部门”编排 GPU 容量,确保为团队分配所需的容量份额,以处理训练工作负载。调度程序还支持工作负载突发,使工作负载能够在集群有额外资源可用时利用额外的 GPU 容量。这可提高 GPU 利用率,同时尊重资源分配、更大限度地提高开发者工作效率并更大限度地缩短价值实现时间。

Screenshot of the run:ai Workloads dashboard showing a PyTorchJob named 'mixtral-pretrain' running for 2 days, 20 hours, and 53 minutes.
图 3。NVIDIA DGX 云上运行的预训练作业期间的 GPU 使用示例

第 1 步:持续预训练

为了在电池研究中建立特定领域知识的坚实基础,对 Llama3.1 70B 模型进行了持续的预训练。这需要在大量精心策划的科学文献语料库上进行训练,使模型能够获得细致入微的理解和专业知识,这对于准确、上下文感知的响应至关重要。

数据管护和处理

预训练语料库包含来自同行评审期刊或预打印资源库的 19M 篇开源论文。有关数据源的详细信息,请参阅表 1。

来自不同来源的 PDF 已转换为纯文本。在训练之前,我们使用 NeMo Curator 提取和处理文档,NeMo Curator 应用了先进的启发式过滤和 GPU 加速的模糊重复数据删除技术,包括 MinHash 和 Locality Sensitive Hashing。这种严格的工作流将原始样本的数量减少了 19M 个,达到 17M 个独特的高质量记录。NeMo Curator 的预处理能力对于消除冗余、过滤低质量数据以及保留丰富的特定领域知识至关重要。

数据源 文档
开源同行评审文献 约 400 万
arXiv 140 万
ChemRxiv 2.6 万
开放研究 1200 万
PubChem 6 万
学术教科书或专著 80
PLOS 20 万
表 1。用于 Domain-Adaptive Pre-training 的数据源细分

模型架构和训练细节

Molecular Universe LLM Base 模型通过调整 LLaMA 3.1 70B 基础模型的预训练权重来构建。NeMo 框架用于模型的持续预训练,利用先进的优化技术,包括 4D 并行、混合精度训练和 Flash Attention。此外,NeMo 上下文并行在使模型能够在不影响内存效率、速度或稳定性的情况下处理多达 8K 个 token 的长序列方面发挥了关键作用。

该模型的输入序列长度为 8192 个 token,每次向前传递时处理 524288 个 token。他们使用 128 个 NVIDIA H100 GPU 训练该模型,在 bfloat16 精度下,总训练时间为 144 小时。域自适应预训练 (DAPT) 仅在原始预训练中使用的部分 token 上执行,而且效率要高得多,只需约 1.5% 的预训练总计算量。

在初始步骤中,训练和验证损失曲线呈现快速下降趋势,反映出领域适应迅速。随着时间的推移,损失趋于稳定,表明收敛是有效的,没有过拟合的迹象。

第 2 步:通过监督式微调进行模型对齐

为了使 Molecular Universe Base 模型与特定领域的知识保持一致,并提高指令遵循能力,我们采用了 Supervised Fine-Tuning (SFT) 。Supervised Fine-Tuning (SFT) 基于已标记示例训练模型,以改进指令遵循和特定于任务的响应生成,尤其是在特定领域的上下文中。

数据管护和处理

SES 利用 NVIDIA Llama 3.1 70B NIM 生成合成数据 (SDG) ,以创建高质量的 SFT 数据集。他们对 50,000 篇论文进行了采样,并在四个任务 (问答、总结、阅读理解和多项选择题) 中生成了 200,000 个指令样本,其中 160,000 个用于训练,40,000 个用于评估。

最终的 SFT 数据集整合了来自 Daring-Anteater 数据集的 9 万个常规聊天样本,总计 25 万个样本,其中 SDG 提供了大多数,这凸显了使用 NIM 生成特定领域训练数据的有效性。

模型架构和训练细节

然后,使用 LLaMA 3.1 70B 标记器对该数据集进行标记化,然后使用 NeMo 框架在多节点系统上使用 SFT 进行微调,从而生成最终的 Molecular Universe Chat 模型。该模型在 DGX Cloud 上使用 128 个 NVIDIA H100 GPU 和 NVIDIA Run:ai 软件进行训练,只需 32 小时即可完成。

训练和验证损失曲线初始下降迅速,稳定在 400 步左右。经过 600 个步骤后,训练损失略有增加,这表明可能存在对学习率的过拟合或敏感性。但是,验证损失保持稳定,表明泛化性能强劲。

第 3 步:使用高质量推理数据进行后训练

虽然基于科学文献的领域自适应预训练和基于指令的微调可增强模型处理一般和特定领域问题的能力,但它们并不擅长解决需要多步骤推理的复杂科学问题。

为了克服这一问题,我们在 s1K Reasoning Data 中的精选样本集 (~25,000 个样本) 上对 Molecular Universe Chat 模型进行了微调。s1K Reasoning Data 包含高质量的难题,并提供了 Gemini Thinking 的推理跟踪和解决方案。s1K 数据集经过筛选,可删除带有格式问题和问题的低质量样本,这些问题可由基础模型轻松回答,例如 Qwen2.5 7B Instruct 和 Qwen2.5 32B Instruct。

此外,还使用了 LLM 将样本聚类为主题类别 (例如 math, science) ,并在应用统一采样时偏向于包含较长推理轨迹的示例,以更好地捕获任务复杂性。对生成的数据样本进行了进一步的去污处理,以删除特定任务的基准测试,例如 GPQA Diamond。

使用 NeMo 框架完成了训练后监督微调,方法是将上下文长度增加到 16k,以考虑推理轨迹。这一步骤在 64 个 H100 GPU 上花费了大约 12 小时进行 5 次训练,不仅提高了事实准确性,而且增强了模型通过复杂想法进行推理的能力,在 GPQA Diamond 上的得分为 0.72。

结果

Molecular Universe Chat and Reasoning 模型在以科学为重点的公共基准 (例如 GPQA Diamond ) 和自定义领域特定基准上进行了评估。它在 GPQA Diamond 上获得了 0.72 分,超越了大多数其他著名、类似规模甚至更大的开源模型,例如 DeepSeek-R1。

在 MMLU、Winogrande、Hellaswag 和 ARC-E 等公共基准测试中,Molecular Universe Reasoning 模型的性能优于 LLama 3.1 70B。从基础起始模型中获得的显著性能提升凸显了持续的域预训练和推理驱动的后训练在提升模型能力方面的价值,而不仅仅是指令对齐。

A bar chart titled "Performance comparison between different SOTA reasoning models on the GPQA."
图 4。GPQA 上不同 SOTA 推理模型的性能比较
模型 # 个参数 电池问答 Battery MCQ Battery RC Battery 摘要 电池推理
GPT-o1 96% 92% 90% 88% 84%
分子宇宙推理 700 亿 96% 89% 90% 86% 82%
克劳德 3.7 十四行诗 94% 86% 89% 86% 80%
Gemini 闪存思维 92% 85% 88% 82% 79%
Molecular Universe 聊天 700 亿 93% 79% 84% 79% 73%
LLaMA 3.1 700 亿 71% 67% 78% 75% 66%
表 2。电池特定任务的性能比较,包括 Q/A、MCQ、阅读理解、总结和推理

Molecular Universe Chat 和推理模型在 40,000 个 SFT 测试集和自定义电池特定推理基准测试中得到了进一步评估。将该模型与 GPT-o1、LLaMA 3.1 70B、Claude 3.7 Sonnet 50B 和 Gemini 60B 等模型进行了比较。

在问答、MCQ、阅读理解、总结和推理等任务中,Molecular Universe Reasoning LLM 的表现始终优于 GPT-o1 以外的所有基准。尽管 GPT-o1 在生成微调数据方面发挥着重要作用,因此处于领先地位,但 Molecular Universe Reasoning 仍凭借更少的参数和更低的训练成本取得了有竞争力的结果,进一步凸显了领域适应和推理对齐的影响。

结论和未来工作

分子宇宙推理 (Molecular Universe Reasoning),一个 70B 参数的科学推理 LLM,在其规模类别中展示了科学任务的先进性能。高效计算的训练策略将域自适应预训练与基于推理的监督微调相结合,显著提高了基准模型的性能,并将额外的计算成本降至最低。

事实证明,结合使用这两种技术很有价值,其性能优于单独使用的任何一种方法,并在一般和电池特定基准测试中取得了与更大模型相美的结果。 Molecular Universe Reasoning 模型的部署使用了 NIM 微服务 支持微调模型 ,实现了模型的可扩展实时服务,允许最终用户同时发送多个并发请求。Molecular Universe LLM 将集成到 SES AI 的材料发现平台 Molecular Universe (MU-0) 中。该平台是一个统一的软件和服务解决方案,旨在帮助电池研究人员和行业专业人员通过一个统一的搜索界面探索候选小分子的庞大数据库。

未来的工作将涉及通过特定领域的推理后训练来完善模型,特别是通过构建以 battery 为中心的专用数据集来增强任务相关推理,并探索使用人类反馈进行强化学习,以进一步提高特定领域的性能。这项工作说明了在不同领域开发经济高效的中型 (<100B) 领域专家模型的路径,这些模型具有强大的专门化能力。

如需详细了解 NVIDIA DGX 云上的 NeMo 框架,请访问 NVIDIA 官方文档和 GitHub。立即开始使用 NVIDIA DGX 云。探索 NVIDIA ALCHEMI 并探索适用于先进机器学习解决方案的 NVIDIA cuML。

感谢 Zihan Wang (NVIDIA) 和 Kang Xu (SES) 的宝贵支持和见解。

 

标签