立即下载 DOCA,开启高性能AI 网络之旅,实现“一站式” 编程
数据中心/云端

Mistral-NeMo-Minitron 8B 基础模型实现准确性巅峰

上个月,NVIDIA 和 Mistral AI 推出了 Mistral NeMo 12B,这是一款领先的先进的大型语言模型(LLM)。Mistral NeMo 12B 在各种基准测试中的表现始终优于类似大小的模型。

今天,我们宣布推出 Mistral-NeMo-Minitron 8B,这是该体量级别中最先进的开放访问模型之一。该模型在9个热门基准测试中始终提供领先的准确度。Mistral-NeMo-Minitron 8B基础模型是通过对Mistral NeMo 12B基础模型的宽度剪枝获得的,然后使用知识提炼进行轻量再训练。这是NVIDIA最初在论文’通过剪枝和知识提炼实现紧凑语言模型‘中提出的成功方法。NVIDIA Minitron 8B和4B以及Llama-3.1-Minitron 4B模型反复验证了这一点。

  训练令牌 Wino-Grande 5-shot ARC
Challenge 25-shot
MMLU 五步 Hella
Swag 10-shot
GSM8K 5-shot TruthfulQA 0-shot XLSum en (20%)
3-shot
MBPP
0-shot
Human
Eval
0-shot
Llama 3.1 8B 15T 77.27% 57.94 65.28 81.80 48.60 欧元 45.06 30.05 42.27 24.76
Gemma 7B 6T 78 61 64 82 50 45 17 39 32
Mistral-NeMo-Minitron 8B 3800 亿 80.35 64.42 69.5% 83.03 58.45 47.56% 31.94 43.77 36.22
Mistral NeMo 12B 不适用 82.24 65.10 68.99% 85.16% 56.41 49.79% 33.43 42.63 23.78
表 1.Mistral-NeMo-Minitron 8B 基础模型与教师 Mistral-NeMo 12B、Gemma 7B 和 Llama-3.1 8B 基础模型相比的准确度。粗体数字表示在 8B 模型类别中取得的最佳结果。

模型剪枝和提炼概述

模型剪枝是指通过减少图层(深度剪枝)或减少神经元、注意力头和嵌入通道(宽度剪枝)来缩小模型并使其更加精简的过程。剪枝通常伴随一定数量的重新训练,以恢复准确性。

模型提炼是一种技术,用于将知识从复杂的大型模型(通常称为教师模型)传输到更小、更简单的学生模型。其目标是创建更高效的模型,该模型保留原始大型模型的大部分预测能力,同时运行速度更快,资源密集程度更低。在本文中,我们采用模型提炼作为剪枝后的轻量再训练过程,其数据集比从头开始训练模型时使用的数据集要小得多。

迭代剪枝和提炼是一种方法,从单个预训练模型开始,可以获得多个逐渐缩小的模型。例如,可以对 15B 模型进行剪枝和提炼,以获得 8B 模型,而这反过来又可以作为剪枝和提炼 4B 模型的起点,以此类推。

模型剪枝与通过蒸进行光线重新训练相结合,是一种经济高效的模型训练方法。对于每个新增的模型,只需使用 100-400 亿个令牌进行重新训练,与从零开始训练相比,减少了 40 倍以上。因此,与从零开始训练所有模型相比,训练一系列模型 (12B、8B 和 4B) 可节省高达 1.95 倍的计算成本。

从广泛的消融研究中学习已总结为结构化权重剪枝与知识提炼相结合的 10 种最佳实践。我们发现,宽度剪枝的性能始终优于深度剪枝;最重要的是,经过剪枝和提炼的模型在质量上优于从零开始训练的模型。

Mistral-NeMo-Minitron 8B

按照我们的最佳实践,我们对 Mistral NeMo 12B 模型进行了宽度剪枝,以获得 8B 目标模型。本节将详细介绍获取 Mistral-NeMo-Minitron 8B 基础模型的步骤和参数,以及该模型的性能。

教师微调

为了修正训练模型所用的原始数据集上的分布偏移,我们首先使用 127B 个令牌在数据集上微调了未剪枝的 Mistral NeMo 12B 模型。实验表明,若不修正分布偏移,教师在提炼数据集时将提供次优指导。

仅宽度的剪枝

鉴于我们的目标是尽可能获得最强的 8B 模型,我们继续进行仅宽度的剪枝。我们沿着宽度轴剪枝了嵌入(隐藏)和 MLP 中间维度,以压缩 Mistral NeMo 12B。具体来说,我们使用基于激活的策略计算了每个注意力头、嵌入通道和 MLP 隐藏维度的重要性分数。遵循重要性估计,我们:

  • 将 MLP 中间维度从 14336 剪枝为 11520
  • 将隐藏大小从 5120 剪枝为 4096
  • 保留了’Attention Head’数量和层数

蒸参数

我们使用峰值学习率 = 1e-4、最低学习率 = 4.5e-7、60 个步骤的线性预热、余弦衰减计划和 3800 亿个令牌的 768 全局批量大小(与教师微调中使用的数据集相同)来提炼模型。

结束语

Mistral-NeMo-Minitron 8B 提供业界领先的精度,并且性能始终优于最近推出的类似大小的先进模型。Mistral-NeMo-Minitron 8B 是我们在 Mistral NeMo 12B 模型提炼方面的首个成果,它为我们结合知识提炼最佳实践的结构化权重剪枝提供了有力的支持。我们计划进一步对模型进行提炼,并获得更小、更准确的模型。该技术的实现将逐步在 NVIDIA NeMo 框架中推出,用于生成式 AI。

如需了解详情,请查看以下资源:

致谢

如果没有 NVIDIA 的许多人的贡献,这项工作是不可能完成的。仅举几例:

基础模型:Sharath Turuvekere Sreenivas、Saurav Muralidharan、Raviraj Joshi、Marcin Chochowski、Pavlo Molchanov、Mostofa Patwary、Daniel Korzekwa、Ashwath Aithal、Mohammad Shoeybi、Bryan Catanzaro 和 Jan Kautz 对齐:Ameya Sunil Mahabaleshwarkar、Hayley Ross、Brandon Rowlett、Oluwatobi Olabiyi、Shizhe Diao 和 Yoshi Suhara 数据集:Sanjeev Satheesh、Jupinder Parmar、Shengyang Sun、Jiaqi Zeng、Zhilin Wang、Yi Dong、Zihan Liu、Rajarshi Roy、Wei Ping、Makesh Narsimhan Sreedhar 和 Oleksii Kuchaiev TensorRT-LLM:Bobby Chen、James Shen 和 Chenhan Yu Hugging Face 支持:Ao Tang、Yoshi Suhara 和 Greg Heinrich

 

标签