AMD OLMo 1B

🚀 AMD-OLMo

AMD-OLMo 是由 AMD 在 AMD Instinct™ MI250 GPU 上从头开始训练的一系列 10 亿参数语言模型。所使用的训练代码基于 OLMo。

我们发布了预训练模型、监督微调模型和 DPO 对齐模型，具体如下：

• AMD-OLMo-1B：在 Dolma v1.7 的一个子集上进行预训练，该子集包含 13 万亿个标记。
• AMD-OLMo-1B-SFT：在 Tulu V2 数据集（第一阶段）上进行监督微调（SFT），然后在 OpenHermes-2.5、WebInstructSub 和 Code-Feedback 数据集（第二阶段）上继续微调。
• AMD-OLMo-1B-SFT-DPO：在 UltraFeedback 数据集上使用直接偏好优化（DPO）与人类偏好进行对齐。

🚀 快速开始

本项目提供了一系列语言模型，你可以按照以下步骤快速使用：

1. 安装依赖：参考后续的安装指南部分，安装必要的软件和库。
2. 运行示例代码：参考使用示例部分，使用提供的代码示例进行模型推理。

✨ 主要特性

• 高性能训练：使用 AMD Instinct™ MI250 GPUs 进行训练，具有较高的训练吞吐量。
• 多阶段微调：经过预训练、监督微调（SFT）和 DPO 对齐等多个阶段，提升模型性能。
• 多任务表现出色：在多个标准基准测试中表现优异，具有良好的常识推理、多任务理解和负责任 AI 能力。

📦 安装指南

在 AMD GPU 上运行 PyTorch

你可以使用 Docker Hub 中的 ROCm Docker 镜像在 AMD GPUs 上运行 PyTorch：

docker pull rocm/pytorch:latest
# 在 Docker 内部
pip install transformers

评估环境设置

# lm-eval-harness
git clone https://github.com/EleutherAI/lm-evaluation-harness
cd lm-evaluation-harness
pip install -e .

# AlpacaEval
pip install git+https://github.com/tatsu-lab/alpaca_eval
cd alpaca_eval
pip install -e .

# MT-Bench
git clone https://github.com/lm-sys/FastChat.git
cd FastChat
pip install -e ".[model_worker,llm_judge]"

训练环境设置

WORK_DIR="<path_to_your_working_directory>"
cd $WORK_DIR
# 克隆 OLMo 代码库
git clone https://github.com/allenai/OLMo.git --branch v0.3.0
cd OLMo
# 克隆包含重现模型训练文件的 AMD-OLMo 仓库
git clone https://huggingface.co/amd/AMD-OLMo

docker pull rocm/pytorch:latest
docker run -it --network=host --device=/dev/kfd --device=/dev/dri --group-add=video --ipc=host --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --shm-size 8G -v $WORK_DIR/OLMo:/OLMo -w /OLMo rocm/pytorch:latest

# 由于 Docker 已安装 ROCm PyTorch，移除第 17 行
sed -i '17d' pyproject.toml
pip install -e .[all]

💻 使用示例

基础用法

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("amd/AMD-OLMo-1B-SFT").to("cuda") # 移除 .to("cuda") 以在 CPU 上加载
tokenizer = AutoTokenizer.from_pretrained("amd/AMD-OLMo-1B-SFT")

prompt = "What is large language model?"
bos = tokenizer.eos_token
template = bos + "<|user|>\n{prompt}\n<|assistant|>\n"

input_text = template.format(prompt=prompt)
inputs = tokenizer([input_text], return_tensors='pt', return_token_type_ids=False).to("cuda")
outputs = model.generate(**inputs, max_new_tokens=1000, do_sample=True, top_k=50, top_p=0.95)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])

📚 详细文档

模型信息

训练硬件和参数

• 硬件：每个计算节点包含 4 个 AMD Instinct™ MI250 GPUs，预训练 AMD-OLMo-1B 使用 16 个节点。
• 训练吞吐量：12,200 标记/秒/ GPU。

模型架构

AMD-OLMo-1B 基于完全开源的 10 亿版本 OLMo-1B 的模型架构和训练设置，具体参数如下：

超参数设置

🔧 技术细节

评估框架

我们使用以下开源评估框架对模型进行评估：

• Language Model Evaluation Harness：用于评估常识推理、多任务理解和负责任 AI 基准。
• AlpacaEval：用于评估聊天模型的指令遵循能力。
• MT-Bench：用于评估聊天模型的多轮对话能力。

评估设置

# lm-eval-harness
HF_MODEL=amd/AMD-OLMo-1B-SFT-DPO
accelerate launch -m lm_eval --model hf \
    --model_args pretrained=$HF_MODEL,trust_remote_code=True \
    --tasks arc_easy,arc_challenge,hellaswag,piqa,boolq,sciq,winogrande,openbookqa,mmlu,gsm8k_cot,bbh_cot_fewshot,toxigen,truthfulqa,crows_pairs \
    --device cuda \
    --batch_size 32 \
    --output_path ./lm-eval-results/$HF_MODEL

训练数据准备

# 下载预训练数据集
DATA_DIR=./datasets/dolma
mkdir -p $DATA_DIR

PARALLEL_DOWNLOADS="<number_of_parallel_downloads>"
cat "AMD-OLMo/dolma_v1_7_subset.txt" | xargs -n 1 -P $PARALLEL_DOWNLOADS wget -q -P $DATA_DIR

