【论文阅读】英伟达ChipNeMo，一个领域大模型的训练案例

ChipNeMo是英伟达针对芯片设计训练的一个领域大模型，使用了微调和RAG的组合技术，论文有针对技术细节的大量对比分析，对于理解微调和RAG技术，是一个非常好的技术案例。

先说要点：

1、领域适配大模型的优势：针对领域需求的性能更好，安全，部署成本低。

2、训练成本方面：领域适配预训练成本，对比从头开始预训练基础模型的总成本不到 1.5%。

3、整个训练过程包括：领域适配分词、领域适配继续预训练、模型对齐（Model Alignment）、检索增强生成，这些步骤一个都不能少。

√领域适配分词：提升专业领域的分词效率，提升对领域的特殊Token的识别。

√领域适配继续预训练：向大模型注入领域知识，是整个方案的重要步骤。

√模型对齐：继续预训练会导致模型丢失原有的在下游任务的能力，需要重新SFT。

√检索增强生成：单纯大模型的性能还不够好，需要结合RAG；另一方面，单靠RAG的性能也还是不行，需要针对领域进行继续预训练，更好的注入领域知识，将领域模型和RAG相结合。

4、领域适配继续预训练，对比PEFT的微调技术，有巨大性能优势。

5、RAG模型，需要领域适配检索模型，使用领域数据微调预训练的检索模型后，检索命中率会显著提高。

6、模型的训练对数据的规模、质量和多样性都有要求，对于数据的多样性，合理的数据比例非常重要，即重采样策略。

《ChipNeMo: Domain-Adapted LLMs for Chip Design》这篇来自英伟达的论文，介绍了一个适用于芯片设计的领域大模型——ChipNeMo。文章比较全面而详细的介绍了适用于特定领域的大模型的训练过程，通过此文，可以一窥大模型训练之究竟。虽然英伟达是针对内部的芯片设计，但是论文阐述的训练过程，对于其他领域的大模型训练无疑有很好的参考价值。

文章发表于去年11月，在arXiv上最新的版本更新到2311.00176v4 [cs.CL] 07 Mar 2024，相比去年的版本，还是有比较多的变化。

摘要
ChipNeMo旨在探索大型语言模型（LLMs）在工业芯片设计中的应用。我们没有直接部署现成的商业或开源 LLMs，而是采用了以下领域自适应技术：领域适配分词、领域适配的继续预训练、模型与特定领域指令的对齐以及领域适配的检索模型。我们在三个选定的芯片设计 LLM 应用程序上评估了这些方法：工程助理聊天机器人、EDA 脚本生成以及错误摘要和分析。我们的评估表明，与基本的 LLaMA2 模型相比，语言模型的领域适配预训练可以在与领域相关的下游任务中带来卓越的性能，而不会降低通用能力。特别是，我们最大的模型 ChipNeMo-70B 在我们的两个用例（即工程助理聊天机器人和 EDA 脚本生成）上优于功能强大的 GPT-4，同时在错误摘要和分析方面表现出有竞争力的性能。这些结果强调了特定领域定制对于提高大语言模型在专业应用程序中的有效性的潜力。

ChipNeMo的模型功能
基于开源模型进行训练后，针对芯片设计任务中许多与语言相关的功能，通过生成式AI模型构建自动化工具，如代码生成、回答工程问题、分析和报告生成，以及错误分类，从而帮助提高芯片设计效率。具体来说，论文针对三个选定的应用场景进行了评估：

-工程助理聊天机器人：为GPU ASIC和架构设计工程师提供帮助，能够理解内部硬件设计并解释复杂的设计主题。

-EDA脚本生成：基于Python和Tcl两个特定领域工具，为VLSI定时分析任务生成脚本。

-bug总结和分析：作为内部错误和问题跟踪系统的一部分，对bug进行总结和分析。

文章的主要观点和方法
针对特定领域训练领域大模型的必要性在于，通过领域适配预训练的LLM模型可以在特定领域任务上胜过通用模型，并且这种定制的LLM避免了通过API将专有芯片设计数据发送给第三方LLMs相关的安全风险。但是，为每个领域从头开始训练一个特定领域的模型成本高昂，论文选择合适的基础模型进行领域适配预训练（DAPT，Domain-Adaptive PreTraining），提出了一系列领域适配技术，包括：

