OLMo - 训练和使用AI2 模型

文章目录

- 一、关于 OLMo
- - 安装
- 二、模型
- - 概览
  - Checkpoints
- 三、推理
- - 关于微调检查点的推理
  - 量化
- 四、Reproducibility
- - 训练
  - 检查训练数据
- 五、微调
- 六、评估

一、关于 OLMo

OLMo: Open Language Model
OLMo是一个用于训练和使用AI2最先进的开放语言模型的存储库。它由科学家为科学家构建。

github : https://github.com/allenai/OLMo （2409 4.4k ⭐️）
官网：https://allenai.org/olmo

安装

首先根据特定于您的操作系统的说明安装PyTorch。

要从源代码安装（建议用于训练/微调），请运行：

git clone https://github.com/allenai/OLMo.git
cd OLMo
pip install -e .[all]

否则，您可以直接从PyPI自行安装模型代码：

pip install ai2-olmo

二、模型

概览

到目前为止，OLMo家族发布的核心模型是（都在Dolma数据集上训练）：

Model	Training Tokens	Context Length	Training Config	W&B Logs	Data Order File(s) ☨
OLMo 1B	3 Trillion	2048	configs/official/OLMo-1B.yaml	wandb.ai/…/OLMo-1B	epoch 1
OLMo 7B	2.5 Trillion	2048	configs/official/OLMo-7B.yaml	wandb.ai/…/OLMo-7B	epoch 1, epoch 2
OLMo 7B Twin 2T	2 Trillion	2048	configs/official/OLMo-7B.yaml	wandb.ai/…/OLMo-7B-Twin-2T	epoch 1
OLMo 7B April 2024	2.05 Trillion	4096	configs/official/OLMo-7B-0424.yaml	Coming soon	Coming soon
OLMo 7B July 2024	2.75 Trillion	4096	configs/official/OLMo-7B-0724.yaml	Coming soon	Coming soon

请参阅下面的检查训练数据的用法。

Checkpoints

模型训练中间步骤检查点的URL可以在checkpoints/official/下的csv文件中找到。这些“目录”URL目前无法直接访问，但目录中的文件可以公开访问。这些URL也可以提供给训练脚本，以便从检查点恢复训练（参见训练）。每个检查点目录包括：

config.yaml：训练步骤中的配置。
model.pt，optim.pt，train.pt：该训练步骤的模型、优化器和训练状态。

有关其他类型的OLMo检查点（包括OLMo高频变压器检查点）的详细信息，请参见Checkpoints.md。

三、推理

您可以利用我们的Hugging Face集成在OLMo变形金刚检查点上运行推理：

from transformers import AutoModelForCausalLM, AutoTokenizer

olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-1.7-7B-hf")
tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-1.7-7B-hf")

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-1.7-7B-hf")
print(olmo_pipe("Language modeling is"))

关于微调检查点的推理

如果您使用微调中的代码微调模型，您可以使用转换脚本将本机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-1.7-7B-hf", torch_dtype=torch.float16, load_in_8bit=True)  # requires bitsandbytes

量化模型对类型/cuda更敏感，因此建议将输入作为输入传递。input_ids.to（‘cuda’）以避免潜在问题。

四、Reproducibility

训练

用于训练官方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启动训练运行。例如，要在单个8x GPU节点上启动1B模型训练，您将运行：

torchrun --nproc_per_node=8 scripts/train.py configs/official/OLMo-1B.yaml

您也可以使用相同的方法启动多节点作业。请参阅torchrun的留档，了解配置会合后端/端点所需的其他参数。

要从检查点恢复训练，您可以传递其路径（本地或URL）到带有--load_path参数的scripts/train.py。例如，要从OLMo1B的步骤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模型期间检查组成特定批次的确切令牌。我们提供了执行此操作的工具，但首先您需要下载上述数据（除非您有R2 API密钥）并相应地更新相应的配置。

然后记下所需数据顺序文件的URL，该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

# 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