# 准备预训练数据集
NUM_WORKERS="<number_of_workers>"
python scripts/prepare_memmap_dataset.py $DATA_DIR/*.json.gz -o $DATA_DIR/memmap_dataset --workers $NUM_WORKERS

# 下载并准备 SFT 数据集
# 第一阶段 SFT 数据集
python AMD-OLMo/prepare_sft_data.py --output_dir ./datasets/tulu --tokenizer tokenizers/allenai_eleuther-ai-gpt-neox-20b-pii-special.json --dataset tulu

# 第二阶段 SFT 数据集
python AMD-OLMo/prepare_sft_data.py --output_dir ./datasets/OpenHermes_WebInstructSub_CodeFeedBack --tokenizer tokenizers/allenai_eleuther-ai-gpt-neox-20b-pii-special.json --dataset 2nd-phase

模型训练

# 单节点训练
HSA_FORCE_FINE_GRAIN_PCIE=1 OMP_NUM_THREADS=128 NCCL_DEBUG=INFO torchrun --nproc_per_node=8 ./scripts/train.py AMD-OLMo/AMD-OLMo-1B.yaml

# 多节点训练
HSA_FORCE_FINE_GRAIN_PCIE=1 OMP_NUM_THREADS=128 NCCL_DEBUG=INFO torchrun --nnodes=$nnodes --node-rank=$node_rank --master_addr=$master_addr --master_port=$master_port --nproc_per_node=8 ./scripts/train.py AMD-OLMo/AMD-OLMo-1B.yaml

DPO 训练

# 安装 trl 库
git clone https://github.com/huggingface/trl.git -b v0.8.6

# 替换 dpo_trainer.py
cp AMD-OLMo/dpo_trainer.py trl/trl/trainer

pip install -e ./trl

# 安装 alignment-handbook
git clone https://github.com/huggingface/alignment-handbook.git hf-align
# 70769f9 是 2024 年 4 月 11 日的主分支。
cd hf-align && git checkout 70769f9 && cd .. 

pip install -e ./hf-align

# 将 AMD OLMo DPO 配方复制到 hf-align/recipes。
cp AMD-OLMo/AMD-OLMo-1B-dpo.yaml hf-align/recipes/

# 准备转换后的 AMD-OLMo SFT Huggingface 模型到 ckpt_dir。
ckpt_dir=amd/AMD-OLMo-1B-SFT
local_tokenizer_dir=${ckpt_dir}

# 设置输出检查点目录。
dpo_ckpt_dir=<your_output_checkpoint_dir>

accelerate launch --config_file hf-align/recipes/accelerate_configs/deepspeed_zero3.yaml \
hf-align/scripts/run_dpo.py hf-align/recipes/AMD-OLMo-1B-dpo.yaml \
--trust_remote_code=true \
--model_name_or_path=${ckpt_dir} \
--tokenizer_name_or_path=${local_tokenizer_dir} \
--output_dir=${dpo_ckpt_dir} \
--num_train_epochs=1 \
--learning_rate=4e-6 \
--beta=0.3 \
--loss_type=sigmoid

主要结果

预训练结果

指令微调结果

聊天基准测试结果

负责任 AI 基准测试结果

评估

我们使用以下开源评估框架对模型进行评估：

• Language Model Evaluation Harness：用于评估常识推理、多任务理解和负责任 AI 基准。
• AlpacaEval：用于评估聊天模型的指令遵循能力。
• MT-Bench：用于评估聊天模型的多轮对话能力。

评估指标

偏差、风险和限制

⚠️ 重要提示

• 这些模型仅用于研究目的，不适用于需要高事实性、安全关键情况、健康或医疗应用、生成虚假信息、促进有毒对话的用例。
• 模型检查点在没有任何安全保证的情况下提供。用户必须根据各自的用例进行全面评估并实施安全过滤机制。
• 可能可以通过提示使模型生成可能事实不准确、有害、暴力、有毒、有偏见或以其他方式令人反感的内容。即使提示并非有意产生此类输出，也可能会生成此类内容。因此，建议用户在使用模型时保持警惕并进行负责任的思考。
• 模型的多语言能力尚未经过测试，因此可能会误解并在不同语言中生成错误的响应。

📄 许可证

Copyright (c) 2018 - 2024 Advanced Micro Devices, Inc. All Rights Reserved.

本项目遵循 Apache License 2.0 许可协议，你可以在以下链接获取许可证副本： http://www.apache.org/licenses/LICENSE-2.0

除非适用法律要求或书面同意，否则根据本许可证分发的软件按“原样”分发，不附带任何形式的明示或暗示的保证或条件。请查看许可证以了解特定语言的许可权限和限制。

引用

如果你使用了我们的 AMD-OLMo 模型，请引用以下文献：

@misc{AMD-OLMo,
    title = {AMD-OLMo: A series of 1B language models trained from scratch by AMD on AMD Instinct™ MI250 GPUs.},
    url = {https://huggingface.co/amd/AMD-OLMo},
    author = {Jiang Liu, Jialian Wu, Prakamya Mishra, Zicheng Liu, Sudhanshu Ranjan, Pratik Prabhanjan Brahma, Yusheng Su, Gowtham Ramesh, Peng Sun, Zhe Li, Dong Li, Lu Tian, Emad Barsoum},
    month = {October},
    year = {2024}
}