1. 领域适配分词（Domain-Adaptive Tokenization）：通过定制分词器来提高对芯片设计数据的分词效率，同时保持对通用数据的处理能力。

2. 领域适配预训练（Domain Adaptive Pretraining, DAPT）：在现有的基础模型上，使用特定领域的数据进行继续预训练，以增强模型对芯片设计相关知识的理解。

3. 模型对齐（Model Alignment）：通过监督式微调（Supervised Fine-Tuning, SFT）和SteerLM等技术，使模型更好地适应特定的任务，如聊天机器人对话和EDA脚本生成。

4. 检索增强生成（Retrieval-Augmented Generation, RAG）：结合检索机制和语言模型生成，以提供更准确的上下文信息，从而生成更准确的回答。

文章评估了这些方法在三个选定的应用场景中的效果：工程助理聊天机器人、EDA脚本生成和错误总结与分析。通过与通用LLMs的比较，研究表明领域适配预训练的模型在特定任务上表现更优，且没有损失通用能力。

在方法论上，文章采用了以下步骤：

• 数据准备：收集和预处理大量与芯片设计相关的内部数据，包括硬件代码、自然语言文档等，并结合了一定比例的公共数据集，以保持模型的通用性。

• 模型训练：使用NVIDIA NeMo框架进行模型训练，应用了如张量并行（tensor parallelism）和flash attention等技术来提高效率。

• 评估和分析：通过自动评估和人类评估的方式，对模型在多个领域特定任务上的性能进行了详细的评估。此外，还进行了多项消融研究，以理解不同领域适配技术对模型性能的具体影响。

文章的结论是，领域适配预训练和其他技术显著提高了LLMs在芯片设计任务上的性能，并且这种定制化的方法对于提高大型语言模型在专业应用的有效性具有潜力。未来的工作将集中在进一步改进ChipNeMo模型和方法，以便在生产环境中使用。

总体流程
论文基于LLaMA2模型的多种变体（LLaMA2 7B/13B/70B）作为基础模型，训练了多个ChipNeMo 基础和聊天模型。

图1展示了训练流程，这里直观的展示出基础模型的预训练、继续预训练（DAPT）和监督式微调（SteerLM/SFT）所需的训练时长和训练数据量需求。

√DAPT，领域适配预训练，就是使用域内数据进行继续预训练，已被证明在生物医学和计算机科学出版物、新闻和评论等领域是有效的。论文中，从专有的硬件相关代码（例如软件、RTL、验证测试平台等）和自然语言数据集（例如硬件规范、文档等）的集合中构建特定于领域的预训练数据集。对于领域数据，需要清理和预处理，然后继续预训练基础模型。DAPT 相比于基础模型的预训练使用的数据少很多，所以成本要低得多，只需要大约 1.5% 的预训练计算。

√领域适配分词器，通过为特定领域的术语（例如 RTL 中常见的关键字）定制规则和模式来提高分词化效率。对于 DAPT，无法从头开始重新训练新的特定领域的分词器，因为这会使基础模型无效。另一方面， ChipNeMo 也不能限制为基础模型使用的预训练通用分词器，需要将预训练的分词器调整到与芯片设计数据集相适应，为特定领域的术语添加新Token。

√SteerLM/SFT：DAPT之后，需要对齐模型以适应聊天等下游任务。论文使用大量公开的通用聊天指令数据集进行多轮聊天，以及少量特定于领域的指令数据集，以对 ChipNeMo 基础模型进行对齐，生成 ChipNeMo 聊天模型。论文分析，与通用聊天指令数据集的一致性足以使 ChipNeMo 基础模型与芯片设计领域的查询保持一致。论文添加了少量特定于任务的指令数据，这进一步改善了对齐方式。

√为了提高工程助理聊天机器人应用程序的性能，还需要利用检索增强生成 （RAG）。RAG 是一种开卷方法，用于为 LLM 提供用户查询的精确上下文。RAG从其数据存储中检索相关的域内知识，以增强给定用户查询的响应生成。这种方法在将模型建立在特定问题的上下文方面有了显著的改进。至关重要的是，论文分析，当使用领域数据微调预训练的检索模型后，检索命中率有了显著提高。这导致了模型质量的进一步提高。

ChipNeMo表现出卓越的性能，有以下要点：

