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