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データセット（第1段階）で教師付き微調整（SFT）され、その後、OpenHermes-2.5、WebInstructSub、およびCode-Feedbackデータセット（第2段階）でさらに微調整されました。
• AMD-OLMo-1B-SFT-DPO：UltraFeedbackデータセットで直接嗜好最適化（DPO）を使用して、人間の嗜好にアライメントされました。

🚀 クイックスタート

このプロジェクトでは一連の言語モデルを提供しており、以下の手順ですぐに使用することができます。

1. 依存関係のインストール：後述のインストールガイドを参照して、必要なソフトウェアとライブラリをインストールしてください。
2. サンプルコードの実行：使用例のセクションを参照して、提供されているコード例を使ってモデル推論を行ってください。

✨ 主な機能

• 高性能訓練：AMD Instinct™ MI250 GPUを使用して訓練されており、高い訓練スループットを実現しています。
• 多段階微調整：事前学習、教師付き微調整（SFT）、およびDPOアライメントなどの複数の段階を経て、モデルの性能が向上しています。
• 多タスクでの優れた性能：複数の標準ベンチマークテストで優れた結果を示し、良好な常識推論、多タスク理解、および責任あるAI能力を備えています。

📦 インストール

AMD GPU上でのPyTorchの実行

Docker HubのROCm Dockerイメージを使用して、AMD GPU上で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 GPUが含まれており、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データセットをダウンロードして準備
# 第1段階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

# 第2段階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}
}