安装

首先，请根据您的操作系统和硬件配置，按照 PyTorch 的官方说明进行安装。

对于开发环境，我们建议从源码安装：

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

或者您也可以通过 PyPI 直接安装：

pip install ai2-olmo-core

此外，还有一些可选依赖项需要单独安装，以支持特定功能，包括：

• flash-attn、ring-flash-attn 和 TransformerEngine，用于相应的注意力后端。
• Liger-Kernel，用于低内存“融合线性”损失实现。
• torchao，用于 float8 训练。
• grouped_gemm，用于无丢弃的专家混合（MoE）模型。在 PR #21 发布之前（v0.1.6 版本之后），可能需要从源码编译。
• QuACK，用于一些基于 CuTe 的内核。

我们发布的 Docker 镜像 包含了所有核心及可选依赖，并且会在我们的内部 H100 集群上定期测试。不过，如果您打算使用这些镜像，需要注意以下几点：

• 这些镜像仅包含 OLMo-core 的依赖项，而未预装 OLMo-core 包本身，以便适应代码的频繁更新。
• 如果您的集群硬件或驱动程序/CUDA 版本不同，这些镜像可能无法正常工作。

如果由于上述原因导致发布的镜像无法满足您的需求，您可以参考我们的 Dockerfile，自行构建适合您环境的镜像。

官方训练脚本

已发布模型的官方训练脚本位于 src/scripts/official/ 目录下。

这些脚本通常需要配合 torchrun 使用，或者如果您有 Beaker 访问权限，则可以使用 OLMo-core 提供的 Beaker 启动 CLI 来运行。

例如：

torchrun --nproc-per-node=8 src/scripts/official/OLMo2/OLMo-2-0325-32B-train.py \
  --save-folder=/path/to/save/checkpoints

您还可以通过命令行覆盖大多数配置选项。例如，要调整学习率，可以这样启动脚本：

torchrun --nproc-per-node=8 src/scripts/official/OLMo2/OLMo-2-0325-32B-train.py \
  --save-folder=/path/to/save/checkpoints \
  --train_module.optim.lr=6e-3

若需从检查点继续训练，我们可以使用另一份专门的脚本，如下所示：

torchrun --nproc-per-node=8 src/scripts/official/OLMo2/OLMo-2-0325-32B-anneal.py \
  --save-folder=/path/to/save/checkpoints \
  --checkpoint=https://olmo-checkpoints.org/ai2-llm/peteish32/step721901

可用的训练脚本

模型系列	目录	描述
OLMo-2	src/scripts/official/OLMo2/	OLMo-2 32B 模型的训练脚本及模型卡片
OLMo-3	src/scripts/official/OLMo3/	OLMo-3 7B 和 32B 模型的训练脚本及模型卡片

推理

使用 Hugging Face Transformers

您可以通过我们与 Hugging Face transformers 的集成，在 OLMo 检查点上进行推理：

pip install transformers>=4.57.0

from transformers import AutoModelForCausalLM, AutoTokenizer
olmo = AutoModelForCausalLM.from_pretrained("allenai/Olmo-3-1125-32B")
tokenizer = AutoTokenizer.from_pretrained("allenai/Olmo-3-1125-32B")
message = ["语言建模是 "]
inputs = tokenizer(message, return_tensors='pt', return_token_type_ids=False)
# inputs = {k: v.to('cuda') for k,v in inputs.items()} # 可选：验证是否使用 GPU
# olmo = olmo.to('cuda')
response = olmo.generate(**inputs, max_new_tokens=100, do_sample=True, temperature=1.0, top_p=0.7)
print(tokenizer.batch_decode(response, skip_special_tokens=True)[0])

或者，您也可以使用 Hugging Face 的管道抽象：

from transformers import pipeline
olmo_pipe = pipeline("text-generation", model="allenai/Olmo-3-1125-32B")
print(olmo_pipe("语言建模是"))

使用 vLLM

vLLM 提供了针对 OLMo 模型的高吞吐量推理能力。您可以用它来进行离线批量推理：

pip install vllm>=0.11.0

from vllm import LLM, SamplingParams
llm = LLM(model="allenai/Olmo-3-1125-32B")
sampling_params = SamplingParams(temperature=1.0, top_p=0.7)
prompts = ["语言建模是"]
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

更多详细信息，请参阅 vLLM 文档。

使用 Olmo-core（beta 版）

自回归生成功能已直接集成在 Olmo-core 中。借助此功能，我们提供了一个聊天循环演示，可用于与模型进行交互式对话：

python -m olmo_core.generate.chat https://olmo-checkpoints.org/ai2-llm/Olmo-3-1025-7B/stage3/step11921/ --max-new-tokens 512

评估

用于评估 OLMo 模型的额外工具可在 OLMo Eval 和 olmes 仓库中找到。

开发

Python 库的源代码位于 src/olmo_core 目录下。相应的测试文件位于 src/test 目录下。库的文档位于 docs 目录中。您可以在本地通过运行 make docs 来构建文档。

代码检查：

