AMD-OLMo are a series of 1B language models trained from scratch by AMD on AMD Instinct™ MI250 GPUs. The training code used is based on
OLMo.
We release the pre-trained model, supervised fine-tuned model, and DPO aligned model as follows:
-
Hardware: Each compute node consists of 4 AMD Instinct™ MI250 GPUs. We use 16 nodes for pretraining AMD-OLMo-1B
-
Training throughput: 12,200 tokens/sec/gpu
-
Model architecture: AMD-OLMo-1B is based on the model architecture and training set up of fully open source 1 billion version of
OLMo-1B with the details below:
| Parameter size | Number of layers | Number of heads | Hidden size | Context length | Vocabulary Size |
|---|
| 1.2B | 16 | 16 | 2048 | 2048 | 50,280 |
-
Hyper-parameters:
| Stage | LR schedule | Peak LR | Warmup steps | Epochs | Batch size (tokens) |
|---|
| Pretraining | Cosine | 4.0e-4 | 2000 | 1 | 4M |
| SFT Phase 1 | Linear | 2.0e-5 | 200 | 3 | 262K |
| SFT Phase 2 | Linear | 2.0e-5 | 200 | 3 | 1024K |
| DPO | Cosine | 4.0e-6 | 47 | 1 | 64K |
For more details, please refer to our
blog.
For running pytorch on AMD GPUs you can use the following rocm docker as in
docker hub
1docker pull rocm/pytorch:latest
2# Inside docker
3pip install transformers
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model = AutoModelForCausalLM.from_pretrained("amd/AMD-OLMo-1B-SFT").to("cuda") # remove .to("cuda") to load on cpu
4tokenizer = AutoTokenizer.from_pretrained("amd/AMD-OLMo-1B-SFT")
5
6prompt = "What is large language model?"
7bos = tokenizer.eos_token
8template = bos + "<|user|>\n{prompt}\n<|assistant|>\n"
9
10input_text = template.format(prompt=prompt)
11inputs = tokenizer([input_text], return_tensors='pt', return_token_type_ids=False).to("cuda")
12outputs = model.generate(**inputs, max_new_tokens=1000, do_sample=True, top_k=50, top_p=0.95)
13print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
*All numbers in above tables were obtained from our evaluations.
We use the following open source evaluation frameworks for evaluating our models:
1# lm-eval-harness
2git clone https://github.com/EleutherAI/lm-evaluation-harness
3cd lm-evaluation-harness
4pip install -e .
5
6# AlpacaEval
7pip install git+https://github.com/tatsu-lab/alpaca_eval
8cd alpaca_eval
9pip install -e .
10
11# MT-Bench
12git clone https://github.com/lm-sys/FastChat.git
13cd FastChat
14pip install -e ".[model_worker,llm_judge]"
1# lm-eval-harness
2HF_MODEL=amd/AMD-OLMo-1B-SFT-DPO
3accelerate launch -m lm_eval --model hf \
4 --model_args pretrained=$HF_MODEL,trust_remote_code=True \
5 --tasks arc_easy,arc_challenge,hellaswag,piqa,boolq,sciq,winogrande,openbookqa,mmlu,gsm8k_cot,bbh_cot_fewshot,toxigen,truthfulqa,crows_pairs \
6 --device cuda \
7 --batch_size 32 \
8 --output_path ./lm-eval-results/$HF_MODEL
1WORK_DIR="<path_to_your_working_directory>"
2cd $WORK_DIR
3# Clone OLMo codebase:
4git clone https://github.com/allenai/OLMo.git --branch v0.3.0
5cd OLMo
6# Clone AMD-OLMo that contains files to reproduce our model training
7git clone https://huggingface.co/amd/AMD-OLMo
8
9docker pull rocm/pytorch:latest
10docker 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
11
12# Remove Line 17 as the docker already has ROCm PyTorch installed
13sed -i '17d' pyproject.toml
14pip install -e .[all]
1# Download
2DATA_DIR=./datasets/dolma
3mkdir -p $DATA_DIR
4
5PARALLEL_DOWNLOADS="<number_of_parallel_downloads>"
6cat "AMD-OLMo/dolma_v1_7_subset.txt" | xargs -n 1 -P $PARALLEL_DOWNLOADS wget -q -P $DATA_DIR
7
8# Prepare
9NUM_WORKERS="<number_of_workers>"
10python scripts/prepare_memmap_dataset.py $DATA_DIR/*.json.gz -o $DATA_DIR/memmap_dataset --workers $NUM_WORKERS
1# 1st phase SFT dataset
2python AMD-OLMo/prepare_sft_data.py --output_dir ./datasets/tulu --tokenizer tokenizers/allenai_eleuther-ai-gpt-neox-20b-pii-special.json --dataset tulu
3
4# 2nd phase SFT dataset
5python 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
1# Single node
2HSA_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
3
4# Multiple nodes
5HSA_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
1# install trl library
2git clone https://github.com/huggingface/trl.git -b v0.8.6
3
4# replace dpo_trainer.py
5cp AMD-OLMo/dpo_trainer.py trl/trl/trainer
6
7pip install -e ./trl
8
9# install alignment-handbook
10git clone https://github.com/huggingface/alignment-handbook.git hf-align
11# 70769f9 is the main branch on 2024-04-11.
12cd hf-align && git checkout 70769f9 && cd ..
13
14pip install -e ./hf-align
15
16# Copy AMD OLMo DPO recipe to hf-align/recipes.
17cp AMD-OLMo/AMD-OLMo-1B-dpo.yaml hf-align/recipes/
18
19# Prepare the converted AMD-OLMo SFT Huggingface model to ckpt_dir.
20ckpt_dir=amd/AMD-OLMo-1B-SFT
21local_tokenizer_dir=${ckpt_dir}
22
23# Set output checkpoint dir.
24dpo_ckpt_dir=<your_output_checkpoint_dir>
25
26accelerate launch --config_file hf-align/recipes/accelerate_configs/deepspeed_zero3.yaml \
27hf-align/scripts/run_dpo.py hf-align/recipes/AMD-OLMo-1B-dpo.yaml \
28--trust_remote_code=true \
29--model_name_or_path=${ckpt_dir} \
30--tokenizer_name_or_path=${local_tokenizer_dir} \
31--output_dir=${dpo_ckpt_dir} \
32--num_train_epochs=1 \
33--learning_rate=4e-6 \
34--beta=0.3 \
35--loss_type=sigmoid
1@article{instella,
2 title={Instella: Fully Open Language Models with Stellar Performance},
3 author={Liu, Jiang and Wu, Jialian and Yu, Xiaodong and Su, Yusheng and Mishra, Prakamya and Ramesh, Gowtham and Ranjan, Sudhanshu and Manem, Chaitanya and Sun, Ximeng and Wang, Ze and Brahma, Pratik Prabhanjan and Liu, Zicheng and Barsoum, Emad},
4 journal={arXiv preprint arXiv:2511.10628},
5 year={2025}
6}
Copyright (c) 2018-2024 Advanced Micro Devices, Inc. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.