•在三个用例上展示了领域适配 LLM 有效性：工程助理聊天机器人、EDA 工具脚本生成以及错误摘要和分析。在7分制（Likert scale）的工程助理聊天机器人上获得了6.0分，在简单EDA脚本生成上的正确率超过70%，在总结和任务识别任务上，专家评估得分在7分制上高于5分。

• ChipNeMo 模型的性能在多项选择特定域AutoEval 基准测试和应用程序人工评估中，显著优于所有通用 LLM。

•使用 SteerLM 对齐方法与传统的 SFT 相比，在 7 分制（Likert scale）上，工程助理聊天机器人的人工评估分数提高了 0.62 分。

• SFT 通过额外的1.4K特定于领域的指令，显著提高了模型生成EDA 工具脚本的正确度，提高了18%。

•领域适配分词可将域数据Token数量减少多达3.3%，不会损害应用程序的有效性。

•使用特定领域的数据微调ChipNeMo检索模型，与预先训练的最先进的检索器相比，检索器命中率提高了 30%，从而提高了 RAG 响应的整体质量。

模型训练

√DAPT：图2显示了 ChipNeMo 在指定超参数下的训练损失。图中可以观察到训练损失的峰值。论文分析这些峰值可以归因于“坏数据”，因为这些不规则性似乎始终发生在同一模型的类似训练步骤中，即使在不同的模型大小中也是如此。论文选择不解决这个问题，因为这些异常似乎并没有阻碍后续的训练步骤（验证损失没有明显的下降），这可能是由于应用了低学习率。

√模型对齐：数据集有大约 1.4k 个样本组成的领域对齐数据集与更大的通用聊天数据集相结合。对于 SFT，将领域指令数据与 128k 商用用途的聊天数据混合在一起，然后在随机洗牌后对单个 epoch 进行微调。针对SFT，论文实验了特定领域数据集的扩充，但是对域内问题模型会出现过拟合。

√对于 SteerLM，在其他论文有详细介绍。对于模型对齐技术，RLHF过于复杂，SteerLM是一种相对简单有效的方法。

√论文试验了在聊天对齐模型（例如 LLaMA2-Chat 模型）上进行 DAPT，发现 DAPT 显著降低了模型的对齐性能，使得生成的模型对下游任务毫无用处。

评估

√ChipNeMo 模型对比两个基线通用模型（LLaMA2-70B-Chat和GPT-4），评估重点归纳如下：

• DAPT对领域本身的任务产生了实质性的积极影响。这种效果体现在内部设计知识和一般电路设计知识的显著改进上。

• DAPT 模型在开放领域学术基准上的性能略有下降。

•使用更大、性能更高的基础模型可以在特定领域的任务上产生更好的零样本结果。此外，使用卓越的基础模型会导致 DAPT 之后的领域模型得到增强，从而提高域内任务的性能。

•归因于域内任务的 DAPT 改进与模型大小呈正相关，较大的模型在特定领域的任务性能方面表现出更明显的增强。

√成本

模型使用 128 个 A100 GPU 进行了训练。ChipNeMo 的领域适配预训练成本，对比从头开始预训练基础模型的总成本不到 1.5%。具体见下表，单位是A100 GPU hours

√RAG

针对检索增强，论文提到：“即使在对检索模型进行微调带来了显著收益的事实面前，检索仍然在无法直接映射到文档语料库中的段落的查询或需要段落中不存在的更多上下文的查询方面存在困难，将检索与领域适应的语言模型相结合是解决这个问题的一种方法”。换句话说，单纯靠RAG的性能还是不行，需要针对领域进行继续预训练，更好的注入领域知识，将领域模型和RAG相结合。

领域适配RAG模型的性能比原始 e5_small_unsupervised 模型好 2 倍，比sentence transformer好30%。

上图对不同模型的人工评估。“仅模型”表示不带 RAG 的结果。RAG（命中）/（未命中）仅包括检索到的段落命中/未命中理想上下文的问题，RAG（平均）包括所有问题。对多种 ChipNeMo 模型和带或不带 RAG 的 LLaMA2 模型进行了评估。然后由人类评估对结果进行评分，重点包括：

•ChipNeMo-70B-Steer 在所有类别中都优于 GPT-4，包括 RAG 未命中和命中。

•ChipNeMo-70B-Steer 在仅模型评估和 RAG 评估中的表现分别优于类似尺寸的 LLaMA2-70b-Chat。