• 我们使用 pytest 运行测试。您可以使用 pytest -v src/test 运行所有测试，也可以指定某个测试文件单独运行。
• 我们使用 isort 和 black 进行代码格式化。理想情况下，您应将这些工具集成到编辑器中，但也可以手动运行它们，或通过 pre-commit 钩子进行配置。要验证所有文件是否格式正确，请运行 make style-check。
• 我们使用 ruff 作为主要的代码静态分析工具。您可以使用 make lint-check 来运行它。
• 我们使用 mypy 作为类型检查工具。您可以使用 make type-check 来运行它。

引用

@misc{olmo20242olmo2furious,
      title={{2 OLMo 2 Furious}},
      author={{Team OLMo} and Pete Walsh and Luca Soldaini and Dirk Groeneveld and Kyle Lo and Shane Arora and Akshita Bhagia and Yuling Gu and Shengyi Huang and Matt Jordan and Nathan Lambert and Dustin Schwenk and Oyvind Tafjord and Taira Anderson and David Atkinson and Faeze Brahman and Christopher Clark and Pradeep Dasigi and Nouha Dziri and Michal Guerquin and Hamish Ivison and Pang Wei Koh and Jiacheng Liu and Saumya Malik and William Merrill and Lester James V. Miranda and Jacob Morrison and Tyler Murray and Crystal Nam and Valentina Pyatkin and Aman Rangapur and Michael Schmitz and Sam Skjonsberg and David Wadden and Christopher Wilhelm and Michael Wilson and Luke Zettlemoyer and Ali Farhadi and Noah A. Smith and Hannaneh Hajishirzi},
      year={2024},
      eprint={2501.00656},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2501.00656},
}

OLMo-core 快速上手指南

OLMo-core 是 AllenAI 提供的用于构建和训练 OLMo 系列大语言模型的核心工具库。本指南将帮助中国开发者快速完成环境配置并运行基础示例。

1. 环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu) 或 macOS。
• Python: 建议 Python 3.9 或更高版本。
• PyTorch: 必须先安装与您的硬件（CPU/GPU）和操作系统匹配的 PyTorch。
  • 访问 PyTorch 官网 获取安装命令。
  • 国内加速: 推荐使用清华或中科大镜像源安装 PyTorch，例如：

    pip install torch torchvision torchaudio --index-url https://pypi.tuna.tsinghua.edu.cn/simple
    

• 可选依赖: 若需使用 Flash Attention、混合专家模型 (MoE) 或 Float8 训练等高级功能，需额外安装 flash-attn、grouped_gemm 等库（详见官方文档）。普通推理和基础训练通常不需要立即安装这些。

2. 安装步骤

您可以选择从源码安装（推荐用于开发）或从 PyPI 安装（推荐用于快速使用）。

方式一：从 PyPI 安装（推荐）

使用国内镜像源可显著提升下载速度：

pip install ai2-olmo-core -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从源码安装（开发模式）

如果您需要修改源码或运行最新示例：

git clone https://github.com/allenai/OLMo-core.git
cd OLMo-core
pip install -e .[all] -i https://pypi.tuna.tsinghua.edu.cn/simple

注意: 官方提供了预构建的 Docker 镜像，包含所有核心及可选依赖。如果您的集群环境与官方 H100 环境不一致，建议参考仓库中的 Dockerfile 自行构建适配本地驱动的镜像。

3. 基本使用

OLMo-core 支持多种使用场景，以下是两种最常用的快速入门方式。

场景一：使用 Hugging Face Transformers 进行推理

这是最简单的使用方式，适合快速测试模型效果。首先确保安装了较新版本的 transformers：

pip install "transformers>=4.57.0" -i https://pypi.tuna.tsinghua.edu.cn/simple

Python 代码示例：

from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型和分词器 (以 OLMo-3 32B 为例)
model_name = "allenai/Olmo-3-1125-32B"
olmo = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 准备输入
message = ["Language modeling is "]
inputs = tokenizer(message, return_tensors='pt', return_token_type_ids=False)

# (可选) 移动到 GPU
# inputs = {k: v.to('cuda') for k,v in inputs.items()}
# olmo = olmo.to('cuda')

# 生成文本
response = olmo.generate(**inputs, max_new_tokens=100, do_sample=True, temperature=1.0, top_p=0.7)

# 输出结果
print(tokenizer.batch_decode(response, skip_special_tokens=True)[0])

或者使用更简洁的 Pipeline API：

from transformers import pipeline

olmo_pipe = pipeline("text-generation", model="allenai/Olmo-3-1125-32B")
print(olmo_pipe("Language modeling is"))

场景二：使用 OLMo-core 原生交互式对话

OLMo-core 内置了自回归生成能力，提供了一个简单的聊天循环演示脚本，可直接与模型交互：

python -m olmo_core.generate.chat https://olmo-checkpoints.org/ai2-llm/Olmo-3-1025-7B/stage3/step11921/ --max-new-tokens 512

场景三：启动官方训练脚本

若需复现官方模型训练，请使用 torchrun 启动位于 src/scripts/official/ 下的脚本。以下是一个启动 32B 模型训练的示例：

torchrun --nproc-per-node=8 src/scripts/official/OLMo2/OLMo-2-0325-32B-train.py \
  --save-folder=/path/to/save/checkpoints

您可以通过命令行参数覆盖默认配置，例如修改学习率：

torchrun --nproc-per-node=8 src/scripts/official/OLMo2/OLMo-2-0325-32B-train.py \
  --save-folder=/path/to/save/checkpoints \
  --train_module.optim.lr=6e-3