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}, 
}