Minitron 8B Base

🚀 Minitron-8B-Base大语言模型

Minitron-8B-Base是一个大语言模型（LLM），它通过对Nemotron-4 15B进行剪枝得到。该模型在训练效率和性能上有显著优势，能节省大量计算成本，且在多项评估中表现出色。本模型仅用于研究和开发。

🚀 快速开始

对该模型的支持将在即将发布的transformers版本中添加。在此期间，请从源代码安装该库：

pip install git+https://github.com/huggingface/transformers

以下代码展示了如何加载Minitron-8B模型并使用它进行文本生成：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the tokenizer and model
model_path = "nvidia/Minitron-8B-Base"
tokenizer = AutoTokenizer.from_pretrained(model_path)

device='cuda'
dtype=torch.bfloat16
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)

# Prepare the input text
prompt = "To be or not to be,"
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)

# Generate the output
output_ids = model.generate(input_ids, max_length=50, num_return_sequences=1)

# Decode and print the output
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(output_text)

✨ 主要特性

• 高效训练：与从头开始训练相比，使用我们的方法从基础的15B模型派生Minitron 8B和4B模型，每个模型所需的训练令牌最多可减少40倍，全模型系列（15B、8B和4B）的训练计算成本节省1.8倍。
• 性能出色：Minitron模型在MMLU分数上比从头开始训练提高了16%，性能与其他社区模型（如Mistral 7B、Gemma 7B和Llama - 3 8B）相当，并且优于文献中的最先进压缩技术。

📦 安装指南

支持该模型的功能将在即将发布的transformers版本中添加。在此期间，你可以从源代码安装该库：

pip install git+https://github.com/huggingface/transformers

💻 使用示例

基础用法

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 加载分词器和模型
model_path = "nvidia/Minitron-8B-Base"
tokenizer = AutoTokenizer.from_pretrained(model_path)

device='cuda'
dtype=torch.bfloat16
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)

# 准备输入文本
prompt = "To be or not to be,"
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)

# 生成输出
output_ids = model.generate(input_ids, max_length=50, num_return_sequences=1)

# 解码并打印输出
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(output_text)

📚 详细文档

模型架构

Minitron-8B-Base使用的模型嵌入大小为4096，有48个注意力头，MLP中间维度为16384。它还使用了分组查询注意力（GQA）和旋转位置嵌入（RoPE）。

数据集与训练

评估结果

• 5次射击性能：使用大规模多任务语言理解评估语言理解能力： | 平均 | | :---- | | 64.5 |

• 零次射击性能：使用LM评估工具包中的选定数据集进行评估： | HellaSwag | Winogrande | GSM8K | ARC - C | XLSum | | :------------- | :------------- | :------------- | :------------- | :------------- | | 81.6 | 80.3 | 54.2 | 49.2 | 31.1 |

• 代码生成性能：使用HumanEval评估： | p@1, 0 - 次射击 | | :------------- | | 31.6 |

更多评估结果请参考论文。

推理

局限性

该模型在包含从互联网爬取的有毒语言、不安全内容和社会偏见的数据上进行训练。因此，模型可能会放大这些偏见并返回有毒回复，尤其是在输入有毒提示时。即使提示本身不包含任何明确的冒犯性内容，模型也可能生成不准确、遗漏关键信息、包含无关或冗余文本的答案，从而产生社会不可接受或不良的文本。

伦理考量

NVIDIA认为可信AI是一项共同责任，我们已经制定了政策和实践，以支持广泛的AI应用开发。当按照我们的服务条款下载或使用该模型时，开发者应与内部模型团队合作，确保该模型满足相关行业和用例的要求，并解决不可预见的产品滥用问题。请在此报告安全漏洞或NVIDIA AI相关问题。

🔧 技术细节

Minitron-8B-Base是通过对Nemotron-4 15B进行剪枝得到的大语言模型。具体来说，我们对模型的嵌入大小、注意力头数量和MLP中间维度进行了剪枝。剪枝后，我们使用940亿个令牌进行蒸馏式的持续训练，以得到最终模型；为此，我们使用了Nemotron-4 15B中使用的连续预训练数据语料库。

📄 许可证

Minitron-8B-Base根据NVIDIA开放模型许可协议发布。

📖 引用

如果你认为我们的工作有帮助，请考虑引用我们的论文：

@article{minitron2024,
      title={Compact Language Models via Pruning and Knowledge Distillation}, 
      author={Saurav Muralidharan and Sharath Turuvekere Sreenivas and Raviraj Joshi and Marcin Chochowski and Mostofa Patwary and Mohammad Shoeybi and Bryan Catanzaro and Jan Kautz and Pavlo Molchanov},
      journal={arXiv preprint arXiv:2407.14679},
      year={2024},
      url={https://arxiv.org/abs/2407.14679}, 
}