OLMo 是一个用于训练和使用 AI2 最先进的开放语言模型的存储库。 它是由科学家建造的,为科学家服务。
首先,根据特定于你的操作系统的说明安装 PyTorch。
要从源代码安装(建议用于训练/微调),请运行:
git clone https://github.com/allenai/OLMo.git
cd OLMo
pip install -e .[all]
否则,你可以使用以下命令直接从 PyPI 自行安装模型代码:
pip install ai2-olmo
目前发布的OLMo系列中的核心模型是(均在Dolma数据集上训练):
型 | 训练令牌 | 上下文长度 | 训练配置 | W&B 日志 | 数据订单文件 ☨ |
---|---|---|---|---|---|
OLMo 1B | 3万亿 | 2048 | configs/official/OLMo-1B.yaml | 纪元 1 | |
OLMo 7B系列 | 2.5万亿 | 2048 | configs/official/OLMo-7B.yaml | wandb.ai/ai2-llm/OLMo-7B | 纪元 1, 纪元 2 |
OLMo 7B 双床 2T | 2 万亿 | 2048 | configs/official/OLMo-7B.yaml | 纪元 1 |
☨ 请参阅下面的检查训练数据以了解使用情况。
你可以利用我们的 Hugging Face 集成在 olmo 检查点上运行推理:
from hf_olmo import * # registers the Auto* classes
from transformers import AutoModelForCausalLM, AutoTokenizer
olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-7B")
tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-7B")
message = ["Language modeling is "]
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 管道抽象:
from transformers import pipeline
olmo_pipe = pipeline("text-generation", model="allenai/OLMo-7B")
print(olmo_pipe("Language modeling is"))
如果使用上述代码微调模型,则可以使用转换脚本将本机 OLMo 检查点转换为与 Hugging Face 兼容的检查点
python hf_olmo/convert_olmo_to_hf.py --checkpoint-dir /path/to/checkpoint
olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-7B", torch_dtype=torch.float16, load_in_8bit=True) # requires bitsandbytes
量化模型对类型 / cuda 更敏感,因此建议将输入作为 inputs.input_ids.to('cuda') 传递,以避免潜在问题。
用于训练官方 OLMo 模型的配置在 configs/official/
目录中提供。
请注意,虽然训练和验证数据是公开的,可以免费下载,但这些配置中的数据路径指向 CloudFlare R2 存储桶,这需要 API 密钥才能进行编程访问。 因此,为了使用这些配置中的任何一个来重现训练运行,你首先必须将相应的数据下载到你选择的位置,然后相应地更新配置中的路径。
你可以通过将 替换为 来从 R2 URL 派生公共 HTTP URL。 例如,如果 R2 数据 URL 为:
r2://olmo-data
https://olmo-data.org
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
在配置中更新数据路径后,可以通过 启动训练运行。例如,要在单个 8 个 GPU 节点上启动 1B 模型训练,你可以运行:
torchrun
torchrun --nproc_per_node=8 scripts/train.py configs/official/OLMo-1B.yaml
你也可以使用相同的方法来启动多节点作业。请参阅文档,了解配置会合后端/端点所需的其他参数。
torchrun
在训练其中一个 OLMo 模型期间,你可能会对检查组成特定批次的确切标记感兴趣。 我们提供了执行此操作的工具,但首先你需要下载上述数据(除非你有 R2 API 密钥)并相应地更新相应的配置。
然后记下所需数据订单文件的 URL,该文件可在“模型概述”(Models Overview) 表格中找到。例如,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
# Update these paths to what you want:
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
# Get all 2048 x 2048 token IDs in the first batch.
get_batch_instances(0)
要使用我们的训练器对 OLMo 模型进行微调,你首先需要通过标记化数据集并将标记 ID 保存到平面 numpy 内存映射数组来准备数据集。有关 Tulu V2 数据集的示例,请参阅 scripts/prepare_tulu_data.py
,该数据集可以轻松针对其他数据集进行修改。
接下来,准备训练配置。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
注意:仅当你没有在配置中更新这些字段时,才需要传递 CLI 覆盖。
--reset_trainer_state
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},
}