结果表明，RAG显著提高了人类得分。RAG 将 ChipNeMo-70B-Steer、GPT-4 和 LLaMA2-70b-Chat 分别提高了 0.56、1.68 和 2.05。即使 RAG 未命中，分数通常也高于不使用检索。包含相关的域内上下文仍然提高了性能，因为检索不是严格意义上的二元结果。此外，虽然 ChipNeMo-70B-SFT 通过传统的监督微调大幅优于 GPT4，但应用 SteerLM 方法让聊天机器人性能进一步提高。

√EDA脚本生成

两种不同类型的基准测试。第一种是“简单”和“中等”难度任务；第二种任务是“困难”任务，来自特殊选择的真实用例场景。


图 8 显示“简单”和“中等”基准测试的结果，在这些基准测试中，检查代码是否完全准确。对于图 9 中的“硬”基准测试，检查代码的部分正确性，由人类用户以 0-10 的等级进行评估。ChipNeMo-70B-Steer 的性能明显优于现成的 GPT-4 和 LLaMA2-70B-Chat 模型。

如图 8 所示，GPT-4 和LLaMA2-70B-Chat 等模型对 Python 工具的准确度接近于零，其中与工具的 API 相关的领域知识是必要的。由此可见DAPT的重要性。如果没有 DAPT，该模型对底层 API 几乎一无所知，并且在自动和人工评估的基准测试中表现不佳。对齐模型进一步改进了 DAPT 的结果，因为领域指令数据有助于指导模型以最有用的方式呈现最终脚本。

从图中看出，非领域模型在 Tcl 工具上的表现比 Python 工具更好，但领域模型的趋势恰恰相反。论文怀疑这是由于 Python 工具的专有性质。如果不了解 API，一般 LLM 很难在专业领域的 Python 工具基准测试中表现良好。由于 ChipNeMo 是使用领域数据训练的，因此基础模型固有的 python 编码能力使 ChipNeMo-70B-Steer 性能更好。这再次凸显了 DAPT 对于小容量或专有编程语言的重要性。

√Bug 总结和分析

研究了三个独立的子任务：以技术细节为重点的总结，以管理细节为重点的总结，以及任务分配的总结建议。针对这三项任务中的每一项对模型的性能进行评分，结果如图 10 所示。尽管 GPT-4 模型在所有三项任务中都表现出色，优于的 ChipNeMo-70B-Steer 模型和 LLaMA2-70B-Chat 模型，但与同等规模的现成 LLaMA 模型相比，ChipNeMo-70B-Steer 确实表现出增强功能。论文将领域适应导致的摘要任务的相对较低的改进归因于与其他用例相比，摘要任务对特定领域知识的需求有限。

数据集
针对芯片设计的专用用途，论文使用的领域特定预训练数据集包括专有的与硬件相关的代码（例如软件、RTL、验证测试集等）和自然语言数据集（例如硬件规格、文档等）。需要注意论文中针对继续预训练数据的两个处理方式：

1、 除了专用数据，继续预训练同时需要一定比例的公共数据，论文使用公共数据集包括自然语言和代码，自然语言部分从维基百科中提取数据，因为其数据的高质量，代码利用了 GitHub 数据。后面的消融实验指出：仅移除公共数据集在大多数任务上只导致了轻微的退步，包括学术基准测试，但在 Verilog 编码方面有明显的差异。

2、 数据多样性：模型的训练对数据的规模、质量和多样性都有要求，对于数据的多样性，合理的数据比例非常重要，即重采样策略。论文提到，训练周期内对代码数据进行了下采样，同时对自然语言数据进行了上采样，特别是设计文档。

SFT阶段的数据，前面已经有描述。在对齐阶段，使用领域的特定指令数据集，以将模型与下游用例对齐，这里使用的数据量非常小。

最后，为了定量评估模型的准确性，需要构建评估标准，这里包括公开的 LLM 基准数据，还有专门构建的涵盖模拟和数字设计主题的通用电路设计知识的额外基准。

与PEFT的对比
论文比较了应用参数高效微调（PEFT）的可能性，使用LoRA 适配器进行了两种参数预算的微调，实验表明与全面参数 DAPT 方法相比，在领域内任务准确性上存在显著差距。图13展示了基于LoRA和全参数DAPT的对比。

论文链接：https://arxiv.org/html/2311.00176v4

(更多论文阅读，请关注微信公众号：老无谈藏书票)