OpenLM语言建模（LM）代码库说明
OpenLM代码库说明，DCLM数据集评估的项目引用这份仓库进行数据处理、模型训练和评估。注意和Megatron-LM的区别。
https://github.com/mlfoundations/open_lm
2025年6月4日
以下是网页正文部分的中文翻译：
OpenLM
OpenLM 是一个简洁但性能出色的语言建模（LM）代码库，旨在便于开展中等规模语言模型的研究。我们已验证 OpenLM 的性能，支持多达 70 亿参数和 256 个 GPU。与其他代码库（如 Megatron）不同，我们的核心建模代码仅依赖 PyTorch、XFormers 或 Triton。
内容目录
-
发行说明
-
快速入门
- 设置
- 处理训练数据
- 运行训练
- 评估模型
- 生成文本
-
预训练模型
-
团队与致谢
发行说明
- 2023 年 9 月 26 日：公开发布，并在 Laion 博客上被提及。
- 2023 年 8 月 18 日：更新了 README.md 文件。
快速入门
接下来，我们将通过一个基础示例，从全新安装开始，下载并预处理一些训练数据，然后训练模型。
设置
我们要求使用 Python >=3.9，并且安装了最新版本的 PyTorch，以及一些其他软件包。完整的依赖列表包含在 requirements.txt 文件中，可以通过以下命令在 Python 环境中安装：
>>> pip install -r requirements.txt
接下来，为了在虚拟环境中随处访问 open_lm，请使用 pip 安装（从顶级 GitHub 仓库目录中运行）：
>>> pip install --editable .
一些注意事项：
- 我们推荐使用 WandB 和 Tensorboard 进行日志记录。我们将在下面说明如何在训练过程中使用它们。
处理训练数据
接下来，您必须指定一组经过分词的数据。在本示例中，我们将使用 HuggingFace 上最近发布的英文维基百科数据集。为了将该数据集下载到本地，我们提供了一个脚本，位于 open_lm/datapreprocess/wiki_download.py。您只需指定一个输出目录，用于存储原始数据：
python open_lm/datapreprocess/wiki_download.py --output-dir path/to/raw_data
接下来，我们通过运行 BPE 分词器处理训练数据，并将其分割成适当长度的片段。默认情况下，我们使用与 GPT-NeoX-20B 捆绑的分词器。为此，请使用以下脚本：
>>> python open_lm/datapreprocess/make_2048.py \
--input-files path_to_raw_data/*.jsonl
--output-dir preproc_data
--num-workers 32
--num-consumers 1
其中，input-files 参数会将其所有（可能很多）参数通过 Python 的 glob 模块传递，支持使用通配符。如果需要，可以通过设置环境变量 S3_BASE 并在脚本中添加 --upload-to-s3 标志，将数据存储到 S3 中。这会将分片数据保存到指定的存储桶中，其前缀为 S3_BASE。例如：
>>> export S3_BASE=preproc_data-v1/
>>> python open_lm/datapreprocess/make2048.py --upload-to-s3 ... # 与前面相同的参数
运行训练
现在，经过分词的数据可以传递给主训练脚本 open_lm/main.py。通过 torchrun 实现分布式计算，超参数通过多种关键字参数指定。我们在此强调其中一些最重要的参数：
train-data：分词训练数据的位置。如果是在本地生成并存储的，则会指向包含类似preproc_data/2048-v1/0/XXXXXXX.tar文件的目录。使用 webdataset 包处理数据，支持使用通配符，例如preproc_data/2048-v1/0/{0000000..0000099}.tar，用于选择前 100 个.tar文件。model：使用的模型。请参阅下表，了解有效的选项以及每个模型的参数规模。train-num-samples：从指定的训练数据集中使用的样本数量。name：此特定训练运行的名称，用于日志记录目的。report-to：如果存在，可以是wandb、tensorboard或all，用于将日志信息存储到 WandB 或 Tensorboard 中。
模型选择包含在下表中，例如，11m 表示一个拥有 1100 万参数的模型，而 1b 表示一个拥有 10 亿参数的模型。
| 模型名称 |
|---|
open_lm_11m |
open_lm_25m |
open_lm_87m |
open_lm_160m |
open_lm_411m |
open_lm_830m |
open_lm_1b |
open_lm_3b |
open_lm_7b |
以下是示例训练命令：
>>> export CUDA_VISIBLE_DEVICES=0,1,2,3
>>> torchrun --nproc-per-node 4 -m open_lm.main \
--model open_lm_3b \
--train-data /preproc_data/shard-{0000000..0000099}.tar \
--train-num-samples 1000000000 \
--workers 8 \
--dataset-resampled \
--precision amp_bfloat16 \
--batch-size 8 \
--grad-checkpointing \
--log-every-n-steps 100 \
--grad-clip-norm 1 \
--data-key txt \
--lr 3e-4 \
--fsdp --fsdp-amp \
--warmup 2000 \
--wd 0.1 \
--beta2 0.95 \
--epochs 100 \
--report-to wandb \
--wandb-project-name open_lm_example \
--name open_lm_ex_$RANDOM \
--resume latest \
--logs path/to/logging/dir/
检查点和最终模型权重将保存到指定的日志目录中。
在训练过程中，上述命令将通过有放回抽样选择分片进行训练。也可以通过无放回抽样选择分片进行训练。为此，必须首先使用以下命令对输入数据集进行预处理：
python -m open_lm.utils.make_wds_manifest --data-dir /preproc_data/
这将在 /preproc_data 目录下创建一个名为 manifest.jsonl 的文件。然后可以通过以下示例命令通过无放回抽样进行训练：
>>> export CUDA_VISIBLE_DEVICES=0,1,2,3
>>> torchrun --nproc-per-node 4 -m open_lm.main \
--model open_lm_3b \
--dataset-manifest /preproc_data/manifest.jsonl \
--train-num-samples 1000000000 \
--workers 8 \
--precision amp_bfloat16 \
--batch-size 8 \
--grad-checkpointing \
--log-every-n-steps 100 \
--grad-clip-norm 1 \
--data-key txt \
--lr 3e-4 \
--fsdp --fsdp-amp \
--warmup 2000 \
--wd 0.1 \
--beta2 0.95 \
--epochs 100 \
--report-to wandb \
--wandb-project-name open_lm_example \
--name open_lm_ex_$RANDOM \
--resume latest \
--logs path/to/logging/dir/
数据集清单
使用 open_lm/utils/make_wds_manifest.py 创建的清单是一个 jsonl 文件，用于描述数据集。该文件的每一行对应数据集的一个分片，是一个包含两个字段的 json 对象：
"shard"：数据集中分片的名称。"num_sequences"：分片中包含的序列数量。每个序列包含一组固定长度的标记。
该清单文件提供了有关数据集的辅助信息，并假定它位于分片所在的同一目录中。
评估模型
训练完成后，我们可以对模型进行评估。这需要使用 LLM Foundry，可以通过以下命令安装：pip install llm-foundry。接下来需要传递一些配置参数给评估器：这些参数的框架位于 eval/in_memory_hf_eval.yaml 文件中。然后只需运行以下脚本，确保将其指向您训练的模型的检查点（以及对应的配置 .json 文件）：
cd eval
python eval_openlm_ckpt.py \
--eval-yaml in_memory_hf_eval.yaml \
--model open_lm_1b \
--checkpoint /path/to/openlm_checkpoint.pt
--positional_embedding_type head_rotary
请注意，--positional-embedding-type head_rotary 仅在使用下面托管的预训练 open_lm_1b 模型时才需要。关于这一点的讨论请参阅下一节。
生成文本
也可以使用训练好的模型生成文本。这可以通过位于 scripts/generate.py 的脚本实现。其参数与评估时使用的类似：
cd scripts
python generate.py \
--model open_lm_1b \
--checkpoint /path/to/openlm_checkpoint.pt \
--positional-embedding-type head_rotary \
--input-text "请给我一个巧克力曲奇的食谱"
同样，--positional-embedding-type head_rotary 仅在使用下面托管的预训练 open_lm_1b 模型时才需要。
预训练模型
OpenLM 1B
OpenLM 1B 是一个约 10 亿参数的模型，训练数据集包含 1.6 万亿个标记，由 RedPajama、Pile、S2ORC、The Pile of Law、Deepmind Math 和 RealNews 组成（完整的训练数据组合细节请参阅此处）。该模型检查点可以从 HuggingFace 下载。用于训练该模型的脚本（用于复制配置）位于此处。下载检查点后，您可以按照上面的“评估模型”部分中的说明进行操作，并传递 --positional-embedding-type head_rotary 或在模型配置中设置 "positional_embedding_type": "head_rotary"（请参阅下面的说明）。
注意：我们在训练该模型时，将旋转位置嵌入应用于 头部 维度，这是截至 2023 年 9 月 1 日 xformers 的默认设置。自这些模型训练完成后，我们已更新 OpenLM，正确地将旋转位置嵌入应用于序列维度（请参阅此问题和此问题了解详情）。要评估这些模型，请确保在模型配置中使用 "positional_embedding_type": "head_rotary"。
| OpenLM-1B | 250B 标记 | 500B 标记 | 750B 标记 | 1T 标记 | 1.25T 标记 | 1.5T 标记 | 1.6T 标记 |
|---|---|---|---|---|---|---|---|
| arc_challenge | 0.27 | 0.28 | 0.29 | 0.28 | 0.29 | 0.31 | 0.31 |
| arc_easy | 0.49 | 0.50 | 0.51 | 0.53 | 0.54 | 0.56 | 0.56 |
| boolq | 0.60 | 0.61 | 0.62 | 0.62 | 0.65 | 0.64 | 0.65 |
| copa | 0.71 | 0.70 | 0.70 | 0.78 | 0.71 | 0.73 | 0.70 |
| hellaswag | 0.50 | 0.54 | 0.54 | 0.57 | 0.59 | 0.61 | 0.61 |
| lambada_openai | 0.56 | 0.57 | 0.61 | 0.61 | 0.65 | 0.65 | 0.66 |
| piqa | 0.70 | 0.70 | 0.71 | 0.72 | 0.73 | 0.74 | 0.74 |
| triviaqa | |||||||
| winogrande | 0.55 | 0.57 | 0.58 | 0.59 | 0.61 | 0.60 | 0.60 |
| MMLU | 0.24 | 0.24 | 0.24 | 0.23 | 0.26 | 0.24 | 0.25 |
| Jeopardy | 0.01 | 0.02 | 0.01 | 0.01 | 0.04 | 0.09 | 0.10 |
| Winograd | 0.75 | 0.77 | 0.77 | 0.79 | 0.81 | 0.80 | 0.79 |
| 平均值 | 0.49 | 0.50 | 0.51 | 0.52 | 0.53 | 0.54 | 0.54 |
| 1B 基线 | OPT-1.3B | Pythia-1B | Neox-1.3B | OPT-IML-1.3B |
|---|---|---|---|---|
| arc_challenge | 0.27 | 0.26 | 0.26 | 0.30 |
| arc_easy | 0.49 | 0.51 | 0.47 | 0.58 |
| boolq | 0.58 | 0.61 | 0.62 | 0.72 |
| copa | 0.75 | 0.68 | 0.72 | 0.73 |
| hellaswag | 0.54 | 0.49 | 0.48 | 0.54 |
| lambada_openai | 0.59 | 0.58 | 0.57 | 0.57 |
| piqa | 0.72 | 0.70 | 0.72 | 0.73 |
| triviaqa | ||||
| winogrande | 0.59 | 0.53 | 0.55 | 0.59 |
| MMLU | 0.25 | 0.26 | 0.26 | 0.30 |
| Jeopardy | 0.01 | 0.00 | 0.00 | 0.12 |
| Winograd | 0.74 | 0.71 | 0.75 | 0.73 |
| 平均值 | 0.50 | 0.48 | 0.49 | 0.54 |
OpenLM 7B
OpenLM 7B 尚未完成训练，但我们已发布了 1.25T 标记的检查点。相关信息与 OpenLM-1B 上述内容相同，包括有关旋转位置嵌入的信息。
| OpenLM-7B | 275B 标记 | 500B 标记 | 675B 标记 | 775B 标记 | 1T 标记 | 1.25T 标记 | 1.5T 标记 | 1.6T 标记 | LLAMA-7B | MPT-7B |
|---|---|---|---|---|---|---|---|---|---|---|
| arc_challenge | 0.35 | 0.35 | 0.36 | 0.37 | 0.39 | 0.39 | 0.41 | 0.39 | ||
| arc_easy | 0.60 | 0.61 | 0.62 | 0.62 | 0.63 | 0.66 | 0.65 | 0.67 | ||
| boolq | 0.67 | 0.66 | 0.69 | 0.69 | 0.70 | 0.70 | 0.77 | 0.75 | ||
| copa | 0.75 | 0.79 | 0.75 | 0.80 | 0.80 | 0.78 | 0.78 | 0.81 | ||
| hellaswag | 0.64 | 0.67 | 0.68 | 0.68 | 0.69 | 0.70 | 0.75 | 0.76 | ||
| lambada_openai | 0.67 | 0.68 | 0.69 | 0.70 | 0.70 | 0.70 | 0.74 | 0.70 | ||
| piqa | 0.75 | 0.76 | 0.76 | 0.76 | 0.77 | 0.77 | 0.79 | 0.80 | ||
| triviaqa | ||||||||||
| winogrande | 0.62 | 0.65 | 0.65 | 0.65 | 0.67 | 0.67 | 0.68 | 0.68 | ||
| MMLU-0 shot | 0.25 | 0.25 | 0.27 | 0.27 | 0.28 | 0.30 | 0.30 | 0.30 | ||
| Jeopardy | 0.15 | 0.18 | 0.23 | 0.22 | 0.16 | 0.21 | 0.33 | 0.31 | ||
| Winograd | 0.82 | 0.81 | 0.84 | 0.84 | 0.85 | 0.86 | 0.81 | 0.88 | ||
| 平均值 | 0.57 | 0.58 | 0.60 | 0.60 | 0.60 | 0.61 | 0.64 | 0.64 | ||
| MMLU-5 shot | 0.34 | 0.34 |
单元测试
对于单元测试，我们使用 pytest。您可以：
或者通过运行以下命令创建 open_lm_tests Conda 环境：
conda env create --file environment-tests.yml
测试位于 tests/ 文件夹中。
运行测试时，请确保您使用的是带有 GPU 的机器，并运行：
团队与致谢
团队（迄今为止，* 表示同等贡献）：Suchin Gururangan*、Mitchell Wortsman*、Samir Yitzhak Gadre*、Achal Dave*、Maciej Kilian、Weijia Shi、Jean Mercat、Georgios Smyrnis、Gabriel Ilharco、Matt Jordan、Reinhard Heckel、Alex Dimakis、Ali Farhadi、Vaishaal Shankar*、Ludwig Schmidt。
代码主要基于由 Ross Wightman、Romain Beaumont、Cade Gordon、Mehdi Cherti、Jenia Jitsev 开发的 open-clip，以及由 Anas Awadalla 和 Irena Gao 开发的 open-flamingo。此外，还受到 lit-llama 的启发。
我们感谢 stability.ai 提供的资源支持。
OpenLM 由来自不同机构的研究人员开发，包括华盛顿大学的 RAIVN 实验室、UWNLP、丰田研究所、哥伦比亚大学等。
引用
如果您的工作中使用了该模型，请使用以下 BibTeX 引用：
@misc{open_lm,
author = {Gururangan, Suchin and Wortsman, Mitchell and Gadre, Samir Yitzhak and Dave, Achal and Kilian, Maciej and Shi, Weijia and Mercat, Jean and Smyrnis, Georgios and Ilharco, Gabriel and Jordan, Matt and Heckel, Reinhard and Dimakis, Alex and Farhadi, Ali and Shankar, Vaishaal and Schmidt, Ludwig},
title = {{open_lm}: 一个简洁但性能出色的语言建模（LM）代码库},
year = {2023},
note = {GitHub 代码库},
url = {https://github.com/mlfoundations/open_lm/}
}
浙公网安备 33010602011771号