OLMo

<div align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/423215fa-e392-4e5e-aaae-bc32ad11f9ff.gif" alt="OLMo标志" width="800" style="margin-left:'auto' margin-right:'auto' display:'block'"/> <br> <br> <h1>OLMo：开放语言模型</h1> </div> <p align="center"> <a href="https://github.com/allenai/OLMo/blob/main/LICENSE"> <img alt="GitHub许可证" src="https://img.shields.io/github/license/allenai/OLMo"> </a> <a href="https://github.com/allenai/OLMo/releases"> <img alt="GitHub发布" src="https://yellow-cdn.veclightyear.com/835a84d5/aba8994f-ca8b-4719-82c9-5ebf29984e50.svg"> </a> <a href="https://arxiv.org/pdf/2402.00838.pdf"> <img alt="论文URL" src="https://img.shields.io/badge/arxiv-2402.00838-blue"> </a> </p>

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