AMD-Llama-135m

AMD-Llama-135m项目介绍

AMD-Llama-135m是一个由AMD公司开发的小型语言模型，基于Llama2模型架构，使用AMD Instinct MI250加速器进行训练。这个项目旨在为开发者和研究人员提供一个高效、灵活的语言模型选择。

模型特点

1. 参数规模：AMD-Llama-135m拥有1.35亿参数，属于小型语言模型范畴。
2. 兼容性：该模型可以与Hugging Face的transformers库无缝集成，可以作为LlamaForCausalLM模型加载使用。
3. 通用性：模型使用与Llama2相同的分词器，使其可以作为Llama2和CodeLlama的推测解码草稿模型。
4. 灵活性：提供了通用版本（AMD-Llama-135m）和代码特化版本（AMD-Llama-135m-code）。

模型架构

AMD-Llama-135m的架构设计精巧：

• 12层transformer结构
• 768维隐藏层
• 12个注意力头
• 使用Swiglu激活函数
• 采用RMSNorm进行层归一化
• 使用RoPE（Rotary Position Embedding）作为位置编码
• 最大上下文窗口大小为2048个token

训练过程

模型的训练分为两个阶段：预训练和代码微调。

1. 预训练：

   • 数据集：使用SlimPajama和Project Gutenberg数据集，总计约6700亿个训练token。
   • 训练细节：使用AdamW优化器，余弦学习率调度，批量大小为1024，训练1个epoch。

2. 代码微调（仅适用于AMD-Llama-135m-code版本）：

   • 数据集：使用StarCoder数据集的Python部分，约200亿个训练token。
   • 微调细节：在预训练模型基础上进行，使用较小的学习率和批量大小，训练1个epoch。

模型评估

AMD-Llama-135m在多个NLP基准测试中表现出色，与同等规模的模型相比具有竞争力：

• 在SciQ、PIQA等任务上表现优异
• 在MMLU、ARC等挑战性任务上也有不错的表现

推测解码应用

AMD-Llama-135m-code可以作为CodeLlama-7b的草稿模型，用于推测解码：

• 在MI250 GPU上，可以实现最高2.80倍的吞吐量提升
• 在Ryzen AI CPU（使用NPU内核）上，可以实现最高2.98倍的吞吐量提升

使用方法

使用Hugging Face transformers库可以轻松加载和使用AMD-Llama-135m模型：

from transformers import LlamaForCausalLM, AutoTokenizer

model = LlamaForCausalLM.from_pretrained("amd/AMD-Llama-135m")
tokenizer = AutoTokenizer.from_pretrained("amd/AMD-Llama-135m")

# 使用模型生成文本
inputs = tokenizer("Tell me a story?\nOnce upon a time", add_special_tokens=False, return_tensors="pt")
tokens = model.generate(**inputs)
output = tokenizer.decode(tokens[0])

结语

AMD-Llama-135m项目为自然语言处理和代码生成任务提供了一个高效、灵活的小型模型选择。无论是作为独立模型使用，还是作为大型模型的辅助模型，AMD-Llama-135m都展现出了优秀的性能和广阔的应用前景。开发者和研究人员可以利用这个模型来探索更多有趣的NLP应用和优化技术。