OLMo:开放语言模型
OLMo是一个用于训练和使用AI2最先进的开放语言模型的代码库。 它由科学家为科学家而建。
安装
首先根据您的操作系统特定说明安装PyTorch。
要从源代码安装(推荐用于训练/微调),请运行:
git clone https://github.com/allenai/OLMo.git
cd OLMo
pip install -e .[all]
或者,您可以直接从PyPI安装模型代码:
pip install ai2-olmo
模型
概览
到目前为止发布的OLMo系列核心模型(均在Dolma数据集上训练)有:
模型 | 训练tokens | 上下文长度 | 训练配置 | W&B日志 | 数据顺序文件 ☨ |
---|---|---|---|---|---|
OLMo 1B | 3万亿 | 2048 | configs/official/OLMo-1B.yaml | wandb.ai/…/OLMo-1B | 第1轮 |
OLMo 7B | 2.5万亿 | 2048 | configs/official/OLMo-7B.yaml | wandb.ai/…/OLMo-7B | 第1轮, 第2轮 |
OLMo 7B Twin 2T | 2万亿 | 2048 | configs/official/OLMo-7B.yaml | wandb.ai/…/OLMo-7B-Twin-2T | 第1轮 |
OLMo 7B 2024年4月 | 2.05万亿 | 4096 | configs/official/OLMo-7B-0424.yaml | 即将推出 | 即将推出 |
OLMo 7B 2024年7月 | 2.75万亿 | 4096 | configs/official/OLMo-7B-0724.yaml | 即将推出 | 即将推出 |
☨ 参见下方检查训练数据了解用法。
检查点
模型训练中间步骤的检查点URL可以在checkpoints/official/
目录下的csv文件中找到。这些"目录"URL目前无法直接访问,但目录内的文件是公开可访问的。这些URL也可以提供给训练脚本以从检查点恢复训练(参见训练)。每个检查点目录包含:
config.yaml
:该训练步骤的配置。model.pt
、optim.pt
、train.pt
:该训练步骤的模型、优化器和训练状态。
有关其他类型OLMo检查点(包括OLMo HF Transformers检查点)的详细信息可以在Checkpoints.md中找到。
推理
您可以使用我们的Hugging Face集成来对OLMo Transformers检查点进行推理:
from transformers import AutoModelForCausalLM, AutoTokenizer
olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-7B-0724-hf")
tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-7B-0724-hf")
message = ["语言建模是"]
inputs = tokenizer(message, return_tensors='pt', return_token_type_ids=False)
response = olmo.generate(**inputs, max_new_tokens=100, do_sample=True, top_k=50, top_p=0.95)
print(tokenizer.batch_decode(response, skip_special_tokens=True)[0])
或者,使用Hugging Face pipeline抽象:
from transformers import pipeline
olmo_pipe = pipeline("text-generation", model="allenai/OLMo-7B-0724-hf")
print(olmo_pipe("语言建模是"))
对微调后的检查点进行推理
如果您使用微调中的代码对模型进行了微调,可以使用转换脚本将原生OLMo检查点转换为兼容Hugging Face的检查点。
python scripts/convert_olmo_to_hf_new.py --input_dir /path/to/olmo/checkpoint --output_dir /path/to/hf/checkpoint/ --tokenizer_json_path tokenizers/allenai_gpt-neox-olmo-dolma-v1_5.json
量化
olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-7B-0724-hf", torch_dtype=torch.float16, load_in_8bit=True) # 需要bitsandbytes
量化后的模型对类型/cuda更敏感,因此建议将输入作为inputs.input_ids.to('cuda')传递,以避免潜在问题。
可复现性
训练
用于训练官方OLMo模型的配置文件位于configs/official/
目录中。
请注意,虽然训练和验证数据是公开和免费下载的,但这些配置文件中的数据路径指向了一个CloudFlare R2存储桶,需要API密钥才能进行程序化访问。 因此,为了使用这些配置文件复现训练过程,您首先需要将相应的数据下载到您选择的位置,然后相应地更新配置中的路径。
您可以通过将r2://olmo-data
替换为https://olmo-data.org
来从R2 URL派生公共HTTP URL。
例如,如果R2数据URL是:
r2://olmo-data/preprocessed/olmo-mix/v1_5/gpt-neox-20b-pii-special/part-000-00000.npy
那么相应的公共URL是:
https://olmo-data.org/preprocessed/olmo-mix/v1_5/gpt-neox-20b-pii-special/part-000-00000.npy
更新配置中的数据路径后,您可以通过torchrun
启动训练。例如,要在单个8 GPU节点上启动1B模型训练,您可以运行:
torchrun --nproc_per_node=8 scripts/train.py configs/official/OLMo-1B.yaml
您也可以使用相同的方法启动多节点作业。请参阅torchrun
的文档以了解配置集合后端/端点所需的额外参数。
要从检查点恢复训练,您可以将其路径(本地或URL)
传递给scripts/train.py
,并使用--load_path
参数。例如,要从OLMo 1B运行的第1000步恢复训练:
torchrun --nproc_per_node=8 scripts/train.py configs/official/OLMo-1B.yaml --load_path=https://olmo-checkpoints.org/ai2-llm/olmo-small/w1r5xfzt/step1000-unsharded
检查训练数据
您可能对检查OLMo模型训练期间特定批次的确切token感兴趣。 我们提供了工具来实现这一点,但首先您需要如上所述下载数据(除非您有R2 API密钥)并相应地更新相应的配置。 然后注意你想要的数据顺序文件的URL,可以在模型概览表格中找到。例如,OLMo-7B模型第一轮训练的数据顺序文件是https://olmo-checkpoints.org/ai2-llm/olmo-medium/wvc30anm/train_data/global_indices.npy。
获得该文件后,你可以使用以下代码片段来检查特定批次内的数据:
import numpy as np
from cached_path import cached_path
from olmo.config import TrainConfig
from olmo.data import build_memmap_dataset
# 更新这些路径为你想要的:
data_order_file_path = cached_path("https://olmo-checkpoints.org/ai2-llm/olmo-medium/wvc30anm/train_data/global_indices.npy")
train_config_path = "configs/official/OLMo-7B.yaml"
cfg = TrainConfig.load(train_config_path)
dataset = build_memmap_dataset(cfg, cfg.data)
batch_size = cfg.global_train_batch_size
global_indices = np.memmap(data_order_file_path, mode="r+", dtype=np.uint32)
def get_batch_instances(batch_idx: int) -> list[list[int]]:
batch_start = batch_idx * batch_size
batch_end = (batch_idx + 1) * batch_size
batch_indices = global_indices[batch_start:batch_end]
batch_instances = []
for index in batch_indices:
token_ids = dataset[index]["input_ids"].tolist()
batch_instances.append(token_ids)
return batch_instances
# 获取第一批次中所有2048 x 2048的标记ID。
get_batch_instances(0)
微调
要使用我们的训练器微调OLMo模型,你首先需要准备你的数据集,将其标记化并将标记ID保存到一个扁平的numpy内存映射数组中。可以参考scripts/prepare_tulu_data.py
中的Tulu V2数据集示例,该示例可以轻松修改以适用于其他数据集。
接下来,准备你的训练配置。在configs/
目录中有许多示例可以作为起点。最重要的是确保模型参数(配置中的model
字段)与你要开始使用的检查点相匹配。为了安全起见,你可以始终从模型检查点附带的配置开始。至少需要对配置进行以下更改,或从命令行提供相应的覆盖:
- 更新
load_path
以指向你要开始使用的检查点。 - 将
reset_trainer_state
设置为true
。 - 更新
data.paths
以指向你生成的token_ids.npy
文件。 - 可选地更新
data.label_mask_paths
以指向你生成的label_mask.npy
文件,除非你不需要对损失进行特殊掩码。 - 更新
evaluators
以添加/删除循环内评估。
一旦你对训练配置感到满意,就可以通过torchrun
启动训练任务。例如:
torchrun --nproc_per_node=8 scripts/train.py {path_to_train_config} \
--data.paths=[{path_to_data}/input_ids.npy] \
--data.label_mask_paths=[{path_to_data}/label_mask.npy] \
--load_path={path_to_checkpoint} \
--reset_trainer_state
注意:只有在你没有在配置中更新这些字段的情况下,才需要传递像--reset_trainer_state
这样的CLI覆盖。
评估
OLMo Eval仓库提供了评估OLMo模型的其他工具。
引用
@article{OLMo,
title={OLMo: Accelerating the Science of Language Models},
author={Dirk Groeneveld and Iz Beltagy and Pete Walsh and Akshita Bhagia and Rodney Kinney and Oyvind Tafjord and A. Jha and Hamish Ivison and Ian Magnusson and Yizhong Wang and Shane Arora and David Atkinson and Russell Authur and Khyathi Raghavi Chandu and Arman Cohan and Jennifer Dumas and Yanai Elazar and Yuling Gu and Jack Hessel and Tushar Khot and William Merrill and Jacob Daniel Morrison and Niklas Muennighoff and Aakanksha Naik and Crystal Nam and Matthew E. Peters and Valentina Pyatkin and Abhilasha Ravichander and Dustin Schwenk and Saurabh Shah and Will Smith and Emma Strubell and Nishant Subramani and Mitchell Wortsman and Pradeep Dasigi and Nathan Lambert and Kyle Richardson and Luke Zettlemoyer and Jesse Dodge and Kyle Lo and Luca Soldaini and Noah A. Smith and Hanna Hajishirzi},
year={2024},
url={https://api.semanticscholar.org/CorpusID:267365485},
journal={arXiv preprint},
}