From 5fba2417993f4392515d7b4ec12c207433fa53e6 Mon Sep 17 00:00:00 2001 From: WangJie <2740469261@qq.com> Date: Sat, 8 Aug 2026 11:20:48 +0800 Subject: [PATCH 1/2] feat: olmo3 1b and scripts --- .gitignore | 3 + reproduce/.gitkeep | 0 reproduce/olmo-core-backend/README.md | 327 ++++++++++++++++++ reproduce/olmo-core-backend/README_zh.md | 313 +++++++++++++++++ .../cfgs/OLMo3-1B-long-context.py | 165 +++++++++ .../cfgs/OLMo3-1B-midtraining.py | 56 +++ .../cfgs/OLMo3-1B-pretrain.py | 26 ++ reproduce/olmo-core-backend/cfgs/_olmo3_1b.py | 230 ++++++++++++ reproduce/olmo-core-backend/requirements.txt | 4 + .../olmo-core-backend/run/envs.sh.example | 24 ++ reproduce/olmo-core-backend/run/run.sh | 129 +++++++ 11 files changed, 1277 insertions(+) delete mode 100644 reproduce/.gitkeep create mode 100644 reproduce/olmo-core-backend/README.md create mode 100644 reproduce/olmo-core-backend/README_zh.md create mode 100644 reproduce/olmo-core-backend/cfgs/OLMo3-1B-long-context.py create mode 100644 reproduce/olmo-core-backend/cfgs/OLMo3-1B-midtraining.py create mode 100644 reproduce/olmo-core-backend/cfgs/OLMo3-1B-pretrain.py create mode 100644 reproduce/olmo-core-backend/cfgs/_olmo3_1b.py create mode 100644 reproduce/olmo-core-backend/requirements.txt create mode 100755 reproduce/olmo-core-backend/run/envs.sh.example create mode 100755 reproduce/olmo-core-backend/run/run.sh diff --git a/.gitignore b/.gitignore index 83972fa..694a92c 100644 --- a/.gitignore +++ b/.gitignore @@ -216,3 +216,6 @@ __marimo__/ # Streamlit .streamlit/secrets.toml + +# Machine-local OLMo reproduction settings (paths and optional credentials). +/reproduce/olmo-core-backend/run/envs.sh diff --git a/reproduce/.gitkeep b/reproduce/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/reproduce/olmo-core-backend/README.md b/reproduce/olmo-core-backend/README.md new file mode 100644 index 0000000..4ca9a30 --- /dev/null +++ b/reproduce/olmo-core-backend/README.md @@ -0,0 +1,327 @@ +# OLMo 3 1B Three-Stage Reproduction + +English | [中文](README_zh.md) + +This directory provides the training recipes and launcher for the three-stage OLMo 3 1B pipeline: + +1. stage 1: pretraining; +2. stage 2: midtraining; +3. stage 3: long-context extension. + +The model implementation, distributed trainer, checkpoint I/O, and dataset implementation all come from +OLMo-core. This directory contains only the configuration and entry point required for a specific +reproduction, keeping the experiment recipe separate from the general-purpose training framework. Changes +to recipes in this directory do not affect OLMo-core, while framework upgrades and fixes do not require +copying the entire framework into this repository. + +This reproduction must use the following custom OLMo-core source branch instead of the general PyPI release: + +- [https://github.com/JT-Ushio/OLMo-core-muon-fix/tree/ready_for_archspace_base](https://github.com/JT-Ushio/OLMo-core-muon-fix/tree/ready_for_archspace_base) + +This branch contains the Muon fixes required by the recipes. The OLMo 3 model is already implemented by +`TransformerConfig.olmo3_1B()`, so no additional modeling files are needed here. + +## 1. Directory structure + +```text +reproduce/olmo-core-backend/ +├── README.md +├── README_zh.md +├── requirements.txt +├── cfgs/ +│ ├── _olmo3_1b.py +│ ├── OLMo3-1B-pretrain.py +│ ├── OLMo3-1B-midtraining.py +│ └── OLMo3-1B-long-context.py +└── run/ + ├── envs.sh.example + └── run.sh +``` + +## 2. Recipe overview + + +| Stage | Data mix | Sequence length | Global batch (tokens) | Parallelism | Default Muon LR | +| --------- | ---------------------------------- | ----------------: | ----------------------: | ------------- | ----------------: | +| stage 1 | `OLMo_mix_0625_150Bsample` | 4,096 | 2,097,152 | HSDP | `1e-3` | +| stage 2 | `OLMo_midtraining_mix_0625_100B` | 4,096 | 2,097,152 | HSDP | `5e-4` | +| stage 3 | `OLMo_longmino_mix_0625` | 65,536 | 4,194,304 | HSDP | `5e-4` | + +All three stages use BF16, FlashAttention-3, and Muon by default, and each trains for one complete data +epoch. Stages 1 and 2 use fixed-length datasets. Stage 3 uses document packing, an intra-document attention +mask, and 8x YaRN RoPE scaling. All three stages use HSDP without context parallelism. +The default FlashAttention-3 configuration targets Hopper GPUs. Other supported GPUs should switch to +FlashAttention-2 as described in section 3.2. + +You can also select the SkipStep AdamW recipe with `adam`, but all three stages and every resume attempt in a +pipeline must use the same optimizer because later stages inherit the optimizer state from the preceding +stage. + +These are experimental configurations scaled from the official OLMo 3 7B recipes to the 1B model. They are +not officially released or tuned OLMo 3 1B recipes. + +## 3. Environment setup + +This project uses PyTorch 2.10 and CUDA 12.8, although any versions compatible with OLMo-core and the other +dependencies should work in principle. + +```bash +pip install --index-url https://download.pytorch.org/whl/cu128 \ + torch==2.10.0 torchvision torchaudio +``` + +### 3.1 Install the custom OLMo-core from source + +We recommend keeping a separate source checkout and installing it in editable mode: + +```bash +export OLMO_CORE_SRC=/path/to/OLMo-core-muon-fix +git clone --branch ready_for_archspace_base --single-branch \ + https://github.com/JT-Ushio/OLMo-core-muon-fix.git "${OLMO_CORE_SRC}" + +pip install -e "${OLMO_CORE_SRC}[all]" +``` + +### 3.2 Install attention kernels + +Install the latest attention kernels without pinning a tag, commit, or package version. First install +FlashAttention-2: + +```bash +MAX_JOBS=8 python -m pip install --upgrade --no-build-isolation flash-attn +``` + +Hopper GPUs such as H100 and H800 can use FlashAttention-3. Install it from the `hopper/` directory on the +default FlashAttention branch: + +```bash +export FLASH_ATTN_SRC=/path/to/flash-attention +git clone --depth 1 --recurse-submodules --shallow-submodules \ + https://github.com/Dao-AILab/flash-attention.git "${FLASH_ATTN_SRC}" + +cd "${FLASH_ATTN_SRC}/hopper" +FLASH_ATTENTION_DISABLE_FP16=TRUE \ +FLASH_ATTENTION_DISABLE_SM80=TRUE \ +MAX_JOBS=8 \ +python setup.py install +cd - +``` + +Other supported GPUs should use FlashAttention-2. All three upstream-synchronized cfgs select `flash_3` by +default. When using FA2, add the following override to the `extra_args` array in `run/run.sh` so it applies to +all three stages: + +```bash +"--model.attn_backend=flash_2" +``` + +`ring-flash-attn` remains available as an optional backend. Install it when a custom configuration enables +ring context parallelism; the current three-stage HSDP recipe without CP does not require it: + +```bash +pip install ring-flash-attn +``` + +Adjust `MAX_JOBS` to the CPU and memory available on the build node. Run the checks that correspond to the +backends you installed: + +```bash +python -m pip check +# FA2 (all installations) +python -c 'from olmo_core.nn.attention.flash_attn_api import has_flash_attn_2; assert has_flash_attn_2()' +# FA3 (Hopper only) +python -c 'from olmo_core.nn.attention.flash_attn_api import has_flash_attn_3; assert has_flash_attn_3()' +# ring-flash-attn (optional installation only) +python -c 'from olmo_core.nn.attention.flash_attn_api import has_ring_flash_attn; assert has_ring_flash_attn()' +python -c 'import dion, torch; print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name(0))' +``` + +## 4. Data preparation + +### 4.1 Data format + +The configurations directly use four `DataMix` manifests installed with the OLMo-core package. Every `.npy` +path listed by these manifests must be a one-dimensional token-ID binary array following the OLMo-core +convention and readable with `numpy.memmap`. Arbitrary text files, or files merely renamed to `.npy`, will not +work. The Dolma 2 tokenizer has a vocabulary size of 100,278, so these recipes infer the array dtype as +`uint32`. Documents must be correctly separated with the EOS token (ID `100257`), because stage 3 document +packing and intra-document masking depend on these boundaries. + +The data comes from the official OLMo release. This repository will provide a Hugging Face redistribution: +LINK TODO. + +### 4.2 Expected `olmo3_data_root` layout + +`run.sh` passes `olmo3_data_root` from `envs.sh` unchanged to all three configurations. OLMo-core then uses it +as the prefix for every relative path in the manifests. The approximate directory layout is shown below; +ellipses represent all sources and shards listed in the manifests: + +```text +olmo3_data_root/ +├── preprocessed/ +│ ├── dolma2-0625/v0.1-150b/ +│ │ └── allenai/dolma2-tokenizer/ +│ │ ├── finemath-3plus/part-000-00000.npy +│ │ └── ... +│ ├── dolma3-dolmino-official/100B/ +│ │ └── allenai/dolma3-tokenizer/ +│ │ ├── code-meta-reasoning/part-00-00000.npy +│ │ └── ... +│ └── dolma3_longmino_0625/ +│ └── allenai/dolma3-tokenizer/ +│ ├── 000000.npy +│ └── ... +└── eval-data/perplexity/ + └── v3_small_dolma2-tokenizer/ + ├── c4_en/val/part-0-00000.npy + ├── dolma_books/val/part-0-00000.npy + └── ... +``` + +Stage 1 uses the first tree, stage 2 the second, and stage 3 the third. The in-loop LM evaluations in stages 1 +and 2 also require the final validation tree. Every manifest filename must match exactly; providing only +similar top-level directories is insufficient. + +`tokenizer_json` is another required path. It must point to a Dolma 2 `tokenizer.json` readable by every node +and is used by the stage 1 and 2 in-loop downstream evaluator. It does not replace the tokenized training +arrays described above. + +## 5. Run training + +```bash +cd reproduce/olmo-core-backend +cp run/envs.sh.example run/envs.sh +# edit run/envs.sh +bash run/run.sh +``` + +`run/envs.sh` is excluded by `.gitignore`. + +### 5.1 W&B + +`envs.sh.example` sets `WANDB_MODE=offline` by default. This mode requires no API key, writes files under each +stage's `trainer/wandb/` directory, and automatically disables remote cancel tags, which only work online. + +The timestamp is used only in the W&B run name and ID. Keep `pipeline_name` unchanged when resuming the same +experiment, but use a new timestamp for each new job attempt to prevent the new W&B segment from overwriting +or mixing with the previous attempt. Every node in the same multi-node attempt must use the same timestamp. + +### 5.2 Output structure + +```text +out_root/ +├── dataset-cache/ +│ ├── olmo3-stage1/... +│ ├── olmo3-stage2/... +│ └── olmo3-stage3/... +└── runs/olmo3-1b/ + ├── stage1/ + │ ├── _SUCCESS + │ ├── checkpoints/ + │ │ └── step/ + │ │ ├── .metadata.json + │ │ ├── config.json + │ │ ├── data_paths.txt + │ │ ├── model_and_optim/ + │ │ │ ├── .metadata + │ │ │ └── ___.distcp + │ │ └── train/ + │ │ └── rank.pt + │ └── trainer/wandb/... + ├── stage2/ + │ └── ... + └── stage3/ + └── ... +``` + +`config.json` is the effective configuration, while `data_paths.txt` records the expanded data files that +were actually used. Preserve both together with the W&B records when archiving a reproduction run. The +configurations write a temporary checkpoint approximately every 1 billion tokens, retain only one temporary +checkpoint, and save a final checkpoint at the end of each stage. + +Node rank 0 creates `_SUCCESS` after `torchrun` exits successfully for that stage. It indicates successful +process completion; it does not revalidate the checkpoint step or metric values. + +### 5.3 Resume and stage transitions + +A normal resume does not require specifying a checkpoint manually: + +```bash +# Keep out_root and pipeline_name unchanged; use a new attempt timestamp. +bash run/run.sh 0809_093000 +``` + +The launcher and OLMo-core behave as follows: + +1. If `stageN/_SUCCESS` exists, that stage is skipped. +2. If `_SUCCESS` does not exist but the current stage has a checkpoint under `checkpoints/`, the model, + optimizer, trainer, data-loader, and RNG states are restored from it. +3. If the current stage 2 or 3 has no checkpoint, the top-level `--load_path` initializes the model and + optimizer from the preceding stage's checkpoint without inheriting that stage's step or epoch progress. +4. If the current stage 1 has no checkpoint, training starts from scratch. + +Therefore, rerunning after a stage 2 interruption skips the completed stage 1 and continues from stage 2's +own latest checkpoint. Stage 3 starts only after stage 2 finishes. + +## 6. Evaluate with OLMES + +Training produces OLMo-core distributed checkpoints, while OLMES expects a Hugging Face model directory for +a local model. Convert the checkpoint first, then run OLMES. Normally, you evaluate the final stage 3 +checkpoint. To compare stages, convert stages 1, 2, and 3 separately. + +### 6.1 Convert to Hugging Face format + +Select a specific `step` directory, not its parent `checkpoints/` directory: + +```bash +export CHECKPOINT=/path/to/out_root/runs/olmo3-1b/stage3/checkpoints/step11921 +export HF_MODEL_DIR=/path/to/out_root/hf/olmo3-1b-stage3-step11921 + +python "${OLMO_CORE_SRC}/src/examples/huggingface/convert_checkpoint_to_hf.py" \ + --checkpoint-input-path "${CHECKPOINT}" \ + --huggingface-output-dir "${HF_MODEL_DIR}" \ + --max-sequence-length 65536 +``` + +The converter reconstructs the OLMo 3 architecture from the checkpoint's `config.json` and uses +`allenai/dolma2-tokenizer` from the configuration by default. In an offline environment, additionally pass +`--tokenizer /path/to/local/hf-tokenizer-directory`. This must be a complete directory loadable by +`AutoTokenizer.from_pretrained()`, not an individual `tokenizer.json` file. + +Numerical validation is enabled by default. Avoid `--skip-validation` unless you have validated the result +separately and explicitly accept the risk. After conversion, run a minimal loading test: + +```bash +python -c 'import os; from transformers import AutoModelForCausalLM, AutoTokenizer; p=os.environ["HF_MODEL_DIR"]; AutoTokenizer.from_pretrained(p); AutoModelForCausalLM.from_pretrained(p); print("HF checkpoint OK")' +``` + +### 6.2 Install and run OLMES + +Use a separate evaluation environment so that the vLLM and Transformers versions do not affect the training +environment: + +```bash +git clone https://github.com/allenai/olmes.git /path/to/olmes +cd /path/to/olmes +python -m pip install -e '.[gpu]' +git rev-parse HEAD +``` + +For a small-scale experiment, start with the OLMo 3 base-easy suites: + +```bash +olmes \ + --model "${HF_MODEL_DIR}" \ + --task \ + olmo3:base_easy:code_bpb \ + olmo3:base_easy:math_bpb \ + olmo3:base_easy:qa_rc \ + olmo3:base_easy:qa_bpb \ + --output-dir /path/to/out_root/evals/olmo3-1b-stage3-base-easy +``` + +If the installed OLMES/vLLM versions support this model, add `--model-type vllm` for higher throughput. A +formal report should preserve the OLMES commit, complete command, task suite, checkpoint step, Hugging Face +conversion arguments, and output directory. The `FAST_TASKS` and PPL in-loop evaluations built into the +training configuration are intended for training monitoring and do not replace a final, version-pinned OLMES +evaluation. diff --git a/reproduce/olmo-core-backend/README_zh.md b/reproduce/olmo-core-backend/README_zh.md new file mode 100644 index 0000000..bfe728d --- /dev/null +++ b/reproduce/olmo-core-backend/README_zh.md @@ -0,0 +1,313 @@ +# OLMo 3 1B 三阶段复现 + +[English](README.md) | 中文 + +本目录提供 OLMo 3 1B 的三阶段训练配方与启动脚本： + +1. stage 1：pretraining； +2. stage 2：midtraining； +3. stage 3：long-context extension。 + +模型实现、分布式训练器、checkpoint I/O 和数据集实现均来自 OLMo-core。本目录只保存某次 +复现所需的配置和运行入口，以便把“实验配方”与“通用训练框架”隔离开：修改本目录中的配方 +不会污染 OLMo-core，升级或修复训练框架也不需要把整个框架复制进本仓库。 + +本复现必须使用以下自定义 OLMo-core 源码分支，而不是 PyPI 上的通用版本： + +- [https://github.com/JT-Ushio/OLMo-core-muon-fix/tree/ready_for_archspace_base](https://github.com/JT-Ushio/OLMo-core-muon-fix/tree/ready_for_archspace_base) + +该分支包含本配方所依赖的 Muon 修复。OLMo 3 模型本身已由 +`TransformerConfig.olmo3_1B()` 实现，因此这里没有额外的 modeling 文件。 + +## 1. 目录结构 + +```text +reproduce/olmo-core-backend/ +├── README.md +├── README_zh.md +├── requirements.txt +├── cfgs/ +│ ├── _olmo3_1b.py +│ ├── OLMo3-1B-pretrain.py +│ ├── OLMo3-1B-midtraining.py +│ └── OLMo3-1B-long-context.py +└── run/ + ├── envs.sh.example + └── run.sh +``` + +## 2. 配方概览 + + +| 阶段 | 数据 mix | 序列长度 | 全局 batch（token） | 并行方式 | 默认 Muon LR | +| --------- | ---------------------------------- | ---------: | --------------------: | ---------- | -------------: | +| stage 1 | `OLMo_mix_0625_150Bsample` | 4,096 | 2,097,152 | HSDP | `1e-3` | +| stage 2 | `OLMo_midtraining_mix_0625_100B` | 4,096 | 2,097,152 | HSDP | `5e-4` | +| stage 3 | `OLMo_longmino_mix_0625` | 65,536 | 4,194,304 | HSDP | `5e-4` | + +三个阶段默认都使用 BF16、FlashAttention-3 和 Muon，并各自训练一个完整数据 +epoch。stage 1/2 使用固定长度数据集；stage 3 使用 document packing、文档内 attention +mask 和 8 倍 YaRN RoPE scaling。三个阶段均使用 HSDP，不使用 context parallelism。 +默认的 FlashAttention-3 配置面向 Hopper GPU；其他支持的 GPU 应按 3.2 节切换到 FlashAttention-2。 + +也可以用 `adam` 选择 SkipStep AdamW 配方，但同一 pipeline 的三个阶段及所有 resume 必须使用 +同一种 optimizer，因为后续阶段会继承前一阶段的 optimizer state。 + +这些是从 OLMo 3 7B 官方配方缩放到 1B 模型的实验配置，并不是官方发布、已调优的 OLMo 3 +1B recipe。 + +## 3. 环境安装 + +本项目采用 PyTorch 2.10 和 CUDA 12.8，但原则上也可使用任何与 OLMo-core 及其他依赖兼容的版本。 + +```bash +pip install --index-url https://download.pytorch.org/whl/cu128 \ + torch==2.10.0 torchvision torchaudio +``` + +### 3.1 从源码安装自定义 OLMo-core + +推荐保留一个独立源码 checkout，并以 editable 方式安装： + +```bash +export OLMO_CORE_SRC=/path/to/OLMo-core-muon-fix +git clone --branch ready_for_archspace_base --single-branch \ + https://github.com/JT-Ushio/OLMo-core-muon-fix.git "${OLMO_CORE_SRC}" + +pip install -e "${OLMO_CORE_SRC}[all]" +``` + +### 3.2 安装 attention kernels + +本节的 attention kernels 都直接安装最新版，不固定 tag、commit 或 package +version。先安装 FlashAttention-2： + +```bash +MAX_JOBS=8 python -m pip install --upgrade --no-build-isolation flash-attn +``` + +Hopper GPU（例如 H100/H800）可以使用 FlashAttention-3，从 FlashAttention 默认分支的 +`hopper/` 目录安装： + +```bash +export FLASH_ATTN_SRC=/path/to/flash-attention +git clone --depth 1 --recurse-submodules --shallow-submodules \ + https://github.com/Dao-AILab/flash-attention.git "${FLASH_ATTN_SRC}" + +cd "${FLASH_ATTN_SRC}/hopper" +FLASH_ATTENTION_DISABLE_FP16=TRUE \ +FLASH_ATTENTION_DISABLE_SM80=TRUE \ +MAX_JOBS=8 \ +python setup.py install +cd - +``` + +其他支持的 GPU 使用 FlashAttention-2。三个 cfg 从上游同步的默认 backend 都是 +`flash_3`；使用 FA2 时，在 `run/run.sh` 的 `extra_args` 数组中加入以下覆盖，使它同时 +作用于三个阶段： + +```bash +"--model.attn_backend=flash_2" +``` + +`ring-flash-attn` 作为可选 backend 保留。当自定义配置启用 ring context parallelism 时 +再安装；当前 HSDP、无 CP 的三阶段配方不需要它： + +```bash +pip install ring-flash-attn +``` + +`MAX_JOBS` 应按编译节点的 CPU 和内存调整。可按实际选择的 backend 分别验证： + +```bash +python -m pip check +# FA2（所有安装） +python -c 'from olmo_core.nn.attention.flash_attn_api import has_flash_attn_2; assert has_flash_attn_2()' +# FA3（仅 Hopper） +python -c 'from olmo_core.nn.attention.flash_attn_api import has_flash_attn_3; assert has_flash_attn_3()' +# ring-flash-attn（仅可选安装） +python -c 'from olmo_core.nn.attention.flash_attn_api import has_ring_flash_attn; assert has_ring_flash_attn()' +python -c 'import dion, torch; print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name(0))' +``` + +## 4. 数据准备 + +### 4.1 数据格式 + +配置直接使用安装在 OLMo-core 包中的四份 `DataMix` manifest。它们列出的每个 `.npy` 路径是 +OLMo-core 约定的、可由 `numpy.memmap` 读取的一维 token-ID 二进制数组，而不是任意文本文件， +也不能只靠把文件改名为 `.npy` 得到。Dolma 2 tokenizer 的词表大小为 100,278，因此本配方会 +推断数组 dtype 为 `uint32`。不同文档需要以 EOS token（ID `100257`）正确分隔，stage 3 的 +document packing 和文档内 mask 依赖这些边界。 + +数据来自olmo官方，本仓库提供 Huggingface 再发布版本：链接TODO + +### 4.2 `olmo3_data_root` 的预期布局 + +`run.sh` 把 `envs.sh` 中的 `olmo3_data_root` 原样传给三个配置。OLMo-core 再把它作为 manifest +内所有相对路径的前缀。大致目录如下；省略号代表 manifest 中的全部 source 和 shard： + +```text +olmo3_data_root/ +├── preprocessed/ +│ ├── dolma2-0625/v0.1-150b/ +│ │ └── allenai/dolma2-tokenizer/ +│ │ ├── finemath-3plus/part-000-00000.npy +│ │ └── ... +│ ├── dolma3-dolmino-official/100B/ +│ │ └── allenai/dolma3-tokenizer/ +│ │ ├── code-meta-reasoning/part-00-00000.npy +│ │ └── ... +│ └── dolma3_longmino_0625/ +│ └── allenai/dolma3-tokenizer/ +│ ├── 000000.npy +│ └── ... +└── eval-data/perplexity/ + └── v3_small_dolma2-tokenizer/ + ├── c4_en/val/part-0-00000.npy + ├── dolma_books/val/part-0-00000.npy + └── ... +``` + +stage 1 使用第一棵树，stage 2 使用第二棵树，stage 3 使用第三棵树；stage 1/2 的 in-loop LM +evaluation 都需要最后一棵 validation 树。manifest 文件名必须逐项匹配，不能只提供相似的 +顶层目录。 + +`tokenizer_json` 是另一项必填路径：它应指向所有节点都能读取的 Dolma 2 `tokenizer.json`， +供 stage 1/2 的 in-loop downstream evaluator 使用。它不替代上述已分词训练数组。 + +## 5. 运行训练 + +```bash +cd reproduce/olmo-core-backend +cp run/envs.sh.example run/envs.sh +# edit run/envs.sh +bash run/run.sh +``` + +`run/envs.sh` 已被 `.gitignore` 排除。 + +### 5.1 W&B + +`envs.sh.example` 默认 `WANDB_MODE=offline`。这种模式不需要 API key，文件写到每个阶段的 +`trainer/wandb/` 下，并自动禁用只能在线工作的 remote cancel tags。 + +timestamp 只参与 W&B run name/ID。resume 同一个 +实验时保持 `pipeline_name` 不变，但每次重新发起任务建议给一个新 timestamp，避免新的 W&B +片段覆盖或混入上一次 attempt。多机同一次 attempt 必须使用同一个 timestamp。 + +### 5.2 输出目录 + +```text +out_root/ +├── dataset-cache/ +│ ├── olmo3-stage1/... +│ ├── olmo3-stage2/... +│ └── olmo3-stage3/... +└── runs/olmo3-1b/ + ├── stage1/ + │ ├── _SUCCESS + │ ├── checkpoints/ + │ │ └── step/ + │ │ ├── .metadata.json + │ │ ├── config.json + │ │ ├── data_paths.txt + │ │ ├── model_and_optim/ + │ │ │ ├── .metadata + │ │ │ └── ___.distcp + │ │ └── train/ + │ │ └── rank.pt + │ └── trainer/wandb/... + ├── stage2/ + │ └── ... + └── stage3/ + └── ... +``` + +`config.json` 是最终生效配置，`data_paths.txt` 记录实际展开的数据文件；复现实验归档时应和 W&B +记录一起保存。配置约每 10 亿 token 写一次临时 checkpoint，只保留一个临时版本，并在阶段 +结束时保存最终 checkpoint。 + +`_SUCCESS` 由 node rank 0 在该阶段 `torchrun` 成功退出之后创建。它表示进程成功完成，不会 +再次检查 checkpoint step 或指标内容。 + +### 5.3 Resume 与阶段衔接 + +正常 resume 不需要手工指定 checkpoint： + +```bash +# out_root、pipeline_name 保持不变；使用新的 attempt timestamp。 +bash run/run.sh 0809_093000 +``` + +启动器和 OLMo-core 的行为是： + +1. 存在 `stageN/_SUCCESS`：直接跳过该阶段； +2. 不存在 `_SUCCESS`，但当前阶段 `checkpoints/` 中有 checkpoint：恢复该阶段的 model、 + optimizer、trainer、data-loader 和 RNG 状态； +3. 当前 stage 2/3 没有 checkpoint：通过顶层 `--load_path` 从前一阶段 checkpoint 初始化 model + 和 optimizer，但不继承前一阶段的 step/epoch 进度； +4. 当前 stage 1 没有 checkpoint：从头开始。 + +因此 stage 2 中断后的重跑会跳过已完成的 stage 1，并从 stage 2 自己的最新 checkpoint 继续； +stage 2 完成后才进入 stage 3。 + +## 6. 使用 OLMES 评测 + +训练输出是 OLMo-core distributed checkpoint，而 OLMES 的本地模型入口使用 Hugging Face +模型目录。因此先转换 checkpoint，再运行 OLMES。通常评测 stage 3 的最终 checkpoint；若要画 +阶段对比，则分别转换 stage 1/2/3。 + +### 6.1 转换为 Hugging Face 格式 + +选中具体的 `step` 目录，而不是它的 `checkpoints/` 父目录： + +```bash +export CHECKPOINT=/path/to/out_root/runs/olmo3-1b/stage3/checkpoints/step11921 +export HF_MODEL_DIR=/path/to/out_root/hf/olmo3-1b-stage3-step11921 + +python "${OLMO_CORE_SRC}/src/examples/huggingface/convert_checkpoint_to_hf.py" \ + --checkpoint-input-path "${CHECKPOINT}" \ + --huggingface-output-dir "${HF_MODEL_DIR}" \ + --max-sequence-length 65536 +``` + +转换器会从 checkpoint 的 `config.json` 恢复 OLMo 3 架构，并默认使用配置中的 +`allenai/dolma2-tokenizer`。离线环境可额外传 +`--tokenizer /path/to/local/hf-tokenizer-directory`；这里应给一个可由 +`AutoTokenizer.from_pretrained()` 加载的完整目录，而不是单独的 `tokenizer.json`。 + +默认转换包含数值验证。除非已经单独验证且明确接受风险，不建议使用 `--skip-validation`。 +转换后可先做最小加载测试： + +```bash +python -c 'import os; from transformers import AutoModelForCausalLM, AutoTokenizer; p=os.environ["HF_MODEL_DIR"]; AutoTokenizer.from_pretrained(p); AutoModelForCausalLM.from_pretrained(p); print("HF checkpoint OK")' +``` + +### 6.2 安装并运行 OLMES + +建议为评测创建独立环境，避免 vLLM/Transformers 版本反向影响训练环境： + +```bash +git clone https://github.com/allenai/olmes.git /path/to/olmes +cd /path/to/olmes +python -m pip install -e '.[gpu]' +git rev-parse HEAD +``` + +小规模实验可从 OLMo 3 base-easy suites 开始： + +```bash +olmes \ + --model "${HF_MODEL_DIR}" \ + --task \ + olmo3:base_easy:code_bpb \ + olmo3:base_easy:math_bpb \ + olmo3:base_easy:qa_rc \ + olmo3:base_easy:qa_bpb \ + --output-dir /path/to/out_root/evals/olmo3-1b-stage3-base-easy +``` + +如 OLMES/vLLM 版本支持该模型，可加 `--model-type vllm` 提高吞吐。正式报告中应保留 OLMES +commit、完整命令、task suite、checkpoint step、HF 转换参数和输出目录。训练配置内置的 +`FAST_TASKS`/PPL in-loop evaluation 用于训练监控，不能替代最终、版本固定的 OLMES 评测。 diff --git a/reproduce/olmo-core-backend/cfgs/OLMo3-1B-long-context.py b/reproduce/olmo-core-backend/cfgs/OLMo3-1B-long-context.py new file mode 100644 index 0000000..1532afb --- /dev/null +++ b/reproduce/olmo-core-backend/cfgs/OLMo3-1B-long-context.py @@ -0,0 +1,165 @@ +""" +OLMo 3 1B stage-3 long-context extension configuration. + +This is a 1B adaptation of the OLMo 3 7B long-context recipe in +`src/scripts/official/OLMo3/OLMo-3-1025-7B-long-context.py`. OLMo 3 does not publish an +officially tuned 1B long-context recipe. + +Parallelism boundaries +---------------------- +DP=data-parallel world size; PP/CP/TP/EP are their degrees. +H_rep=HSDP replicas, H_shard=HSDP shard degree, L=seqlen, M=microbatch tokens, +B=global batch tokens; heads=16; n_layers=16. + +Mesh: world_size = PP*CP*TP*DP; world_size % (PP*CP*TP) = 0 +HSDP: DP = H_rep*H_shard; DP % H_shard = 0 +Batch: M % L = 0; B % (M*DP) = 0; grad_accum = B/(M*DP) +CP: local_L = L/CP; exact split requires L % CP = 0 +Ulysses CP: q_heads % CP = kv_heads % CP = 0 +TP: tensor_dim % TP = 0 for every sharded dimension +PP: world_size % PP = 0; num_stages % PP = 0; num_stages <= n_layers +EP: MoE and HSDP only; EP = H_shard; TP = 1 (off) + +Optimizer / parallelism matrix: +| Mode | AdamW | Muon | +|--------------|-----------------|--------------------------------------| +| FSDP | yes | yes: heads % (DP*CP) = 0 | +| HSDP, CP off | yes | yes: heads % H_shard = 0 | +| HSDP + CP | yes | no: dp_shard is flattened into dp_cp | +| TP | yes | no: hard error | +| PP | yes (beta) | beta; changes DP | +| EP | MoE + HSDP only | no: flattened/3D expert parameters | + +Other conflicts: flash_3 has no CP, use flash_2 for CP; +TP + EP is forbidden. Multi-stage PP + tied embeddings is forbidden. +Stage 2/3 optimizer states must have the same optimizer type unless loading is disabled. + +Examples: world_size=64 (GPUs), PP=1 (off), TP=1 (off), heads=16 +B=2^22 (tokens), M=L=65,536 (tokens) +| Optim | DP layout | CP | Muon mesh | Result | +|-------|-------------------------|----|-----------|-------------------------------| +| AdamW | HSDP H_rep=16,H_shard=1 | 4 | - | valid | +| AdamW | HSDP H_rep=8,H_shard=1 | 8 | - | valid | +| Muon | HSDP H_rep=8,H_shard=8 | 1 | 8 | valid:16(heads)%8(mesh)=0 | +| Muon | FSDP DP=64 | 1 | 64 | invalid:16(heads)%64(mesh)!=0 | +| Muon | FSDP DP=16 | 4 | DP*CP=64 | invalid:16(heads)%64(mesh)!=0 | +""" + +import argparse +from typing import List + +from _olmo3_1b import build_common_config, build_optim_config, get_olmo3_1b_cli_parser + +from olmo_core.config import DType +from olmo_core.data import ( + DataMix, + NumpyDataLoaderConfig, + NumpyPackedFSLDatasetConfig, + TokenizerConfig, +) +from olmo_core.distributed.parallel import DataParallelType +from olmo_core.nn.attention import AttentionBackendName +from olmo_core.nn.rope import YaRNRoPEScalingConfig +from olmo_core.nn.transformer import TransformerConfig +from olmo_core.optim import LinearWithWarmup +from olmo_core.script_utils import ExperimentConfig, main +from olmo_core.train.common import LoadStrategy +from olmo_core.train.train_module import ( + TransformerContextParallelConfig, # noqa: F401 - used by the optional cp_config below + TransformerDataParallelConfig, + TransformerDataParallelWrappingStrategy, + TransformerTrainModuleConfig, +) + +DEFAULT_SEQUENCE_LENGTH = 65536 +GLOBAL_BATCH_SIZE = 2**22 # 4M tokens +# MAX_TOKENS = 50_000_000_000 # 50B +# Muon retains the 1B recipe; AdamW follows the official stage-3 schedule. +MUON_LR = 5e-4 +ADAM_LR = 5e-4 +SEED = 4123 + + +def build_config(opts: argparse.Namespace, overrides: List[str]) -> ExperimentConfig: + """Build stage 3 from its required components and the shared trainer.""" + # Long context changes the model, dataset, loader, and train module as whole + # units, so this stage does not mutate the stage-1 versions of those components. + sequence_length = opts.sequence_length or DEFAULT_SEQUENCE_LENGTH + tokenizer_config = TokenizerConfig.dolma2() + + model = TransformerConfig.olmo3_1B( + vocab_size=tokenizer_config.padded_vocab_size(), # pad to a multiple of 128 + attn_backend=AttentionBackendName.flash_3, + ).with_rope_scaling( + YaRNRoPEScalingConfig( + factor=8, + beta_fast=32, + beta_slow=1, + old_context_len=8192, + ) + ) + + dataset = NumpyPackedFSLDatasetConfig.from_data_mix( + DataMix.OLMo_longmino_mix_0625, + mix_base_dir=opts.data_root, + work_dir=opts.work_dir, + tokenizer=tokenizer_config, + sequence_length=sequence_length, + generate_doc_lengths=True, # enables intra-document masking + source_group_size=8, + source_permutation_seed=123, + ) + + data_loader = NumpyDataLoaderConfig( + global_batch_size=GLOBAL_BATCH_SIZE, + seed=SEED, + num_workers=8, + prefetch_factor=4, + ) + + train_module = TransformerTrainModuleConfig( + rank_microbatch_size=sequence_length, + max_sequence_length=sequence_length, + optim=build_optim_config( + opts.optim, + muon_lr=MUON_LR, + adam_lr=ADAM_LR, + ), + scheduler=LinearWithWarmup(warmup=200, alpha_f=0.0), + compile_model=True, + dp_config=TransformerDataParallelConfig( + name=DataParallelType.hsdp, + param_dtype=DType.bfloat16, + reduce_dtype=DType.float32, + wrapping_strategy=TransformerDataParallelWrappingStrategy.full, + ), + # cp_config=TransformerContextParallelConfig.llama3(degree=4, head_stride=4), + ac_config=None, + float8_config=None, + # float8_config=Float8Config(enabled=True, ao=AOFloat8LinearConfig.recommended()), + z_loss_multiplier=1e-5, + max_grad_norm=1.0, + ) + + # Only the trainer and its common callbacks are inherited from stage 1. + config = build_common_config( + opts, + model=model, + dataset=dataset, + data_loader=data_loader, + train_module=train_module, + ) + + config.trainer.load_strategy = LoadStrategy.always + # script_utils.main probes save_folder before Trainer.fit() and uses this value, so + # require trainer state for a same-stage resume. The launcher supplies the parent + # stage through ExperimentConfig.load_path, which explicitly skips trainer state. + config.trainer.load_trainer_state = True + config.trainer.load_optim_state = True + + config.init_seed = SEED + return config.merge(overrides) + + +if __name__ == "__main__": + main(build_config, parser=get_olmo3_1b_cli_parser()) diff --git a/reproduce/olmo-core-backend/cfgs/OLMo3-1B-midtraining.py b/reproduce/olmo-core-backend/cfgs/OLMo3-1B-midtraining.py new file mode 100644 index 0000000..6e30a76 --- /dev/null +++ b/reproduce/olmo-core-backend/cfgs/OLMo3-1B-midtraining.py @@ -0,0 +1,56 @@ +""" +OLMo 3 1B stage-2 midtraining configuration. + +This is a 1B adaptation of the official OLMo-3-1025-7B midtraining recipe in +`src/scripts/official/OLMo3/OLMo-3-1025-7B-midtrain.py`. OLMo 3 does not +publish an officially tuned 1B midtraining recipe, so the data schedule and +optimization settings below intentionally retain the official 7B values. +""" + +import argparse +from typing import List + +from _olmo3_1b import build_optim_config, build_pretrain_config, get_olmo3_1b_cli_parser + +from olmo_core.data import DataMix +from olmo_core.optim import LinearWithWarmup +from olmo_core.script_utils import ExperimentConfig, main +from olmo_core.train.common import LoadStrategy + +# MAX_TOKENS = 100_000_000_000 # 100B +# Muon retains the 1B recipe; AdamW follows the official stage-2 schedule. +MUON_LR = 5e-4 +ADAM_LR = 5e-4 +SEED = 1337 + + +def build_config(opts: argparse.Namespace, overrides: List[str]) -> ExperimentConfig: + """Build stage 2 by applying its differences to the pretraining configuration.""" + config = build_pretrain_config(opts) + + # Model shape, including the padded vocabulary size, batching, and callbacks + # remain identical to stage 1. Only data order and optimization are stage-specific. + config.dataset.mix = DataMix.OLMo_midtraining_mix_0625_100B + config.data_loader.seed = SEED + + # Optimizer state is restored from stage 1, so the selected recipe must match. + config.train_module.optim = build_optim_config( + opts.optim, + muon_lr=MUON_LR, + adam_lr=ADAM_LR, + ) + config.train_module.scheduler = LinearWithWarmup(warmup=0, alpha_f=0.0) + + config.trainer.load_strategy = LoadStrategy.always + # script_utils.main probes save_folder before Trainer.fit() and uses this value, so + # require trainer state for a same-stage resume. The launcher supplies the parent + # stage through ExperimentConfig.load_path, which explicitly skips trainer state. + config.trainer.load_trainer_state = True + config.trainer.load_optim_state = True + + config.init_seed = SEED + return config.merge(overrides) + + +if __name__ == "__main__": + main(build_config, parser=get_olmo3_1b_cli_parser()) diff --git a/reproduce/olmo-core-backend/cfgs/OLMo3-1B-pretrain.py b/reproduce/olmo-core-backend/cfgs/OLMo3-1B-pretrain.py new file mode 100644 index 0000000..a9e993e --- /dev/null +++ b/reproduce/olmo-core-backend/cfgs/OLMo3-1B-pretrain.py @@ -0,0 +1,26 @@ +""" +OLMo 3 1B stage-1 pretraining configuration for the local 150B data sample. + +This is a 1B adaptation of the official OLMo-3-1025-7B stage-1 recipe in +``src/scripts/official/OLMo3/OLMo-3-1025-7B-pretrain-1.py``. OLMo 3 does not +publish an official tuned 1B pretraining recipe, so the batch size, learning +rate, and warmup below intentionally retain the official 7B values. +""" + +import argparse +from typing import List + +from _olmo3_1b import build_pretrain_config, get_olmo3_1b_cli_parser +from olmo_core.script_utils import ExperimentConfig, main + + +def build_config(opts: argparse.Namespace, overrides: List[str]) -> ExperimentConfig: + """Build the OLMo 3 1B stage-1 pretraining configuration.""" + # This complete stage-1 recipe, including the padded vocabulary size, is also + # the baseline imported by stage 2. + # Merge CLI overrides only after the shared defaults have been assembled. + return build_pretrain_config(opts).merge(overrides) + + +if __name__ == "__main__": + main(build_config, parser=get_olmo3_1b_cli_parser()) diff --git a/reproduce/olmo-core-backend/cfgs/_olmo3_1b.py b/reproduce/olmo-core-backend/cfgs/_olmo3_1b.py new file mode 100644 index 0000000..4881adc --- /dev/null +++ b/reproduce/olmo-core-backend/cfgs/_olmo3_1b.py @@ -0,0 +1,230 @@ +"""Shared configuration for the OLMo 3 1B training stages.""" + +import argparse + +from olmo_core.config import DType +from olmo_core.data import ( + DataMix, + NumpyDataLoaderConfig, + NumpyDatasetConfig, + NumpyFSLDatasetConfig, + NumpyPaddedFSLDatasetConfig, + TokenizerConfig, +) +from olmo_core.distributed.parallel import DataParallelType +from olmo_core.eval.task_groups import FAST_TASKS +from olmo_core.float8 import Float8Config +from olmo_core.nn.attention import AttentionBackendName +from olmo_core.nn.transformer import TransformerConfig +from olmo_core.optim import ( + CosWithWarmup, + MuonConfig, + OptimConfig, + OptimGroupOverride, + SkipStepAdamWConfig, +) +from olmo_core.script_utils import ExperimentConfig, get_cli_parser +from olmo_core.train import Duration, TrainerConfig +from olmo_core.train.callbacks import ( + CheckpointerCallback, + CometCallback, + ConfigSaverCallback, + DownstreamEvaluatorCallbackConfig, + LMEvaluatorCallbackConfig, + MonkeyPatcherCallback, + WandBCallback, +) +from olmo_core.train.train_module import ( + TransformerDataParallelConfig, + TransformerDataParallelWrappingStrategy, + TransformerTrainModuleConfig, +) + +DEFAULT_SEQUENCE_LENGTH = 4096 +GLOBAL_BATCH_SIZE = 2**21 # 2M tokens +SEED = 34521 +EVAL_LM_STEPS = 500 # 500 steps (~1B token) for 150B data, 2500 steps (~5B token) for 6T data. +EVAL_DOWN_STEPS = 12500 # 12.5K steps (25B tokens) for 150B data +# Keep the current Muon recipe and the official OLMo 3 AdamW recipe independent. +MUON_LR = 1e-3 +ADAM_LR = 1e-3 + + +def get_olmo3_1b_cli_parser() -> argparse.ArgumentParser: + """Build the CLI parser shared by the OLMo 3 1B stages.""" + parser = get_cli_parser() + parser.add_argument( + "--optim", + choices=("adam", "muon"), + default="muon", + help="Optimizer recipe to use; adam selects SkipStep AdamW (default: muon).", + ) + return parser + + +def build_optim_config( + name: str, + *, + muon_lr: float, + adam_lr: float, +) -> OptimConfig: + """Build the selected optimizer with its stage-specific learning rate.""" + # Equivalent whole-object CLI override; define `lr` in the shell first: + # "--train_module.optim={type: muon, lr: ${lr}, weight_decay: 0.033, betas: [0.9, 0.95]}" + if name == "muon": + return MuonConfig( + lr=muon_lr, + weight_decay=0.033, + betas=(0.9, 0.95), + ) + # Equivalent whole-object CLI override; define `lr` in the shell first: + # "--train_module.optim={type: skip_step_adamw, lr: ${lr}, weight_decay: 0.033, betas: [0.9, 0.95], group_overrides: [{params: [embeddings.weight], opts: {weight_decay: 0.0}}]}" + if name == "adam": + # Match the official OLMo 3 AdamW recipe, including no decay on embeddings. + return SkipStepAdamWConfig( + lr=adam_lr, + weight_decay=0.033, + betas=(0.9, 0.95), + group_overrides=[OptimGroupOverride(params=["embeddings.weight"], opts={"weight_decay": 0.0})], + ) + raise ValueError(f"Unknown optimizer '{name}'") + + +def build_common_config( + opts: argparse.Namespace, + *, + model: TransformerConfig, + dataset: NumpyDatasetConfig, + data_loader: NumpyDataLoaderConfig, + train_module: TransformerTrainModuleConfig, +) -> ExperimentConfig: + """Build an experiment from required stage components and the shared trainer.""" + # Temporary checkpoint approximately every 1B tokens. + ephemeral_save_interval = round(2**30 / data_loader.global_batch_size / 10) * 10 + + trainer = ( + TrainerConfig( + save_folder=opts.save_folder, + work_dir=opts.work_dir, + save_overwrite=True, + metrics_collect_interval=10, + cancel_check_interval=10, + max_duration=Duration.epochs(1), + ) + .with_callback("monkey_patcher", MonkeyPatcherCallback()) + .with_callback( + "checkpointer", + CheckpointerCallback( + save_interval=None, # Only save the final ckpt + ephemeral_save_interval=ephemeral_save_interval, + max_checkpoints=1, + # pre_train_checkpoint=False, + # save_async=False, + ), + ) + .with_callback( + "comet", + CometCallback( + name=opts.name, + cancel_check_interval=10, + enabled=False, + ), + ) + .with_callback( + "wandb", + WandBCallback( + name=opts.name, + cancel_check_interval=10, + enabled=False, + ), + ) + .with_callback("config_saver", ConfigSaverCallback()) + ) + + return ExperimentConfig( + model=model, + dataset=dataset, + data_loader=data_loader, + train_module=train_module, + trainer=trainer, + ) + + +def build_pretrain_config(opts: argparse.Namespace) -> ExperimentConfig: + """Build the OLMo 3 1B stage-1 pretraining configuration.""" + sequence_length = opts.sequence_length or DEFAULT_SEQUENCE_LENGTH + tokenizer = TokenizerConfig.dolma2() + + model = TransformerConfig.olmo3_1B( + vocab_size=tokenizer.padded_vocab_size(), # pad to a multiple of 128 + attn_backend=AttentionBackendName.flash_3, + ) + + dataset = NumpyFSLDatasetConfig.from_data_mix( + DataMix.OLMo_mix_0625_150Bsample, + tokenizer=tokenizer, + mix_base_dir=opts.data_root, + sequence_length=sequence_length, + max_target_sequence_length=max(8192, sequence_length), + work_dir=opts.work_dir, + ) + + data_loader = NumpyDataLoaderConfig( + global_batch_size=GLOBAL_BATCH_SIZE, + seed=SEED, + num_workers=8, + prefetch_factor=2, + ) + + train_module = TransformerTrainModuleConfig( + rank_microbatch_size=4 * DEFAULT_SEQUENCE_LENGTH, + max_sequence_length=sequence_length, + optim=build_optim_config( + opts.optim, + muon_lr=MUON_LR, + adam_lr=ADAM_LR, + ), + scheduler=CosWithWarmup(warmup_steps=2000), + compile_model=True, + dp_config=TransformerDataParallelConfig( + name=DataParallelType.hsdp, + param_dtype=DType.bfloat16, + reduce_dtype=DType.float32, + wrapping_strategy=TransformerDataParallelWrappingStrategy.blocks, + ), + float8_config=Float8Config(enabled=False), + z_loss_multiplier=1e-5, + max_grad_norm=1.0, + ) + + config = build_common_config( + opts, + model=model, + dataset=dataset, + data_loader=data_loader, + train_module=train_module, + ) + config.trainer = config.trainer.with_callback( + "lm_evaluator", + LMEvaluatorCallbackConfig( + eval_dataset=NumpyPaddedFSLDatasetConfig.from_data_mix( + DataMix.v3_small_ppl_validation, + mix_base_dir=opts.data_root, + sequence_length=sequence_length, + tokenizer=tokenizer, + work_dir=opts.work_dir, + ), + eval_interval=EVAL_LM_STEPS, + # eval_interval=50, + ), + ).with_callback( + "downstream_evaluator", + DownstreamEvaluatorCallbackConfig( + tasks=sorted(FAST_TASKS), + tokenizer=tokenizer, + eval_interval=EVAL_DOWN_STEPS, + # eval_interval=50, + ), + ) + config.init_seed = SEED + return config diff --git a/reproduce/olmo-core-backend/requirements.txt b/reproduce/olmo-core-backend/requirements.txt new file mode 100644 index 0000000..f9bf4d3 --- /dev/null +++ b/reproduce/olmo-core-backend/requirements.txt @@ -0,0 +1,4 @@ +# Install OLMo-core from the source branch that contains the Muon fixes used by +# this reproduction. Hardware-specific FlashAttention must be installed +# separately; see README.md. +ai2-olmo-core[all] @ git+https://github.com/JT-Ushio/OLMo-core-muon-fix.git@ready_for_archspace_base diff --git a/reproduce/olmo-core-backend/run/envs.sh.example b/reproduce/olmo-core-backend/run/envs.sh.example new file mode 100755 index 0000000..0bb8d15 --- /dev/null +++ b/reproduce/olmo-core-backend/run/envs.sh.example @@ -0,0 +1,24 @@ +#!/usr/bin/env bash + +# Copy this file to envs.sh and replace the placeholders with paths available +# on every node. envs.sh is ignored by Git because it is machine-specific. + +# Root containing every relative path referenced by OLMo-core's built-in OLMo 3 +# stage-1, stage-2, stage-3, and perplexity-validation data-mix manifests. +olmo3_data_root=/path/to/olmo3-data + +# Checkpoints, trainer artifacts, W&B files, and dataset caches are written here. +# Use shared, persistent storage for multi-node training and resume. +out_root=/path/to/training-output + +# Local Dolma 2 tokenizer JSON used by the in-loop downstream evaluator. The LM +# dataset tokenizer metadata still comes from TokenizerConfig.dolma2(). +tokenizer_json=/path/to/tokenizer.json + +# W&B destination. Offline mode is the safe default and does not need an API key. +wandb_entity=YOUR_ENTITY +wandb_project=YOUR_PROJECT +export WANDB_MODE=${WANDB_MODE:-offline} + +# For online logging, export the secret in the calling shell; do not put it here. +# export WANDB_API_KEY=YOUR_SECRET diff --git a/reproduce/olmo-core-backend/run/run.sh b/reproduce/olmo-core-backend/run/run.sh new file mode 100755 index 0000000..a9d3ec6 --- /dev/null +++ b/reproduce/olmo-core-backend/run/run.sh @@ -0,0 +1,129 @@ +#!/usr/bin/env bash +{ + set -euo pipefail + + script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) + reproduce_dir=$(cd -- "${script_dir}/.." && pwd) + env_file=${script_dir}/envs.sh + [[ -f "${env_file}" ]] || { + echo "Local environment file not found. Copy ${script_dir}/envs.sh.example to ${env_file}." >&2 + exit 1 + } + # shellcheck source=/dev/null + source "${env_file}" + + # All nodes in one distributed attempt must receive the same timestamp and + # base port. Give every resume attempt a new timestamp so its W&B ID differs. + timestamp=${1:-$(date +'%m%d_%H%M%S')} + base_port=${2:-29500} + pipeline_name=olmo3-1b + config_basename=${reproduce_dir}/cfgs/OLMo3-1B + extra_args=( + # For a short smoke run, uncomment both overrides. Do not use them for + # the full reproduction. + # "--trainer.max_duration.value=10" + # "--trainer.max_duration.unit=steps" + ) + + olmo3_data_root=${olmo3_data_root:?Set olmo3_data_root in envs.sh} + out_root=${out_root:?Set out_root in envs.sh} + tokenizer_json=${tokenizer_json:?Set tokenizer_json in envs.sh} + pipeline_root=${out_root}/runs/${pipeline_name} + + for stage_index in 1 2 3; do + stage="stage${stage_index}" + previous_save_folder= + stage_args=() + + case "${stage}" in + stage1) + config_file=${config_basename}-pretrain.py + stage_args=( + "--trainer.callbacks.downstream_evaluator.tokenizer.identifier=${tokenizer_json}" + ) + ;; + stage2) + config_file=${config_basename}-midtraining.py + stage_args=( + "--trainer.callbacks.downstream_evaluator.tokenizer.identifier=${tokenizer_json}" + ) + previous_save_folder=${pipeline_root}/stage1/checkpoints + ;; + stage3) + config_file=${config_basename}-long-context.py + previous_save_folder=${pipeline_root}/stage2/checkpoints + ;; + esac + + stage_port=$((base_port + stage_index)) + run_name=${pipeline_name}-${stage} + run_root=${pipeline_root}/${stage} + data_work_dir=${out_root}/dataset-cache/olmo3-${stage} + save_folder=${run_root}/checkpoints + success_marker=${run_root}/_SUCCESS + if [[ "${DRY_RUN:-0}" != "1" && -f "${success_marker}" ]]; then + echo "Skipping ${stage}; success marker already exists at '${success_marker}'" + continue + fi + + train_args=( + "--name=${run_name}" + "--data-root=${olmo3_data_root}" + "--save-folder=${save_folder}" + "--work-dir=${data_work_dir}" + "--trainer.work_dir=${run_root}/trainer" + ) + if [[ -n "${previous_save_folder}" ]]; then + # On a fresh stage this initializes model and optimizer state from the + # parent stage without loading parent trainer progress. If the current + # stage already has a checkpoint, OLMo-core resumes full local state. + train_args+=("--load_path=${previous_save_folder}") + fi + + if [[ "${ENABLE_WANDB:-1}" == "1" ]]; then + train_args+=( + "--trainer.callbacks.wandb.enabled=true" + "--trainer.callbacks.wandb.entity=${wandb_entity:?Set wandb_entity in envs.sh}" + "--trainer.callbacks.wandb.project=${wandb_project:?Set wandb_project in envs.sh}" + "--trainer.callbacks.wandb.group=${pipeline_name}" + "--trainer.callbacks.wandb.name=${run_name}_${timestamp}" # wandb.id = wandb.name + ) + if [[ "${WANDB_MODE:-online}" != "offline" ]]; then + : "${WANDB_API_KEY:?Export WANDB_API_KEY before launching for online W&B logging}" + export WANDB_API_KEY + else + # Remote cancel tags cannot be observed by an offline W&B run. + train_args+=("--trainer.callbacks.wandb.cancel_tags=null") + fi + fi + train_args+=( + "${stage_args[@]}" + "${extra_args[@]}" + ) + + echo "Starting ${stage} for pipeline '${pipeline_name}' on port ${stage_port}" + if [[ "${DRY_RUN:-0}" == "1" ]]; then + python "${config_file}" --dry-run "${train_args[@]}" + continue + fi + + if [[ "${NNODES:-1}" == "1" ]]; then + torchrun_args=(--standalone "--nproc-per-node=${NPROC_PER_NODE:-gpu}") + else + torchrun_args=( + "--nnodes=${NNODES}" + "--node-rank=${NODE_RANK}" + "--nproc-per-node=${NPROC_PER_NODE}" + "--master-addr=${MASTER_ADDR}" + "--master-port=${stage_port}" + ) + fi + mkdir -p "${run_root}" + torchrun "${torchrun_args[@]}" "${config_file}" -- "${train_args[@]}" + + [[ "${NODE_RANK:-0}" == "0" ]] && touch "${success_marker}" + done + + echo "Pipeline '${pipeline_name}' completed all three stages" + exit +} From 59faf0157102b53b84847e97ab1c9c529ed3f804 Mon Sep 17 00:00:00 2001 From: WangJie <2740469261@qq.com> Date: Mon, 17 Aug 2026 15:57:21 +0800 Subject: [PATCH 2/2] feat: olmo3 1b stage4-5 --- .gitignore | 2 +- .gitmodules | 4 + reproduce/olmo-core-backend/README.md | 327 -------------- reproduce/olmo-core-backend/README_zh.md | 313 ------------- .../cfgs/OLMo3-1B-long-context.py | 165 ------- .../cfgs/OLMo3-1B-midtraining.py | 56 --- .../cfgs/OLMo3-1B-pretrain.py | 26 -- reproduce/olmo-core-backend/requirements.txt | 4 - .../olmo-core-backend/run/envs.sh.example | 24 - reproduce/train-olmo-core/README.md | 418 ++++++++++++++++++ reproduce/train-olmo-core/README_zh.md | 399 +++++++++++++++++ .../cfgs/OLMo3-1B-long-context.py | 42 ++ .../cfgs/OLMo3-1B-midtraining.py | 65 +++ .../train-olmo-core/cfgs/OLMo3-1B-pretrain.py | 43 ++ .../train-olmo-core/cfgs/OLMo3-1B-sft.py | 123 ++++++ .../cfgs/_olmo3_1b_base.py} | 116 +++-- .../train-olmo-core/cfgs/_olmo3_1b_long.py | 192 ++++++++ reproduce/train-olmo-core/run/envs.sh.example | 25 ++ .../run/run.sh | 70 ++- .../train-olmo-core/third_party/OLMo-core | 1 + 20 files changed, 1436 insertions(+), 979 deletions(-) create mode 100644 .gitmodules delete mode 100644 reproduce/olmo-core-backend/README.md delete mode 100644 reproduce/olmo-core-backend/README_zh.md delete mode 100644 reproduce/olmo-core-backend/cfgs/OLMo3-1B-long-context.py delete mode 100644 reproduce/olmo-core-backend/cfgs/OLMo3-1B-midtraining.py delete mode 100644 reproduce/olmo-core-backend/cfgs/OLMo3-1B-pretrain.py delete mode 100644 reproduce/olmo-core-backend/requirements.txt delete mode 100755 reproduce/olmo-core-backend/run/envs.sh.example create mode 100644 reproduce/train-olmo-core/README.md create mode 100644 reproduce/train-olmo-core/README_zh.md create mode 100644 reproduce/train-olmo-core/cfgs/OLMo3-1B-long-context.py create mode 100644 reproduce/train-olmo-core/cfgs/OLMo3-1B-midtraining.py create mode 100644 reproduce/train-olmo-core/cfgs/OLMo3-1B-pretrain.py create mode 100644 reproduce/train-olmo-core/cfgs/OLMo3-1B-sft.py rename reproduce/{olmo-core-backend/cfgs/_olmo3_1b.py => train-olmo-core/cfgs/_olmo3_1b_base.py} (62%) create mode 100644 reproduce/train-olmo-core/cfgs/_olmo3_1b_long.py create mode 100644 reproduce/train-olmo-core/run/envs.sh.example rename reproduce/{olmo-core-backend => train-olmo-core}/run/run.sh (65%) create mode 160000 reproduce/train-olmo-core/third_party/OLMo-core diff --git a/.gitignore b/.gitignore index 694a92c..c8f682c 100644 --- a/.gitignore +++ b/.gitignore @@ -218,4 +218,4 @@ __marimo__/ .streamlit/secrets.toml # Machine-local OLMo reproduction settings (paths and optional credentials). -/reproduce/olmo-core-backend/run/envs.sh +/reproduce/train-olmo-core/run/envs.sh diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..9b264f6 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,4 @@ +[submodule "reproduce/train-olmo-core/third_party/OLMo-core"] + path = reproduce/train-olmo-core/third_party/OLMo-core + url = https://github.com/JT-Ushio/OLMo-core-muon-fix.git + branch = ready_for_archspace_base diff --git a/reproduce/olmo-core-backend/README.md b/reproduce/olmo-core-backend/README.md deleted file mode 100644 index 4ca9a30..0000000 --- a/reproduce/olmo-core-backend/README.md +++ /dev/null @@ -1,327 +0,0 @@ -# OLMo 3 1B Three-Stage Reproduction - -English | [中文](README_zh.md) - -This directory provides the training recipes and launcher for the three-stage OLMo 3 1B pipeline: - -1. stage 1: pretraining; -2. stage 2: midtraining; -3. stage 3: long-context extension. - -The model implementation, distributed trainer, checkpoint I/O, and dataset implementation all come from -OLMo-core. This directory contains only the configuration and entry point required for a specific -reproduction, keeping the experiment recipe separate from the general-purpose training framework. Changes -to recipes in this directory do not affect OLMo-core, while framework upgrades and fixes do not require -copying the entire framework into this repository. - -This reproduction must use the following custom OLMo-core source branch instead of the general PyPI release: - -- [https://github.com/JT-Ushio/OLMo-core-muon-fix/tree/ready_for_archspace_base](https://github.com/JT-Ushio/OLMo-core-muon-fix/tree/ready_for_archspace_base) - -This branch contains the Muon fixes required by the recipes. The OLMo 3 model is already implemented by -`TransformerConfig.olmo3_1B()`, so no additional modeling files are needed here. - -## 1. Directory structure - -```text -reproduce/olmo-core-backend/ -├── README.md -├── README_zh.md -├── requirements.txt -├── cfgs/ -│ ├── _olmo3_1b.py -│ ├── OLMo3-1B-pretrain.py -│ ├── OLMo3-1B-midtraining.py -│ └── OLMo3-1B-long-context.py -└── run/ - ├── envs.sh.example - └── run.sh -``` - -## 2. Recipe overview - - -| Stage | Data mix | Sequence length | Global batch (tokens) | Parallelism | Default Muon LR | -| --------- | ---------------------------------- | ----------------: | ----------------------: | ------------- | ----------------: | -| stage 1 | `OLMo_mix_0625_150Bsample` | 4,096 | 2,097,152 | HSDP | `1e-3` | -| stage 2 | `OLMo_midtraining_mix_0625_100B` | 4,096 | 2,097,152 | HSDP | `5e-4` | -| stage 3 | `OLMo_longmino_mix_0625` | 65,536 | 4,194,304 | HSDP | `5e-4` | - -All three stages use BF16, FlashAttention-3, and Muon by default, and each trains for one complete data -epoch. Stages 1 and 2 use fixed-length datasets. Stage 3 uses document packing, an intra-document attention -mask, and 8x YaRN RoPE scaling. All three stages use HSDP without context parallelism. -The default FlashAttention-3 configuration targets Hopper GPUs. Other supported GPUs should switch to -FlashAttention-2 as described in section 3.2. - -You can also select the SkipStep AdamW recipe with `adam`, but all three stages and every resume attempt in a -pipeline must use the same optimizer because later stages inherit the optimizer state from the preceding -stage. - -These are experimental configurations scaled from the official OLMo 3 7B recipes to the 1B model. They are -not officially released or tuned OLMo 3 1B recipes. - -## 3. Environment setup - -This project uses PyTorch 2.10 and CUDA 12.8, although any versions compatible with OLMo-core and the other -dependencies should work in principle. - -```bash -pip install --index-url https://download.pytorch.org/whl/cu128 \ - torch==2.10.0 torchvision torchaudio -``` - -### 3.1 Install the custom OLMo-core from source - -We recommend keeping a separate source checkout and installing it in editable mode: - -```bash -export OLMO_CORE_SRC=/path/to/OLMo-core-muon-fix -git clone --branch ready_for_archspace_base --single-branch \ - https://github.com/JT-Ushio/OLMo-core-muon-fix.git "${OLMO_CORE_SRC}" - -pip install -e "${OLMO_CORE_SRC}[all]" -``` - -### 3.2 Install attention kernels - -Install the latest attention kernels without pinning a tag, commit, or package version. First install -FlashAttention-2: - -```bash -MAX_JOBS=8 python -m pip install --upgrade --no-build-isolation flash-attn -``` - -Hopper GPUs such as H100 and H800 can use FlashAttention-3. Install it from the `hopper/` directory on the -default FlashAttention branch: - -```bash -export FLASH_ATTN_SRC=/path/to/flash-attention -git clone --depth 1 --recurse-submodules --shallow-submodules \ - https://github.com/Dao-AILab/flash-attention.git "${FLASH_ATTN_SRC}" - -cd "${FLASH_ATTN_SRC}/hopper" -FLASH_ATTENTION_DISABLE_FP16=TRUE \ -FLASH_ATTENTION_DISABLE_SM80=TRUE \ -MAX_JOBS=8 \ -python setup.py install -cd - -``` - -Other supported GPUs should use FlashAttention-2. All three upstream-synchronized cfgs select `flash_3` by -default. When using FA2, add the following override to the `extra_args` array in `run/run.sh` so it applies to -all three stages: - -```bash -"--model.attn_backend=flash_2" -``` - -`ring-flash-attn` remains available as an optional backend. Install it when a custom configuration enables -ring context parallelism; the current three-stage HSDP recipe without CP does not require it: - -```bash -pip install ring-flash-attn -``` - -Adjust `MAX_JOBS` to the CPU and memory available on the build node. Run the checks that correspond to the -backends you installed: - -```bash -python -m pip check -# FA2 (all installations) -python -c 'from olmo_core.nn.attention.flash_attn_api import has_flash_attn_2; assert has_flash_attn_2()' -# FA3 (Hopper only) -python -c 'from olmo_core.nn.attention.flash_attn_api import has_flash_attn_3; assert has_flash_attn_3()' -# ring-flash-attn (optional installation only) -python -c 'from olmo_core.nn.attention.flash_attn_api import has_ring_flash_attn; assert has_ring_flash_attn()' -python -c 'import dion, torch; print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name(0))' -``` - -## 4. Data preparation - -### 4.1 Data format - -The configurations directly use four `DataMix` manifests installed with the OLMo-core package. Every `.npy` -path listed by these manifests must be a one-dimensional token-ID binary array following the OLMo-core -convention and readable with `numpy.memmap`. Arbitrary text files, or files merely renamed to `.npy`, will not -work. The Dolma 2 tokenizer has a vocabulary size of 100,278, so these recipes infer the array dtype as -`uint32`. Documents must be correctly separated with the EOS token (ID `100257`), because stage 3 document -packing and intra-document masking depend on these boundaries. - -The data comes from the official OLMo release. This repository will provide a Hugging Face redistribution: -LINK TODO. - -### 4.2 Expected `olmo3_data_root` layout - -`run.sh` passes `olmo3_data_root` from `envs.sh` unchanged to all three configurations. OLMo-core then uses it -as the prefix for every relative path in the manifests. The approximate directory layout is shown below; -ellipses represent all sources and shards listed in the manifests: - -```text -olmo3_data_root/ -├── preprocessed/ -│ ├── dolma2-0625/v0.1-150b/ -│ │ └── allenai/dolma2-tokenizer/ -│ │ ├── finemath-3plus/part-000-00000.npy -│ │ └── ... -│ ├── dolma3-dolmino-official/100B/ -│ │ └── allenai/dolma3-tokenizer/ -│ │ ├── code-meta-reasoning/part-00-00000.npy -│ │ └── ... -│ └── dolma3_longmino_0625/ -│ └── allenai/dolma3-tokenizer/ -│ ├── 000000.npy -│ └── ... -└── eval-data/perplexity/ - └── v3_small_dolma2-tokenizer/ - ├── c4_en/val/part-0-00000.npy - ├── dolma_books/val/part-0-00000.npy - └── ... -``` - -Stage 1 uses the first tree, stage 2 the second, and stage 3 the third. The in-loop LM evaluations in stages 1 -and 2 also require the final validation tree. Every manifest filename must match exactly; providing only -similar top-level directories is insufficient. - -`tokenizer_json` is another required path. It must point to a Dolma 2 `tokenizer.json` readable by every node -and is used by the stage 1 and 2 in-loop downstream evaluator. It does not replace the tokenized training -arrays described above. - -## 5. Run training - -```bash -cd reproduce/olmo-core-backend -cp run/envs.sh.example run/envs.sh -# edit run/envs.sh -bash run/run.sh -``` - -`run/envs.sh` is excluded by `.gitignore`. - -### 5.1 W&B - -`envs.sh.example` sets `WANDB_MODE=offline` by default. This mode requires no API key, writes files under each -stage's `trainer/wandb/` directory, and automatically disables remote cancel tags, which only work online. - -The timestamp is used only in the W&B run name and ID. Keep `pipeline_name` unchanged when resuming the same -experiment, but use a new timestamp for each new job attempt to prevent the new W&B segment from overwriting -or mixing with the previous attempt. Every node in the same multi-node attempt must use the same timestamp. - -### 5.2 Output structure - -```text -out_root/ -├── dataset-cache/ -│ ├── olmo3-stage1/... -│ ├── olmo3-stage2/... -│ └── olmo3-stage3/... -└── runs/olmo3-1b/ - ├── stage1/ - │ ├── _SUCCESS - │ ├── checkpoints/ - │ │ └── step/ - │ │ ├── .metadata.json - │ │ ├── config.json - │ │ ├── data_paths.txt - │ │ ├── model_and_optim/ - │ │ │ ├── .metadata - │ │ │ └── ___.distcp - │ │ └── train/ - │ │ └── rank.pt - │ └── trainer/wandb/... - ├── stage2/ - │ └── ... - └── stage3/ - └── ... -``` - -`config.json` is the effective configuration, while `data_paths.txt` records the expanded data files that -were actually used. Preserve both together with the W&B records when archiving a reproduction run. The -configurations write a temporary checkpoint approximately every 1 billion tokens, retain only one temporary -checkpoint, and save a final checkpoint at the end of each stage. - -Node rank 0 creates `_SUCCESS` after `torchrun` exits successfully for that stage. It indicates successful -process completion; it does not revalidate the checkpoint step or metric values. - -### 5.3 Resume and stage transitions - -A normal resume does not require specifying a checkpoint manually: - -```bash -# Keep out_root and pipeline_name unchanged; use a new attempt timestamp. -bash run/run.sh 0809_093000 -``` - -The launcher and OLMo-core behave as follows: - -1. If `stageN/_SUCCESS` exists, that stage is skipped. -2. If `_SUCCESS` does not exist but the current stage has a checkpoint under `checkpoints/`, the model, - optimizer, trainer, data-loader, and RNG states are restored from it. -3. If the current stage 2 or 3 has no checkpoint, the top-level `--load_path` initializes the model and - optimizer from the preceding stage's checkpoint without inheriting that stage's step or epoch progress. -4. If the current stage 1 has no checkpoint, training starts from scratch. - -Therefore, rerunning after a stage 2 interruption skips the completed stage 1 and continues from stage 2's -own latest checkpoint. Stage 3 starts only after stage 2 finishes. - -## 6. Evaluate with OLMES - -Training produces OLMo-core distributed checkpoints, while OLMES expects a Hugging Face model directory for -a local model. Convert the checkpoint first, then run OLMES. Normally, you evaluate the final stage 3 -checkpoint. To compare stages, convert stages 1, 2, and 3 separately. - -### 6.1 Convert to Hugging Face format - -Select a specific `step` directory, not its parent `checkpoints/` directory: - -```bash -export CHECKPOINT=/path/to/out_root/runs/olmo3-1b/stage3/checkpoints/step11921 -export HF_MODEL_DIR=/path/to/out_root/hf/olmo3-1b-stage3-step11921 - -python "${OLMO_CORE_SRC}/src/examples/huggingface/convert_checkpoint_to_hf.py" \ - --checkpoint-input-path "${CHECKPOINT}" \ - --huggingface-output-dir "${HF_MODEL_DIR}" \ - --max-sequence-length 65536 -``` - -The converter reconstructs the OLMo 3 architecture from the checkpoint's `config.json` and uses -`allenai/dolma2-tokenizer` from the configuration by default. In an offline environment, additionally pass -`--tokenizer /path/to/local/hf-tokenizer-directory`. This must be a complete directory loadable by -`AutoTokenizer.from_pretrained()`, not an individual `tokenizer.json` file. - -Numerical validation is enabled by default. Avoid `--skip-validation` unless you have validated the result -separately and explicitly accept the risk. After conversion, run a minimal loading test: - -```bash -python -c 'import os; from transformers import AutoModelForCausalLM, AutoTokenizer; p=os.environ["HF_MODEL_DIR"]; AutoTokenizer.from_pretrained(p); AutoModelForCausalLM.from_pretrained(p); print("HF checkpoint OK")' -``` - -### 6.2 Install and run OLMES - -Use a separate evaluation environment so that the vLLM and Transformers versions do not affect the training -environment: - -```bash -git clone https://github.com/allenai/olmes.git /path/to/olmes -cd /path/to/olmes -python -m pip install -e '.[gpu]' -git rev-parse HEAD -``` - -For a small-scale experiment, start with the OLMo 3 base-easy suites: - -```bash -olmes \ - --model "${HF_MODEL_DIR}" \ - --task \ - olmo3:base_easy:code_bpb \ - olmo3:base_easy:math_bpb \ - olmo3:base_easy:qa_rc \ - olmo3:base_easy:qa_bpb \ - --output-dir /path/to/out_root/evals/olmo3-1b-stage3-base-easy -``` - -If the installed OLMES/vLLM versions support this model, add `--model-type vllm` for higher throughput. A -formal report should preserve the OLMES commit, complete command, task suite, checkpoint step, Hugging Face -conversion arguments, and output directory. The `FAST_TASKS` and PPL in-loop evaluations built into the -training configuration are intended for training monitoring and do not replace a final, version-pinned OLMES -evaluation. diff --git a/reproduce/olmo-core-backend/README_zh.md b/reproduce/olmo-core-backend/README_zh.md deleted file mode 100644 index bfe728d..0000000 --- a/reproduce/olmo-core-backend/README_zh.md +++ /dev/null @@ -1,313 +0,0 @@ -# OLMo 3 1B 三阶段复现 - -[English](README.md) | 中文 - -本目录提供 OLMo 3 1B 的三阶段训练配方与启动脚本： - -1. stage 1：pretraining； -2. stage 2：midtraining； -3. stage 3：long-context extension。 - -模型实现、分布式训练器、checkpoint I/O 和数据集实现均来自 OLMo-core。本目录只保存某次 -复现所需的配置和运行入口，以便把“实验配方”与“通用训练框架”隔离开：修改本目录中的配方 -不会污染 OLMo-core，升级或修复训练框架也不需要把整个框架复制进本仓库。 - -本复现必须使用以下自定义 OLMo-core 源码分支，而不是 PyPI 上的通用版本： - -- [https://github.com/JT-Ushio/OLMo-core-muon-fix/tree/ready_for_archspace_base](https://github.com/JT-Ushio/OLMo-core-muon-fix/tree/ready_for_archspace_base) - -该分支包含本配方所依赖的 Muon 修复。OLMo 3 模型本身已由 -`TransformerConfig.olmo3_1B()` 实现，因此这里没有额外的 modeling 文件。 - -## 1. 目录结构 - -```text -reproduce/olmo-core-backend/ -├── README.md -├── README_zh.md -├── requirements.txt -├── cfgs/ -│ ├── _olmo3_1b.py -│ ├── OLMo3-1B-pretrain.py -│ ├── OLMo3-1B-midtraining.py -│ └── OLMo3-1B-long-context.py -└── run/ - ├── envs.sh.example - └── run.sh -``` - -## 2. 配方概览 - - -| 阶段 | 数据 mix | 序列长度 | 全局 batch（token） | 并行方式 | 默认 Muon LR | -| --------- | ---------------------------------- | ---------: | --------------------: | ---------- | -------------: | -| stage 1 | `OLMo_mix_0625_150Bsample` | 4,096 | 2,097,152 | HSDP | `1e-3` | -| stage 2 | `OLMo_midtraining_mix_0625_100B` | 4,096 | 2,097,152 | HSDP | `5e-4` | -| stage 3 | `OLMo_longmino_mix_0625` | 65,536 | 4,194,304 | HSDP | `5e-4` | - -三个阶段默认都使用 BF16、FlashAttention-3 和 Muon，并各自训练一个完整数据 -epoch。stage 1/2 使用固定长度数据集；stage 3 使用 document packing、文档内 attention -mask 和 8 倍 YaRN RoPE scaling。三个阶段均使用 HSDP，不使用 context parallelism。 -默认的 FlashAttention-3 配置面向 Hopper GPU；其他支持的 GPU 应按 3.2 节切换到 FlashAttention-2。 - -也可以用 `adam` 选择 SkipStep AdamW 配方，但同一 pipeline 的三个阶段及所有 resume 必须使用 -同一种 optimizer，因为后续阶段会继承前一阶段的 optimizer state。 - -这些是从 OLMo 3 7B 官方配方缩放到 1B 模型的实验配置，并不是官方发布、已调优的 OLMo 3 -1B recipe。 - -## 3. 环境安装 - -本项目采用 PyTorch 2.10 和 CUDA 12.8，但原则上也可使用任何与 OLMo-core 及其他依赖兼容的版本。 - -```bash -pip install --index-url https://download.pytorch.org/whl/cu128 \ - torch==2.10.0 torchvision torchaudio -``` - -### 3.1 从源码安装自定义 OLMo-core - -推荐保留一个独立源码 checkout，并以 editable 方式安装： - -```bash -export OLMO_CORE_SRC=/path/to/OLMo-core-muon-fix -git clone --branch ready_for_archspace_base --single-branch \ - https://github.com/JT-Ushio/OLMo-core-muon-fix.git "${OLMO_CORE_SRC}" - -pip install -e "${OLMO_CORE_SRC}[all]" -``` - -### 3.2 安装 attention kernels - -本节的 attention kernels 都直接安装最新版，不固定 tag、commit 或 package -version。先安装 FlashAttention-2： - -```bash -MAX_JOBS=8 python -m pip install --upgrade --no-build-isolation flash-attn -``` - -Hopper GPU（例如 H100/H800）可以使用 FlashAttention-3，从 FlashAttention 默认分支的 -`hopper/` 目录安装： - -```bash -export FLASH_ATTN_SRC=/path/to/flash-attention -git clone --depth 1 --recurse-submodules --shallow-submodules \ - https://github.com/Dao-AILab/flash-attention.git "${FLASH_ATTN_SRC}" - -cd "${FLASH_ATTN_SRC}/hopper" -FLASH_ATTENTION_DISABLE_FP16=TRUE \ -FLASH_ATTENTION_DISABLE_SM80=TRUE \ -MAX_JOBS=8 \ -python setup.py install -cd - -``` - -其他支持的 GPU 使用 FlashAttention-2。三个 cfg 从上游同步的默认 backend 都是 -`flash_3`；使用 FA2 时，在 `run/run.sh` 的 `extra_args` 数组中加入以下覆盖，使它同时 -作用于三个阶段： - -```bash -"--model.attn_backend=flash_2" -``` - -`ring-flash-attn` 作为可选 backend 保留。当自定义配置启用 ring context parallelism 时 -再安装；当前 HSDP、无 CP 的三阶段配方不需要它： - -```bash -pip install ring-flash-attn -``` - -`MAX_JOBS` 应按编译节点的 CPU 和内存调整。可按实际选择的 backend 分别验证： - -```bash -python -m pip check -# FA2（所有安装） -python -c 'from olmo_core.nn.attention.flash_attn_api import has_flash_attn_2; assert has_flash_attn_2()' -# FA3（仅 Hopper） -python -c 'from olmo_core.nn.attention.flash_attn_api import has_flash_attn_3; assert has_flash_attn_3()' -# ring-flash-attn（仅可选安装） -python -c 'from olmo_core.nn.attention.flash_attn_api import has_ring_flash_attn; assert has_ring_flash_attn()' -python -c 'import dion, torch; print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name(0))' -``` - -## 4. 数据准备 - -### 4.1 数据格式 - -配置直接使用安装在 OLMo-core 包中的四份 `DataMix` manifest。它们列出的每个 `.npy` 路径是 -OLMo-core 约定的、可由 `numpy.memmap` 读取的一维 token-ID 二进制数组，而不是任意文本文件， -也不能只靠把文件改名为 `.npy` 得到。Dolma 2 tokenizer 的词表大小为 100,278，因此本配方会 -推断数组 dtype 为 `uint32`。不同文档需要以 EOS token（ID `100257`）正确分隔，stage 3 的 -document packing 和文档内 mask 依赖这些边界。 - -数据来自olmo官方，本仓库提供 Huggingface 再发布版本：链接TODO - -### 4.2 `olmo3_data_root` 的预期布局 - -`run.sh` 把 `envs.sh` 中的 `olmo3_data_root` 原样传给三个配置。OLMo-core 再把它作为 manifest -内所有相对路径的前缀。大致目录如下；省略号代表 manifest 中的全部 source 和 shard： - -```text -olmo3_data_root/ -├── preprocessed/ -│ ├── dolma2-0625/v0.1-150b/ -│ │ └── allenai/dolma2-tokenizer/ -│ │ ├── finemath-3plus/part-000-00000.npy -│ │ └── ... -│ ├── dolma3-dolmino-official/100B/ -│ │ └── allenai/dolma3-tokenizer/ -│ │ ├── code-meta-reasoning/part-00-00000.npy -│ │ └── ... -│ └── dolma3_longmino_0625/ -│ └── allenai/dolma3-tokenizer/ -│ ├── 000000.npy -│ └── ... -└── eval-data/perplexity/ - └── v3_small_dolma2-tokenizer/ - ├── c4_en/val/part-0-00000.npy - ├── dolma_books/val/part-0-00000.npy - └── ... -``` - -stage 1 使用第一棵树，stage 2 使用第二棵树，stage 3 使用第三棵树；stage 1/2 的 in-loop LM -evaluation 都需要最后一棵 validation 树。manifest 文件名必须逐项匹配，不能只提供相似的 -顶层目录。 - -`tokenizer_json` 是另一项必填路径：它应指向所有节点都能读取的 Dolma 2 `tokenizer.json`， -供 stage 1/2 的 in-loop downstream evaluator 使用。它不替代上述已分词训练数组。 - -## 5. 运行训练 - -```bash -cd reproduce/olmo-core-backend -cp run/envs.sh.example run/envs.sh -# edit run/envs.sh -bash run/run.sh -``` - -`run/envs.sh` 已被 `.gitignore` 排除。 - -### 5.1 W&B - -`envs.sh.example` 默认 `WANDB_MODE=offline`。这种模式不需要 API key，文件写到每个阶段的 -`trainer/wandb/` 下，并自动禁用只能在线工作的 remote cancel tags。 - -timestamp 只参与 W&B run name/ID。resume 同一个 -实验时保持 `pipeline_name` 不变，但每次重新发起任务建议给一个新 timestamp，避免新的 W&B -片段覆盖或混入上一次 attempt。多机同一次 attempt 必须使用同一个 timestamp。 - -### 5.2 输出目录 - -```text -out_root/ -├── dataset-cache/ -│ ├── olmo3-stage1/... -│ ├── olmo3-stage2/... -│ └── olmo3-stage3/... -└── runs/olmo3-1b/ - ├── stage1/ - │ ├── _SUCCESS - │ ├── checkpoints/ - │ │ └── step/ - │ │ ├── .metadata.json - │ │ ├── config.json - │ │ ├── data_paths.txt - │ │ ├── model_and_optim/ - │ │ │ ├── .metadata - │ │ │ └── ___.distcp - │ │ └── train/ - │ │ └── rank.pt - │ └── trainer/wandb/... - ├── stage2/ - │ └── ... - └── stage3/ - └── ... -``` - -`config.json` 是最终生效配置，`data_paths.txt` 记录实际展开的数据文件；复现实验归档时应和 W&B -记录一起保存。配置约每 10 亿 token 写一次临时 checkpoint，只保留一个临时版本，并在阶段 -结束时保存最终 checkpoint。 - -`_SUCCESS` 由 node rank 0 在该阶段 `torchrun` 成功退出之后创建。它表示进程成功完成，不会 -再次检查 checkpoint step 或指标内容。 - -### 5.3 Resume 与阶段衔接 - -正常 resume 不需要手工指定 checkpoint： - -```bash -# out_root、pipeline_name 保持不变；使用新的 attempt timestamp。 -bash run/run.sh 0809_093000 -``` - -启动器和 OLMo-core 的行为是： - -1. 存在 `stageN/_SUCCESS`：直接跳过该阶段； -2. 不存在 `_SUCCESS`，但当前阶段 `checkpoints/` 中有 checkpoint：恢复该阶段的 model、 - optimizer、trainer、data-loader 和 RNG 状态； -3. 当前 stage 2/3 没有 checkpoint：通过顶层 `--load_path` 从前一阶段 checkpoint 初始化 model - 和 optimizer，但不继承前一阶段的 step/epoch 进度； -4. 当前 stage 1 没有 checkpoint：从头开始。 - -因此 stage 2 中断后的重跑会跳过已完成的 stage 1，并从 stage 2 自己的最新 checkpoint 继续； -stage 2 完成后才进入 stage 3。 - -## 6. 使用 OLMES 评测 - -训练输出是 OLMo-core distributed checkpoint，而 OLMES 的本地模型入口使用 Hugging Face -模型目录。因此先转换 checkpoint，再运行 OLMES。通常评测 stage 3 的最终 checkpoint；若要画 -阶段对比，则分别转换 stage 1/2/3。 - -### 6.1 转换为 Hugging Face 格式 - -选中具体的 `step` 目录，而不是它的 `checkpoints/` 父目录： - -```bash -export CHECKPOINT=/path/to/out_root/runs/olmo3-1b/stage3/checkpoints/step11921 -export HF_MODEL_DIR=/path/to/out_root/hf/olmo3-1b-stage3-step11921 - -python "${OLMO_CORE_SRC}/src/examples/huggingface/convert_checkpoint_to_hf.py" \ - --checkpoint-input-path "${CHECKPOINT}" \ - --huggingface-output-dir "${HF_MODEL_DIR}" \ - --max-sequence-length 65536 -``` - -转换器会从 checkpoint 的 `config.json` 恢复 OLMo 3 架构，并默认使用配置中的 -`allenai/dolma2-tokenizer`。离线环境可额外传 -`--tokenizer /path/to/local/hf-tokenizer-directory`；这里应给一个可由 -`AutoTokenizer.from_pretrained()` 加载的完整目录，而不是单独的 `tokenizer.json`。 - -默认转换包含数值验证。除非已经单独验证且明确接受风险，不建议使用 `--skip-validation`。 -转换后可先做最小加载测试： - -```bash -python -c 'import os; from transformers import AutoModelForCausalLM, AutoTokenizer; p=os.environ["HF_MODEL_DIR"]; AutoTokenizer.from_pretrained(p); AutoModelForCausalLM.from_pretrained(p); print("HF checkpoint OK")' -``` - -### 6.2 安装并运行 OLMES - -建议为评测创建独立环境，避免 vLLM/Transformers 版本反向影响训练环境： - -```bash -git clone https://github.com/allenai/olmes.git /path/to/olmes -cd /path/to/olmes -python -m pip install -e '.[gpu]' -git rev-parse HEAD -``` - -小规模实验可从 OLMo 3 base-easy suites 开始： - -```bash -olmes \ - --model "${HF_MODEL_DIR}" \ - --task \ - olmo3:base_easy:code_bpb \ - olmo3:base_easy:math_bpb \ - olmo3:base_easy:qa_rc \ - olmo3:base_easy:qa_bpb \ - --output-dir /path/to/out_root/evals/olmo3-1b-stage3-base-easy -``` - -如 OLMES/vLLM 版本支持该模型，可加 `--model-type vllm` 提高吞吐。正式报告中应保留 OLMES -commit、完整命令、task suite、checkpoint step、HF 转换参数和输出目录。训练配置内置的 -`FAST_TASKS`/PPL in-loop evaluation 用于训练监控，不能替代最终、版本固定的 OLMES 评测。 diff --git a/reproduce/olmo-core-backend/cfgs/OLMo3-1B-long-context.py b/reproduce/olmo-core-backend/cfgs/OLMo3-1B-long-context.py deleted file mode 100644 index 1532afb..0000000 --- a/reproduce/olmo-core-backend/cfgs/OLMo3-1B-long-context.py +++ /dev/null @@ -1,165 +0,0 @@ -""" -OLMo 3 1B stage-3 long-context extension configuration. - -This is a 1B adaptation of the OLMo 3 7B long-context recipe in -`src/scripts/official/OLMo3/OLMo-3-1025-7B-long-context.py`. OLMo 3 does not publish an -officially tuned 1B long-context recipe. - -Parallelism boundaries ----------------------- -DP=data-parallel world size; PP/CP/TP/EP are their degrees. -H_rep=HSDP replicas, H_shard=HSDP shard degree, L=seqlen, M=microbatch tokens, -B=global batch tokens; heads=16; n_layers=16. - -Mesh: world_size = PP*CP*TP*DP; world_size % (PP*CP*TP) = 0 -HSDP: DP = H_rep*H_shard; DP % H_shard = 0 -Batch: M % L = 0; B % (M*DP) = 0; grad_accum = B/(M*DP) -CP: local_L = L/CP; exact split requires L % CP = 0 -Ulysses CP: q_heads % CP = kv_heads % CP = 0 -TP: tensor_dim % TP = 0 for every sharded dimension -PP: world_size % PP = 0; num_stages % PP = 0; num_stages <= n_layers -EP: MoE and HSDP only; EP = H_shard; TP = 1 (off) - -Optimizer / parallelism matrix: -| Mode | AdamW | Muon | -|--------------|-----------------|--------------------------------------| -| FSDP | yes | yes: heads % (DP*CP) = 0 | -| HSDP, CP off | yes | yes: heads % H_shard = 0 | -| HSDP + CP | yes | no: dp_shard is flattened into dp_cp | -| TP | yes | no: hard error | -| PP | yes (beta) | beta; changes DP | -| EP | MoE + HSDP only | no: flattened/3D expert parameters | - -Other conflicts: flash_3 has no CP, use flash_2 for CP; -TP + EP is forbidden. Multi-stage PP + tied embeddings is forbidden. -Stage 2/3 optimizer states must have the same optimizer type unless loading is disabled. - -Examples: world_size=64 (GPUs), PP=1 (off), TP=1 (off), heads=16 -B=2^22 (tokens), M=L=65,536 (tokens) -| Optim | DP layout | CP | Muon mesh | Result | -|-------|-------------------------|----|-----------|-------------------------------| -| AdamW | HSDP H_rep=16,H_shard=1 | 4 | - | valid | -| AdamW | HSDP H_rep=8,H_shard=1 | 8 | - | valid | -| Muon | HSDP H_rep=8,H_shard=8 | 1 | 8 | valid:16(heads)%8(mesh)=0 | -| Muon | FSDP DP=64 | 1 | 64 | invalid:16(heads)%64(mesh)!=0 | -| Muon | FSDP DP=16 | 4 | DP*CP=64 | invalid:16(heads)%64(mesh)!=0 | -""" - -import argparse -from typing import List - -from _olmo3_1b import build_common_config, build_optim_config, get_olmo3_1b_cli_parser - -from olmo_core.config import DType -from olmo_core.data import ( - DataMix, - NumpyDataLoaderConfig, - NumpyPackedFSLDatasetConfig, - TokenizerConfig, -) -from olmo_core.distributed.parallel import DataParallelType -from olmo_core.nn.attention import AttentionBackendName -from olmo_core.nn.rope import YaRNRoPEScalingConfig -from olmo_core.nn.transformer import TransformerConfig -from olmo_core.optim import LinearWithWarmup -from olmo_core.script_utils import ExperimentConfig, main -from olmo_core.train.common import LoadStrategy -from olmo_core.train.train_module import ( - TransformerContextParallelConfig, # noqa: F401 - used by the optional cp_config below - TransformerDataParallelConfig, - TransformerDataParallelWrappingStrategy, - TransformerTrainModuleConfig, -) - -DEFAULT_SEQUENCE_LENGTH = 65536 -GLOBAL_BATCH_SIZE = 2**22 # 4M tokens -# MAX_TOKENS = 50_000_000_000 # 50B -# Muon retains the 1B recipe; AdamW follows the official stage-3 schedule. -MUON_LR = 5e-4 -ADAM_LR = 5e-4 -SEED = 4123 - - -def build_config(opts: argparse.Namespace, overrides: List[str]) -> ExperimentConfig: - """Build stage 3 from its required components and the shared trainer.""" - # Long context changes the model, dataset, loader, and train module as whole - # units, so this stage does not mutate the stage-1 versions of those components. - sequence_length = opts.sequence_length or DEFAULT_SEQUENCE_LENGTH - tokenizer_config = TokenizerConfig.dolma2() - - model = TransformerConfig.olmo3_1B( - vocab_size=tokenizer_config.padded_vocab_size(), # pad to a multiple of 128 - attn_backend=AttentionBackendName.flash_3, - ).with_rope_scaling( - YaRNRoPEScalingConfig( - factor=8, - beta_fast=32, - beta_slow=1, - old_context_len=8192, - ) - ) - - dataset = NumpyPackedFSLDatasetConfig.from_data_mix( - DataMix.OLMo_longmino_mix_0625, - mix_base_dir=opts.data_root, - work_dir=opts.work_dir, - tokenizer=tokenizer_config, - sequence_length=sequence_length, - generate_doc_lengths=True, # enables intra-document masking - source_group_size=8, - source_permutation_seed=123, - ) - - data_loader = NumpyDataLoaderConfig( - global_batch_size=GLOBAL_BATCH_SIZE, - seed=SEED, - num_workers=8, - prefetch_factor=4, - ) - - train_module = TransformerTrainModuleConfig( - rank_microbatch_size=sequence_length, - max_sequence_length=sequence_length, - optim=build_optim_config( - opts.optim, - muon_lr=MUON_LR, - adam_lr=ADAM_LR, - ), - scheduler=LinearWithWarmup(warmup=200, alpha_f=0.0), - compile_model=True, - dp_config=TransformerDataParallelConfig( - name=DataParallelType.hsdp, - param_dtype=DType.bfloat16, - reduce_dtype=DType.float32, - wrapping_strategy=TransformerDataParallelWrappingStrategy.full, - ), - # cp_config=TransformerContextParallelConfig.llama3(degree=4, head_stride=4), - ac_config=None, - float8_config=None, - # float8_config=Float8Config(enabled=True, ao=AOFloat8LinearConfig.recommended()), - z_loss_multiplier=1e-5, - max_grad_norm=1.0, - ) - - # Only the trainer and its common callbacks are inherited from stage 1. - config = build_common_config( - opts, - model=model, - dataset=dataset, - data_loader=data_loader, - train_module=train_module, - ) - - config.trainer.load_strategy = LoadStrategy.always - # script_utils.main probes save_folder before Trainer.fit() and uses this value, so - # require trainer state for a same-stage resume. The launcher supplies the parent - # stage through ExperimentConfig.load_path, which explicitly skips trainer state. - config.trainer.load_trainer_state = True - config.trainer.load_optim_state = True - - config.init_seed = SEED - return config.merge(overrides) - - -if __name__ == "__main__": - main(build_config, parser=get_olmo3_1b_cli_parser()) diff --git a/reproduce/olmo-core-backend/cfgs/OLMo3-1B-midtraining.py b/reproduce/olmo-core-backend/cfgs/OLMo3-1B-midtraining.py deleted file mode 100644 index 6e30a76..0000000 --- a/reproduce/olmo-core-backend/cfgs/OLMo3-1B-midtraining.py +++ /dev/null @@ -1,56 +0,0 @@ -""" -OLMo 3 1B stage-2 midtraining configuration. - -This is a 1B adaptation of the official OLMo-3-1025-7B midtraining recipe in -`src/scripts/official/OLMo3/OLMo-3-1025-7B-midtrain.py`. OLMo 3 does not -publish an officially tuned 1B midtraining recipe, so the data schedule and -optimization settings below intentionally retain the official 7B values. -""" - -import argparse -from typing import List - -from _olmo3_1b import build_optim_config, build_pretrain_config, get_olmo3_1b_cli_parser - -from olmo_core.data import DataMix -from olmo_core.optim import LinearWithWarmup -from olmo_core.script_utils import ExperimentConfig, main -from olmo_core.train.common import LoadStrategy - -# MAX_TOKENS = 100_000_000_000 # 100B -# Muon retains the 1B recipe; AdamW follows the official stage-2 schedule. -MUON_LR = 5e-4 -ADAM_LR = 5e-4 -SEED = 1337 - - -def build_config(opts: argparse.Namespace, overrides: List[str]) -> ExperimentConfig: - """Build stage 2 by applying its differences to the pretraining configuration.""" - config = build_pretrain_config(opts) - - # Model shape, including the padded vocabulary size, batching, and callbacks - # remain identical to stage 1. Only data order and optimization are stage-specific. - config.dataset.mix = DataMix.OLMo_midtraining_mix_0625_100B - config.data_loader.seed = SEED - - # Optimizer state is restored from stage 1, so the selected recipe must match. - config.train_module.optim = build_optim_config( - opts.optim, - muon_lr=MUON_LR, - adam_lr=ADAM_LR, - ) - config.train_module.scheduler = LinearWithWarmup(warmup=0, alpha_f=0.0) - - config.trainer.load_strategy = LoadStrategy.always - # script_utils.main probes save_folder before Trainer.fit() and uses this value, so - # require trainer state for a same-stage resume. The launcher supplies the parent - # stage through ExperimentConfig.load_path, which explicitly skips trainer state. - config.trainer.load_trainer_state = True - config.trainer.load_optim_state = True - - config.init_seed = SEED - return config.merge(overrides) - - -if __name__ == "__main__": - main(build_config, parser=get_olmo3_1b_cli_parser()) diff --git a/reproduce/olmo-core-backend/cfgs/OLMo3-1B-pretrain.py b/reproduce/olmo-core-backend/cfgs/OLMo3-1B-pretrain.py deleted file mode 100644 index a9e993e..0000000 --- a/reproduce/olmo-core-backend/cfgs/OLMo3-1B-pretrain.py +++ /dev/null @@ -1,26 +0,0 @@ -""" -OLMo 3 1B stage-1 pretraining configuration for the local 150B data sample. - -This is a 1B adaptation of the official OLMo-3-1025-7B stage-1 recipe in -``src/scripts/official/OLMo3/OLMo-3-1025-7B-pretrain-1.py``. OLMo 3 does not -publish an official tuned 1B pretraining recipe, so the batch size, learning -rate, and warmup below intentionally retain the official 7B values. -""" - -import argparse -from typing import List - -from _olmo3_1b import build_pretrain_config, get_olmo3_1b_cli_parser -from olmo_core.script_utils import ExperimentConfig, main - - -def build_config(opts: argparse.Namespace, overrides: List[str]) -> ExperimentConfig: - """Build the OLMo 3 1B stage-1 pretraining configuration.""" - # This complete stage-1 recipe, including the padded vocabulary size, is also - # the baseline imported by stage 2. - # Merge CLI overrides only after the shared defaults have been assembled. - return build_pretrain_config(opts).merge(overrides) - - -if __name__ == "__main__": - main(build_config, parser=get_olmo3_1b_cli_parser()) diff --git a/reproduce/olmo-core-backend/requirements.txt b/reproduce/olmo-core-backend/requirements.txt deleted file mode 100644 index f9bf4d3..0000000 --- a/reproduce/olmo-core-backend/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ -# Install OLMo-core from the source branch that contains the Muon fixes used by -# this reproduction. Hardware-specific FlashAttention must be installed -# separately; see README.md. -ai2-olmo-core[all] @ git+https://github.com/JT-Ushio/OLMo-core-muon-fix.git@ready_for_archspace_base diff --git a/reproduce/olmo-core-backend/run/envs.sh.example b/reproduce/olmo-core-backend/run/envs.sh.example deleted file mode 100755 index 0bb8d15..0000000 --- a/reproduce/olmo-core-backend/run/envs.sh.example +++ /dev/null @@ -1,24 +0,0 @@ -#!/usr/bin/env bash - -# Copy this file to envs.sh and replace the placeholders with paths available -# on every node. envs.sh is ignored by Git because it is machine-specific. - -# Root containing every relative path referenced by OLMo-core's built-in OLMo 3 -# stage-1, stage-2, stage-3, and perplexity-validation data-mix manifests. -olmo3_data_root=/path/to/olmo3-data - -# Checkpoints, trainer artifacts, W&B files, and dataset caches are written here. -# Use shared, persistent storage for multi-node training and resume. -out_root=/path/to/training-output - -# Local Dolma 2 tokenizer JSON used by the in-loop downstream evaluator. The LM -# dataset tokenizer metadata still comes from TokenizerConfig.dolma2(). -tokenizer_json=/path/to/tokenizer.json - -# W&B destination. Offline mode is the safe default and does not need an API key. -wandb_entity=YOUR_ENTITY -wandb_project=YOUR_PROJECT -export WANDB_MODE=${WANDB_MODE:-offline} - -# For online logging, export the secret in the calling shell; do not put it here. -# export WANDB_API_KEY=YOUR_SECRET diff --git a/reproduce/train-olmo-core/README.md b/reproduce/train-olmo-core/README.md new file mode 100644 index 0000000..4ffa22a --- /dev/null +++ b/reproduce/train-olmo-core/README.md @@ -0,0 +1,418 @@ +# OLMo 3 1B Five-Stage Reproduction + +English | [中文](README_zh.md) + +This directory provides the training recipes and launcher for the five-stage OLMo 3 1B pipeline: + +1. stage 1: pretraining; +2. stage 2: midtraining; +3. stage 3: long-context extension; +4. stage 4: Think SFT; +5. stage 5: Instruct SFT. + +OLMo-core provides the model definitions, distributed training, checkpoint I/O, and data loading. The Muon +fixes required by these recipes are pinned in the `third_party/OLMo-core` submodule, and the model +configuration uses `TransformerConfig.olmo3_1B()`. + +## 1. Directory structure + +```text +reproduce/train-olmo-core/ +├── README.md +├── README_zh.md +├── cfgs/ +│ ├── _olmo3_1b_base.py +│ ├── _olmo3_1b_long.py +│ ├── OLMo3-1B-pretrain.py +│ ├── OLMo3-1B-midtraining.py +│ ├── OLMo3-1B-long-context.py +│ └── OLMo3-1B-sft.py +├── run/ +│ ├── envs.sh.example +│ └── run.sh +└── third_party/ + └── OLMo-core/ # Editable Git submodule +``` + +## 2. Recipe overview + +| Stage | Data | Sequence length | Rank microbatch (tokens) | Global batch (tokens) | Epochs | Default Muon LR | +| ------- | ------------------------------------- | --------------: | -----------------------: | --------------------: | -----: | --------------: | +| stage 1 | `OLMo_mix_0625_150Bsample` | 4,096 | 16,384 | 2,097,152 | 1 | `5e-3` | +| stage 2 | `OLMo_midtraining_mix_0625_100B` | 4,096 | 16,384 | 1,048,576 | 1 | `2.071235285e-4` | +| stage 3 | `OLMo_longmino_mix_0625` | 32,768 | 32,768 | 2,097,152 | 1 | `2.071235285e-4` | +| stage 4 | `Dolci-Think-SFT-7B` paired NPY files | 32,768 | 32,768 | 1,048,576 | 2 | `5e-5` | +| stage 5 | `Dolci-Instruct-SFT` paired NPY files | 32,768 | 32,768 | 1,048,576 | 2 | `8e-5` | + +All five stages use BF16, FlashAttention-3, and Muon by default. Stages 1 and 2 use fixed-length datasets. +Stage 3 uses Longmino document packing, an intra-document attention mask, and 8x YaRN RoPE scaling from an +old context length of 4,096. Stages 4 and 5 inherit that long-context model and use packed, assistant-masked +SFT data. All five stages use HSDP with context parallelism disabled. +The default FlashAttention-3 configuration targets Hopper GPUs. Other supported GPUs should switch to +FlashAttention-2 as described in section 3.3. + +To select the SkipStep AdamW recipe, add `"--optim=adam"` to `all_stage_args` in `run/run.sh`. All five stages +and every resume attempt in a pipeline must use the same optimizer because later stages inherit the optimizer +state from the preceding stage. + +These experimental configurations adapt the official OLMo 3 7B recipes to the 1B model; they have not been +officially released or tuned as OLMo 3 1B recipes. + +## 3. Environment setup with uv + +This workflow targets Python 3.12, PyTorch 2.10.0, and CUDA 12.8. Install +[uv](https://docs.astral.sh/uv/getting-started/installation/) before following the steps below; uv creates +the virtual environments and installs their packages. + +For a new ArchSpace checkout, initialize OLMo-core as part of the clone: + +```bash +git clone --recurse-submodules --branch arch/base \ + https://github.com/InternLM/archspace.git +cd archspace/reproduce/train-olmo-core +``` + +For an existing checkout or a clone created without `--recurse-submodules`, initialize OLMo-core from the +workflow directory: + +```bash +git -C ../.. submodule sync --recursive +git -C ../.. submodule update --init --recursive +git -C ../.. submodule status --recursive +``` + +The parent repository pins `third_party/OLMo-core` to commit +`45f248d361f0e292c39298d6fba4b4450aed3cc5`. Run the recursive update command again after the parent +repository updates this gitlink. + +### 3.1 Create and activate the training environment + +From `reproduce/train-olmo-core`, create the local training environment. If Python 3.12 is unavailable, +uv obtains a compatible interpreter: + +```bash +uv venv --python 3.12 .venv +source .venv/bin/activate +python --version +``` + +Keep this environment active when running `run/run.sh`; the launcher uses its `python` and `torchrun`. +Reactivate it with `source .venv/bin/activate` in each new shell. + +### 3.2 Install PyTorch and OLMo-core + +Install the CUDA build of PyTorch and the kernel build tools, then install the pinned OLMo-core submodule in +editable mode: + +```bash +uv pip install --index-url https://download.pytorch.org/whl/cu128 \ + torch==2.10.0 torchvision torchaudio + +uv pip install 'setuptools<70' wheel packaging ninja +uv pip install --editable 'third_party/OLMo-core[all]' +``` + +The editable checkout at `third_party/OLMo-core` is the OLMo-core package used by this workflow, so local +source changes are immediately available to the training commands. + +### 3.3 Install attention kernels + +Install the pinned FlashAttention-2 version used by the OLMo-core environment: + +```bash +MAX_JOBS=8 \ +uv pip install --no-build-isolation 'flash-attn==2.8.2' +``` + +The default five-stage configurations use FlashAttention-3 and target Hopper GPUs such as H100 and H800. +Install FA3 from the source commit used by the OLMo-core environment. `FLASH_ATTENTION_FORCE_BUILD=TRUE` +builds the extension from that source, while `--no-cache` prevents reuse of a wheel built with different +feature flags: + +```bash +FLASH_ATTENTION_FORCE_BUILD=TRUE \ +FLASH_ATTENTION_DISABLE_FP16=TRUE \ +FLASH_ATTENTION_DISABLE_SM80=TRUE \ +MAX_JOBS=8 \ +uv pip install --no-cache --no-build-isolation \ + 'flash-attn-3 @ git+https://github.com/Dao-AILab/flash-attention.git@92ca9da8d66f7b34ff50dc080ec0fef9661260d6#subdirectory=hopper' +``` + +For other supported GPUs, use the FA2 installation above and add the following override to `all_stage_args` +in `run/run.sh`: + +```bash +"--model.block.sequence_mixer.backend=flash_2" +``` + +The default HSDP recipe has context parallelism disabled. A custom ring-CP configuration also needs the +version of `ring-flash-attn` used by the OLMo-core environment: + +```bash +uv pip install --no-build-isolation 'ring-flash-attn==0.1.8' +``` + +Adjust `MAX_JOBS` to the CPU and memory available on the build node. + +### 3.4 Verify the training environment + +Run the checks that correspond to the backends installed on this machine: + +```bash +uv pip check +# FA2 (all installations) +python -c 'from olmo_core.nn.attention.flash_attn_api import has_flash_attn_2; assert has_flash_attn_2()' +# FA3 (Hopper only) +python -c 'from olmo_core.nn.attention.flash_attn_api import has_flash_attn_3; assert has_flash_attn_3()' +# ring-flash-attn (optional installation only) +python -c 'from olmo_core.nn.attention.flash_attn_api import has_ring_flash_attn; assert has_ring_flash_attn()' +python -c 'import dion, torch; print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name(0))' +``` + +## 4. Prepare the dataset + +Follow the Dataset Card for +[ArchSpace-Collection/OLMo3-1B-Dataset](https://huggingface.co/datasets/ArchSpace-Collection/OLMo3-1B-Dataset) +to select, download, and extract the data. It also documents storage requirements, layout, provenance, and +license. + +Set `olmo3_data_root` in `run/envs.sh` to the output root passed to the dataset's `extract.sh`. Place this +directory on shared storage at the same path on every training node. + +The selected `run.sh` stages determine which extracted groups must be present: + +| `run.sh` stage | Configuration | Dataset Card groups | Evaluator input | Parent checkpoint when starting here | +| -------------- | ------------- | ------------------- | --------------- | ------------------------------------ | +| stage 1 | pretraining | stage 1 and eval | Dolma 2 `tokenizer_json` | — | +| stage 2 | midtraining | stage 2 and eval | Dolma 2 `tokenizer_json` | stage 1 | +| stage 3 | long-context | stage 3 | — | stage 2 | +| stage 4 | Think SFT | stage 4 | — | stage 3 | +| stage 5 | Instruct SFT | stage 5 | — | stage 4 | + +The default launcher selects all five stages and therefore uses every dataset group. A narrowed stage loop +uses the corresponding groups. If it starts at stage 2–5, make the preceding stage checkpoint available at +the default pipeline path or through `previous_save_folder`. + +## 5. Run training + +### 5.1 Configure the environment and recipe + +```bash +cd reproduce/train-olmo-core +cp run/envs.sh.example run/envs.sh +# edit run/envs.sh +``` + +`run/envs.sh` is excluded by `.gitignore`. Configure these machine-specific values before launching: + +| Setting | Meaning | +| ------- | ------- | +| `olmo3_data_root` | The dataset extraction output root described in Section 4; use the same shared path on every node. | +| `out_root` | A shared, persistent root for checkpoints, trainer state, W&B files, and dataset caches. | +| `tokenizer_json` | A Dolma 2 `tokenizer.json` readable on every node. Required at launcher startup and used by the stage 1/2 downstream evaluator. | +| `wandb_entity`, `wandb_project`, `WANDB_MODE` | W&B destination and mode; see Section 5.3. | + +The launcher accepts two optional positional arguments: + +```bash +bash run/run.sh [TIMESTAMP] [BASE_PORT] +``` + +Configure stage selection, optimizer choice, and Python dotlist overrides in the editable lowercase block +near the top of `run/run.sh`: + +| `run.sh` setting | Purpose | +| ---------------- | ------- | +| `pipeline_name` | Names the training output under `${out_root}/runs/`; edit it when starting a distinct recipe. | +| `all_stage_args` | Python CLI overrides applied to every selected stage, such as optimizer or attention backend. | +| `stage1_args` ... `stage5_args` | Overrides owned by one stage, such as that stage's learning rate. | +| `for stage_index in 1 2 3 4 5` | The selected stages and their execution order. | + +Keep a given dotlist option in one array. `all_stage_args` is appended after the per-stage arguments, so a +duplicate there can override the stage-specific value. Optimizer changes should normally be made in +`all_stage_args` and kept consistent across the checkpoint chain. + +To run only stages 3–5, edit the loop to: + +```bash +for stage_index in 3 4 5; do +``` + +If the first selected stage is stage 2–5, it needs the preceding stage's checkpoint. By default, stage N +loads from `${out_root}/runs/${pipeline_name}/stage(N-1)/checkpoints`. To start from a checkpoint outside +that pipeline, set `previous_save_folder=/path/to/parent/checkpoints` in the corresponding `case` branch in +`run.sh`. + +If you change `data_loader.global_batch_size`, edit it in the stage's Python configuration and then dry-run +the whole pipeline. The configuration derives its temporary checkpoint interval from this value before +merging dotlist arguments. + +### 5.2 Validate the configuration + +Before allocating a full training run, resolve all selected configurations and overrides: + +```bash +ENABLE_WANDB=0 DRY_RUN=1 bash run/run.sh 0816_120000 29500 +``` + +Dry-run visits every selected stage, including stages with an existing `_SUCCESS` marker, and leaves training +outputs unchanged. It validates configuration construction. Before training, separately verify the data +shards, parent checkpoints, distributed topology, and ports. + +### 5.3 W&B + +`envs.sh.example` sets `WANDB_MODE=offline` by default. This mode requires no API key, writes files under each +stage's `trainer/wandb/` directory, and automatically disables remote cancel tags, which only work online. + +Disable W&B for a launch with `ENABLE_WANDB=0`. For online logging, set `WANDB_MODE=online`, `wandb_entity`, +and `wandb_project` in `run/envs.sh`, then export `WANDB_API_KEY` in the calling shell. + +The timestamp identifies the W&B run segment, while `out_root` and `pipeline_name` identify the training +output. Keep `pipeline_name` unchanged when resuming an experiment, and use a new timestamp for each job +attempt so its W&B segment remains distinct. Every node in one multi-node attempt must use the same +timestamp. + +### 5.4 Outputs and shared caches + +```text +out_root/ +├── dataset-cache/ +│ ├── olmo3-stage1/... +│ ├── olmo3-stage2/... +│ ├── olmo3-stage3/... +│ ├── olmo3-stage4/... +│ └── olmo3-stage5/... +└── runs/olmo3-1b/ + ├── stage1/ + │ ├── _SUCCESS + │ ├── checkpoints/ + │ │ └── step/ + │ │ ├── .metadata.json + │ │ ├── config.json + │ │ ├── data_paths.txt + │ │ ├── model_and_optim/ + │ │ │ ├── .metadata + │ │ │ └── ___.distcp + │ │ └── train/ + │ │ └── rank.pt + │ └── trainer/wandb/... + ├── stage2/ + │ └── ... + ├── stage3/ + │ └── ... + ├── stage4/ + │ └── ... + └── stage5/ + └── ... +``` + +`config.json` is the effective configuration, while `data_paths.txt` records the expanded data files that +were actually used. Preserve both together with the W&B records when archiving a reproduction run. The +configurations write a temporary checkpoint approximately every 1 billion tokens, retain only one temporary +checkpoint, and save a final checkpoint at the end of each stage. + +Node rank 0 creates `_SUCCESS` after `torchrun` exits successfully for that stage. Treat it as a process +completion marker; inspect the checkpoint and W&B records for the saved step and metrics. + +`pipeline_name` namespaces stage outputs under `${out_root}/runs/`. Dataset caches use the separate path +`${out_root}/dataset-cache/olmo3-stageN`. When source data, packing, or configuration fingerprints change, +use a new `out_root` or assign a new cache namespace through `data_work_dir` in `run.sh`. + +### 5.5 Resume and stage transitions + +> **Compatibility note:** This five-stage recipe uses dataset fingerprints and data-loader state that differ +> from the earlier three-stage recipe. Start it with a new `pipeline_name` and cache namespace. Reuse earlier +> stage outputs, caches, or `_SUCCESS` markers only after validating the model, optimizer, data, and loader +> state together. + +To resume within the five-stage recipe, keep `out_root` and `pipeline_name` unchanged and choose a new +attempt timestamp: + +```bash +bash run/run.sh 0809_093000 +``` + +The launcher and OLMo-core behave as follows: + +1. A `stageN/_SUCCESS` marker skips that stage during training. +2. Otherwise, a checkpoint under the current stage's `checkpoints/` restores the model, optimizer, trainer, + data-loader, and RNG states. +3. For stages 2–5 with neither marker nor same-stage checkpoint, the top-level `--load_path` initializes the + model and optimizer from the preceding stage without inheriting its step or epoch progress. +4. Stage 1 starts from scratch when it has neither marker nor checkpoint. + +Therefore, rerunning after a stage 2 interruption skips the completed stage 1 and continues from stage 2's +own latest checkpoint. Each later stage starts only after its parent finishes, forming the chain +stage 1 → stage 2 → stage 3 → stage 4 → stage 5. + +## 6. Evaluate with OLMES + +Training produces OLMo-core distributed checkpoints, while OLMES expects a Hugging Face model directory for +a local model. Convert the checkpoint first, then run OLMES. For base-model evaluation, normally use the +final stage 3 checkpoint. For chat or instruction-following evaluation, convert the relevant stage 4 or +stage 5 checkpoint and select a suite appropriate to that model stage. Convert stages separately when making +stage-to-stage comparisons. + +### 6.1 Convert to Hugging Face format + +Set `CHECKPOINT` to a specific `step` directory: + +```bash +export CHECKPOINT=/path/to/out_root/runs/olmo3-1b/stage3/checkpoints/stepNNNNN +export HF_MODEL_DIR=/path/to/out_root/hf/olmo3-1b-stage3-stepNNNNN + +python third_party/OLMo-core/src/examples/huggingface/convert_checkpoint_to_hf.py \ + --checkpoint-input-path "${CHECKPOINT}" \ + --huggingface-output-dir "${HF_MODEL_DIR}" \ + --max-sequence-length 32768 +``` + +The converter reconstructs the OLMo 3 architecture from the checkpoint's `config.json` and uses +`allenai/dolma2-tokenizer` from the configuration by default. In an offline environment, pass +`--tokenizer /path/to/local/hf-tokenizer-directory` with a complete directory loadable by +`AutoTokenizer.from_pretrained()`. + +Keep the default numerical validation enabled. After conversion, run a minimal loading test: + +```bash +python -c 'import os; from transformers import AutoModelForCausalLM, AutoTokenizer; p=os.environ["HF_MODEL_DIR"]; AutoTokenizer.from_pretrained(p); AutoModelForCausalLM.from_pretrained(p); print("HF checkpoint OK")' +``` + +### 6.2 Install and run OLMES + +OLMES requires PyTorch 2.8 and therefore uses a separate uv environment. From +`reproduce/train-olmo-core`, install OLMES with its optional vLLM dependencies from the selected source +commit: + +```bash +uv venv --python 3.12 venv/olmes +uv pip install --python venv/olmes/bin/python \ + --index-url https://download.pytorch.org/whl/cu128 'torch==2.8.0' +uv pip install --python venv/olmes/bin/python \ + 'ai2-olmes[gpu] @ git+https://github.com/allenai/olmes.git@5a51f502d463b8cdc4a2dcad7d7096c41ff1197e' +uv pip check --python venv/olmes/bin/python +source venv/olmes/bin/activate +``` + +The Git URL pins the OLMES source itself. Record +`uv pip freeze --python venv/olmes/bin/python` with the evaluation results to capture its resolved +transitive dependencies. + +For a small-scale stage 3 base-model experiment, start with the OLMo 3 base-easy suites: + +```bash +olmes \ + --model "${HF_MODEL_DIR}" \ + --task \ + olmo3:base_easy:code_bpb \ + olmo3:base_easy:math_bpb \ + olmo3:base_easy:qa_rc \ + olmo3:base_easy:qa_bpb \ + --output-dir /path/to/out_root/evals/olmo3-1b-stage3-base-easy +``` + +If the installed OLMES/vLLM versions support this model, add `--model-type vllm` for higher throughput. A +formal report should preserve the OLMES commit, complete command, task suite, checkpoint step, Hugging Face +conversion arguments, and output directory. Use the configuration's `FAST_TASKS` and PPL in-loop evaluations +for training monitoring, and use a version-pinned OLMES run for final evaluation. For stage 4/5 checkpoints, +choose and record a suitable chat or instruction-following suite. diff --git a/reproduce/train-olmo-core/README_zh.md b/reproduce/train-olmo-core/README_zh.md new file mode 100644 index 0000000..8ee4d12 --- /dev/null +++ b/reproduce/train-olmo-core/README_zh.md @@ -0,0 +1,399 @@ +# OLMo 3 1B 五阶段复现 + +[English](README.md) | 中文 + +本目录提供 OLMo 3 1B 的五阶段训练配方与启动脚本： + +1. stage 1：pretraining； +2. stage 2：midtraining； +3. stage 3：long-context extension； +4. stage 4：Think SFT； +5. stage 5：Instruct SFT。 + +模型定义、分布式训练、checkpoint I/O 和数据加载均由 OLMo-core 提供。本配方所需的 Muon 修复 +固定在 `third_party/OLMo-core` Git 子模块中，模型配置使用 +`TransformerConfig.olmo3_1B()`。 + +## 1. 目录结构 + +```text +reproduce/train-olmo-core/ +├── README.md +├── README_zh.md +├── cfgs/ +│ ├── _olmo3_1b_base.py +│ ├── _olmo3_1b_long.py +│ ├── OLMo3-1B-pretrain.py +│ ├── OLMo3-1B-midtraining.py +│ ├── OLMo3-1B-long-context.py +│ └── OLMo3-1B-sft.py +├── run/ +│ ├── envs.sh.example +│ └── run.sh +└── third_party/ + └── OLMo-core/ # 以可编辑模式安装的 Git 子模块 +``` + +## 2. 配方概览 + +| 阶段 | 数据 | 序列长度 | rank microbatch（token） | 全局 batch（token） | epoch | 默认 Muon LR | +| ------- | ------------------------------------- | -------: | ----------------------: | --------------------: | ----: | -------------: | +| stage 1 | `OLMo_mix_0625_150Bsample` | 4,096 | 16,384 | 2,097,152 | 1 | `5e-3` | +| stage 2 | `OLMo_midtraining_mix_0625_100B` | 4,096 | 16,384 | 1,048,576 | 1 | `2.071235285e-4` | +| stage 3 | `OLMo_longmino_mix_0625` | 32,768 | 32,768 | 2,097,152 | 1 | `2.071235285e-4` | +| stage 4 | `Dolci-Think-SFT-7B` 成对 NPY 文件 | 32,768 | 32,768 | 1,048,576 | 2 | `5e-5` | +| stage 5 | `Dolci-Instruct-SFT` 成对 NPY 文件 | 32,768 | 32,768 | 1,048,576 | 2 | `8e-5` | + +五个阶段默认都使用 BF16、FlashAttention-3 和 Muon。stage 1/2 使用固定长度数据集； +stage 3 使用 Longmino 文档打包和文档内 attention mask，并以 4,096 为原始上下文长度进行 8 倍 +YaRN RoPE 缩放。stage 4/5 继承该长上下文模型，并使用打包且带 assistant mask 的 SFT 数据。 +五个阶段均使用 HSDP，并关闭上下文并行。 +默认的 FlashAttention-3 配置面向 Hopper GPU；其他支持的 GPU 应按 3.3 节切换到 FlashAttention-2。 + +如需选择 SkipStep AdamW 配方，请在 `run/run.sh` 的 `all_stage_args` 中加入 `"--optim=adam"`。 +同一流水线的五个阶段及所有续训任务必须使用同一种优化器，因为后续阶段会继承前一阶段的 +优化器状态。 + +这些实验配置由 OLMo 3 7B 官方配方缩放至 1B 模型，尚未作为正式调优的 OLMo 3 1B 配方发布。 + +## 3. 使用 uv 安装环境 + +本工作流面向 Python 3.12、PyTorch 2.10.0 和 CUDA 12.8。请先安装 +[uv](https://docs.astral.sh/uv/getting-started/installation/)；后续步骤由 uv 创建虚拟环境并安装依赖。 + +首次克隆 ArchSpace 时，可同时初始化 OLMo-core： + +```bash +git clone --recurse-submodules --branch arch/base \ + https://github.com/InternLM/archspace.git +cd archspace/reproduce/train-olmo-core +``` + +已有仓库或首次 clone 未使用 `--recurse-submodules` 时，在当前工作流目录手动初始化 OLMo-core： + +```bash +git -C ../.. submodule sync --recursive +git -C ../.. submodule update --init --recursive +git -C ../.. submodule status --recursive +``` + +父仓库将 `third_party/OLMo-core` 固定在 commit +`45f248d361f0e292c39298d6fba4b4450aed3cc5`。父仓库更新该 gitlink 后，再执行一次递归更新命令 +即可同步。 + +### 3.1 创建并激活训练虚拟环境 + +在 `reproduce/train-olmo-core` 目录创建训练环境；如本机缺少 Python 3.12，uv 会自动获取兼容的 +解释器： + +```bash +uv venv --python 3.12 .venv +source .venv/bin/activate +python --version +``` + +运行 `run/run.sh` 时应保持该环境处于激活状态，启动脚本会调用其中的 `python` 和 `torchrun`。 +每次打开新的 shell 后，执行 `source .venv/bin/activate` 重新激活。 + +### 3.2 安装 PyTorch 与 OLMo-core + +先安装 CUDA 版 PyTorch 和扩展编译工具，再以可编辑模式安装固定版本的 OLMo-core Git 子模块： + +```bash +uv pip install --index-url https://download.pytorch.org/whl/cu128 \ + torch==2.10.0 torchvision torchaudio + +uv pip install 'setuptools<70' wheel packaging ninja +uv pip install --editable 'third_party/OLMo-core[all]' +``` + +训练命令使用 `third_party/OLMo-core` 中以可编辑模式安装的源码，因此本地源码修改会立即 +生效。 + +### 3.3 安装 attention kernel + +先安装与 OLMo-core 配套的固定 FlashAttention-2 版本： + +```bash +MAX_JOBS=8 \ +uv pip install --no-build-isolation 'flash-attn==2.8.2' +``` + +默认五阶段配置使用 FlashAttention-3，面向 H100、H800 等 Hopper GPU。请从 OLMo-core 环境采用的 +commit 安装 FA3。`FLASH_ATTENTION_FORCE_BUILD=TRUE` 会从该源码构建扩展，`--no-cache` 则避免 +复用由其他 feature flags 构建的 wheel： + +```bash +FLASH_ATTENTION_FORCE_BUILD=TRUE \ +FLASH_ATTENTION_DISABLE_FP16=TRUE \ +FLASH_ATTENTION_DISABLE_SM80=TRUE \ +MAX_JOBS=8 \ +uv pip install --no-cache --no-build-isolation \ + 'flash-attn-3 @ git+https://github.com/Dao-AILab/flash-attention.git@92ca9da8d66f7b34ff50dc080ec0fef9661260d6#subdirectory=hopper' +``` + +其他受支持的 GPU 使用上述 FA2，并在 `run/run.sh` 的 `all_stage_args` 中加入以下覆盖： + +```bash +"--model.block.sequence_mixer.backend=flash_2" +``` + +默认 HSDP 配方关闭上下文并行。自定义 ring-CP 配置还需安装与 OLMo-core 配套的 +`ring-flash-attn` 版本： + +```bash +uv pip install --no-build-isolation 'ring-flash-attn==0.1.8' +``` + +`MAX_JOBS` 应按编译节点的 CPU 和内存调整。 + +### 3.4 验证训练环境 + +按本机实际安装的后端执行对应检查： + +```bash +uv pip check +# FA2（所有安装） +python -c 'from olmo_core.nn.attention.flash_attn_api import has_flash_attn_2; assert has_flash_attn_2()' +# FA3（仅 Hopper） +python -c 'from olmo_core.nn.attention.flash_attn_api import has_flash_attn_3; assert has_flash_attn_3()' +# ring-flash-attn（仅可选安装） +python -c 'from olmo_core.nn.attention.flash_attn_api import has_ring_flash_attn; assert has_ring_flash_attn()' +python -c 'import dion, torch; print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name(0))' +``` + +## 4. 准备数据集 + +请按照 +[ArchSpace-Collection/OLMo3-1B-Dataset](https://huggingface.co/datasets/ArchSpace-Collection/OLMo3-1B-Dataset) +的 Dataset Card 选择、下载并解压数据；其中也包含空间需求、文件布局、来源和许可信息。 + +将数据集 `extract.sh` 的输出根目录填入 `run/envs.sh` 的 `olmo3_data_root`。该目录应位于共享 +存储，并在所有训练节点上以同一路径可读。 + +`run.sh` 选中的阶段决定需要准备哪些解压分组： + +| `run.sh` 阶段 | 配置 | Dataset Card 分组 | 评测输入 | 从该阶段开始所需的上游 checkpoint | +| ------------- | ---- | ----------------- | --------------- | ------------------------------------ | +| stage 1 | pretraining | stage 1 与 eval | Dolma 2 `tokenizer_json` | — | +| stage 2 | midtraining | stage 2 与 eval | Dolma 2 `tokenizer_json` | stage 1 | +| stage 3 | long-context | stage 3 | — | stage 2 | +| stage 4 | Think SFT | stage 4 | — | stage 3 | +| stage 5 | Instruct SFT | stage 5 | — | stage 4 | + +默认启动脚本依次运行五个阶段，因此会使用全部数据分组。缩小阶段循环后，只需准备对应分组。 +如果从 stage 2–5 开始，请将前一阶段的 checkpoint 放在默认流水线路径，或通过 +`previous_save_folder` 指定。 + +## 5. 运行训练 + +### 5.1 配置环境与配方 + +```bash +cd reproduce/train-olmo-core +cp run/envs.sh.example run/envs.sh +# 编辑 run/envs.sh +``` + +`.gitignore` 会忽略 `run/envs.sh`。启动前请配置以下机器相关字段： + +| 字段 | 含义 | +| ---- | ---- | +| `olmo3_data_root` | 第 4 节所述的数据集解压输出根目录；所有节点使用同一个共享路径。 | +| `out_root` | checkpoint、训练器状态、W&B 文件和数据集缓存的共享持久化根目录。 | +| `tokenizer_json` | 所有节点均可读的 Dolma 2 `tokenizer.json`。启动脚本的必填项，供 stage 1/2 的下游评测器使用。 | +| `wandb_entity`、`wandb_project`、`WANDB_MODE` | W&B 目标与模式，见第 5.3 节。 | + +启动脚本接受两个可选位置参数： + +```bash +bash run/run.sh [TIMESTAMP] [BASE_PORT] +``` + +阶段选择、优化器和 Python dotlist 覆盖参数统一在 `run/run.sh` 顶部可编辑的小写变量块中配置： + +| `run.sh` 字段 | 用途 | +| ------------ | ---- | +| `pipeline_name` | 决定 `${out_root}/runs/` 下的训练输出名称；启动不同配方时应修改。 | +| `all_stage_args` | 应用于所有选中阶段的 Python CLI 覆盖参数，例如优化器或 attention 后端。 | +| `stage1_args` ... `stage5_args` | 只应用于某一阶段的覆盖参数，例如该阶段的学习率。 | +| `for stage_index in 1 2 3 4 5` | 选择要运行的阶段及其执行顺序。 | + +同一个 dotlist 配置项只应放在一个数组中。`all_stage_args` 在各阶段参数之后追加，因此重复项可能覆盖 +阶段专用值。优化器变更通常应放进 `all_stage_args`，并在整条 checkpoint 链中保持一致。 + +例如，只运行 stage 3–5 时，将循环改为： + +```bash +for stage_index in 3 4 5; do +``` + +如果第一个选中阶段是 stage 2–5，则需要对应的上游 checkpoint。默认情况下，stage N 从 +`${out_root}/runs/${pipeline_name}/stage(N-1)/checkpoints` 加载。若从该流水线之外的 checkpoint +开始，请在 `run.sh` 对应的 `case` 分支中设置 +`previous_save_folder=/path/to/parent/checkpoints`。 + +如需调整 `data_loader.global_batch_size`，请修改对应阶段的 Python 配置，然后重新 dry-run 整条流水线。 +配置会在合并 dotlist 参数之前，根据该值计算临时 checkpoint 间隔。 + +### 5.2 验证配置 + +分配完整训练资源前，先解析所有选中配置及覆盖项： + +```bash +ENABLE_WANDB=0 DRY_RUN=1 bash run/run.sh 0816_120000 29500 +``` + +`dry-run` 会解析每个选中阶段的配置，包括已有 `_SUCCESS` 标记的阶段，且不会修改训练输出。 +正式训练前还需确认数据分片、上游 checkpoint、分布式拓扑和端口。 + +### 5.3 W&B + +`envs.sh.example` 默认 `WANDB_MODE=offline`。这种模式不需要 API 密钥，文件写到每个阶段的 +`trainer/wandb/` 下，并自动禁用只能在线工作的远程取消标签。 + +启动时设置 `ENABLE_WANDB=0` 可关闭 W&B。使用在线记录时，在 `run/envs.sh` 中设置 +`WANDB_MODE=online`、`wandb_entity` 和 `wandb_project`，并在调用 shell 中导出 +`WANDB_API_KEY`。 + +时间戳标识 W&B 运行记录，`out_root` 与 `pipeline_name` 标识训练输出。续训时保持 +`pipeline_name` 不变，并为每次任务使用新的时间戳，使 W&B 记录彼此独立。同一次多机任务的 +所有节点必须使用相同时间戳。 + +### 5.4 输出与共享缓存 + +```text +out_root/ +├── dataset-cache/ +│ ├── olmo3-stage1/... +│ ├── olmo3-stage2/... +│ ├── olmo3-stage3/... +│ ├── olmo3-stage4/... +│ └── olmo3-stage5/... +└── runs/olmo3-1b/ + ├── stage1/ + │ ├── _SUCCESS + │ ├── checkpoints/ + │ │ └── step/ + │ │ ├── .metadata.json + │ │ ├── config.json + │ │ ├── data_paths.txt + │ │ ├── model_and_optim/ + │ │ │ ├── .metadata + │ │ │ └── ___.distcp + │ │ └── train/ + │ │ └── rank.pt + │ └── trainer/wandb/... + ├── stage2/ + │ └── ... + ├── stage3/ + │ └── ... + ├── stage4/ + │ └── ... + └── stage5/ + └── ... +``` + +`config.json` 是最终生效配置，`data_paths.txt` 记录实际展开的数据文件；复现实验归档时应和 W&B +记录一起保存。配置约每 10 亿 token 写一次临时 checkpoint，只保留一个临时版本，并在阶段 +结束时保存最终 checkpoint。 + +`_SUCCESS` 由 node rank 0 在该阶段 `torchrun` 成功退出后创建，用作进程完成标记。保存步数与 +指标应从 checkpoint 和 W&B 记录中确认。 + +`pipeline_name` 为 `${out_root}/runs/` 下的阶段输出划分命名空间；数据集缓存使用独立路径 +`${out_root}/dataset-cache/olmo3-stageN`。当源数据、数据打包方式或配置指纹改变时，请使用 +新的 `out_root`，或通过 `run.sh` 中的 `data_work_dir` 指定新的缓存命名空间。 + +### 5.5 续训与阶段衔接 + +> **兼容性说明：** 五阶段配方的数据集指纹和数据加载器状态与早期三阶段配方不同。请使用新的 +> `pipeline_name` 和缓存命名空间。复用旧阶段输出、缓存或 `_SUCCESS` 标记前，需同时验证模型、 +> 优化器、数据和数据加载器状态的兼容性。 + +在五阶段配方内续训时，保持 `out_root` 和 `pipeline_name` 不变，并使用新的任务时间戳： + +```bash +bash run/run.sh 0809_093000 +``` + +启动器和 OLMo-core 的行为是： + +1. 训练时若存在 `stageN/_SUCCESS`，启动脚本会跳过该阶段； +2. 否则，当前阶段 `checkpoints/` 中的 checkpoint 会恢复模型、优化器、训练器、数据加载器和 + RNG 状态； +3. stage 2–5 既无标记也无本阶段 checkpoint 时，顶层 `--load_path` 会从前一阶段初始化 + 模型和优化器，并从新阶段的 step/epoch 起点开始计数； +4. stage 1 既无标记也无 checkpoint 时从头开始。 + +因此 stage 2 中断后的重跑会跳过已完成的 stage 1，并从 stage 2 自己的最新 checkpoint 继续； +后续阶段只会在各自上游阶段完成后启动，形成 stage 1 → stage 2 → stage 3 → stage 4 → stage 5 +的完整链路。 + +## 6. 使用 OLMES 评测 + +训练输出是 OLMo-core 分布式 checkpoint，而 OLMES 的本地模型入口使用 Hugging Face 模型目录。 +请先转换 checkpoint，再运行 OLMES。基础模型评测通常使用 stage 3 的最终 checkpoint；对话或 +指令遵循评测则转换相应的 stage 4/5 checkpoint，并选择与该阶段匹配的评测套件。做阶段对比时， +应分别转换对应 checkpoint。 + +### 6.1 转换为 Hugging Face 格式 + +将 `CHECKPOINT` 设为具体的 `step` 目录： + +```bash +export CHECKPOINT=/path/to/out_root/runs/olmo3-1b/stage3/checkpoints/stepNNNNN +export HF_MODEL_DIR=/path/to/out_root/hf/olmo3-1b-stage3-stepNNNNN + +python third_party/OLMo-core/src/examples/huggingface/convert_checkpoint_to_hf.py \ + --checkpoint-input-path "${CHECKPOINT}" \ + --huggingface-output-dir "${HF_MODEL_DIR}" \ + --max-sequence-length 32768 +``` + +转换器会从 checkpoint 的 `config.json` 恢复 OLMo 3 架构，并默认使用配置中的 +`allenai/dolma2-tokenizer`。离线环境可传入 +`--tokenizer /path/to/local/hf-tokenizer-directory`，该路径应是可由 +`AutoTokenizer.from_pretrained()` 加载的完整目录。 + +请保留默认开启的数值验证。转换后可先做最小加载测试： + +```bash +python -c 'import os; from transformers import AutoModelForCausalLM, AutoTokenizer; p=os.environ["HF_MODEL_DIR"]; AutoTokenizer.from_pretrained(p); AutoModelForCausalLM.from_pretrained(p); print("HF checkpoint OK")' +``` + +### 6.2 安装并运行 OLMES + +OLMES 需要 PyTorch 2.8，因此使用独立的 uv 虚拟环境。在 `reproduce/train-olmo-core` 中，从固定 +commit 安装 OLMES 及其可选 vLLM 依赖： + +```bash +uv venv --python 3.12 venv/olmes +uv pip install --python venv/olmes/bin/python \ + --index-url https://download.pytorch.org/whl/cu128 'torch==2.8.0' +uv pip install --python venv/olmes/bin/python \ + 'ai2-olmes[gpu] @ git+https://github.com/allenai/olmes.git@5a51f502d463b8cdc4a2dcad7d7096c41ff1197e' +uv pip check --python venv/olmes/bin/python +source venv/olmes/bin/activate +``` + +Git URL 固定的是 OLMES 源码本身。请将 `uv pip freeze --python venv/olmes/bin/python` 与评测结果 +一同保存，以记录实际解析出的间接依赖版本。 + +stage 3 基础模型的小规模实验可从 OLMo 3 base-easy suites 开始： + +```bash +olmes \ + --model "${HF_MODEL_DIR}" \ + --task \ + olmo3:base_easy:code_bpb \ + olmo3:base_easy:math_bpb \ + olmo3:base_easy:qa_rc \ + olmo3:base_easy:qa_bpb \ + --output-dir /path/to/out_root/evals/olmo3-1b-stage3-base-easy +``` + +如 OLMES/vLLM 版本支持该模型，可加 `--model-type vllm` 提高吞吐。正式报告中应保留 OLMES +提交号、完整命令、评测套件、checkpoint step、HF 转换参数和输出目录。`FAST_TASKS`/PPL 用于 +训练中的快速评测与监控，最终评测应使用固定版本的 OLMES。对于 stage 4/5 checkpoint，请选择并 +记录合适的对话或指令遵循评测套件。 diff --git a/reproduce/train-olmo-core/cfgs/OLMo3-1B-long-context.py b/reproduce/train-olmo-core/cfgs/OLMo3-1B-long-context.py new file mode 100644 index 0000000..fd07630 --- /dev/null +++ b/reproduce/train-olmo-core/cfgs/OLMo3-1B-long-context.py @@ -0,0 +1,42 @@ +""" +OLMo 3 1B stage-3 long-context extension configuration. + +This is a 1B adaptation of the OLMo 3 7B long-context recipe in +`src/scripts/official/OLMo3/OLMo-3-1025-7B-long-context.py`. OLMo 3 does not publish an officially +tuned 1B recipe. The local pipeline adds Muon and uses weight decay 0.033 instead of the published +AdamW value 0.1; its Python-default LR follows the published value of approximately 2.071e-4. +""" + +import argparse +from typing import List + +from _olmo3_1b_base import get_olmo3_1b_cli_parser +from _olmo3_1b_long import build_long_context_config + +from olmo_core.script_utils import ExperimentConfig, main + +# Default variables for this stage; supported overrides are shown inline. +lr = 0.00020712352850360292 # Override with `--train_module.optim.lr=1e-3`. +sequence_length = 32768 # Override with `--sequence-length=16384`. +rank_microbatch_size = ( + sequence_length # Override with `--train_module.rank_microbatch_size=16384`. +) +# DO NOT override with `--data_loader.global_batch_size=...`; the ephemeral checkpoint interval +# depends on this value. Edit this default directly instead. +global_batch_size = 2**21 # 2M tokens + + +def build_config(opts: argparse.Namespace, overrides: List[str]) -> ExperimentConfig: + """Apply CLI overrides to the long-context defaults.""" + config = build_long_context_config( + opts, + lr=lr, + sequence_length=sequence_length, + rank_microbatch_size=rank_microbatch_size, + global_batch_size=global_batch_size, + ) + return config.merge(overrides) + + +if __name__ == "__main__": + main(build_config, parser=get_olmo3_1b_cli_parser()) diff --git a/reproduce/train-olmo-core/cfgs/OLMo3-1B-midtraining.py b/reproduce/train-olmo-core/cfgs/OLMo3-1B-midtraining.py new file mode 100644 index 0000000..56af971 --- /dev/null +++ b/reproduce/train-olmo-core/cfgs/OLMo3-1B-midtraining.py @@ -0,0 +1,65 @@ +""" +OLMo 3 1B stage-2 midtraining configuration. + +This is a 1B adaptation of the official OLMo-3-1025-7B midtraining recipe in +`src/scripts/official/OLMo3/OLMo-3-1025-7B-midtrain.py`. OLMo 3 does not publish an officially +tuned 1B recipe. The local pipeline adds Muon and uses weight decay 0.033 instead of the published +AdamW value 0.1; its Python-default LR follows the published value of approximately 2.071e-4. +""" + +import argparse +from typing import List + +from _olmo3_1b_base import ( + build_pretrain_config, + configure_stage_continuation, + get_olmo3_1b_cli_parser, +) + +from olmo_core.data import DataMix +from olmo_core.optim import LinearWithWarmup +from olmo_core.script_utils import ExperimentConfig, main +from olmo_core.train.train_module import ( + TransformerActivationCheckpointingConfig, # noqa: F401 - optional commented config below + TransformerActivationCheckpointingMode, # noqa: F401 - optional commented config below +) + +# Default variables for this stage; supported overrides are shown inline. +lr = 0.00020712352850360292 # Override with `--train_module.optim.lr=1e-3`. +sequence_length = 4096 # Override with `--sequence-length=8192`. +# Override with `--train_module.rank_microbatch_size=8192`. +rank_microbatch_size = 4 * sequence_length +# DO NOT override with `--data_loader.global_batch_size=...`; the ephemeral checkpoint interval +# depends on this value. Edit this default directly instead. +global_batch_size = 2**20 # 1M tokens + + +def build_config(opts: argparse.Namespace, overrides: List[str]) -> ExperimentConfig: + """Apply the midtraining delta and CLI overrides to the short-context baseline.""" + config = build_pretrain_config( + opts, + lr=lr, + sequence_length=sequence_length, + rank_microbatch_size=rank_microbatch_size, + global_batch_size=global_batch_size, + ) + + # Model shape, batching, callbacks, and optimizer parameter groups remain identical to stage 1. + # The data mix, scheduler, and checkpoint continuation policy are stage-specific. + config.dataset.mix = DataMix.OLMo_midtraining_mix_0625_100B + + # Optimizer state is restored from stage 1, so the selected recipe must match. + config.train_module.scheduler = LinearWithWarmup(warmup=0, alpha_f=0.0) + # Official 7B activation checkpointing; intentionally disabled in this local 1B pipeline. + # CLI: `'--train_module.ac_config={mode: selected_modules, modules: ["blocks.*.feed_forward"]}'`. + # config.train_module.ac_config = TransformerActivationCheckpointingConfig( + # mode=TransformerActivationCheckpointingMode.selected_modules, + # modules=["blocks.*.feed_forward"], + # ) + + configure_stage_continuation(config.trainer) + return config.merge(overrides) + + +if __name__ == "__main__": + main(build_config, parser=get_olmo3_1b_cli_parser()) diff --git a/reproduce/train-olmo-core/cfgs/OLMo3-1B-pretrain.py b/reproduce/train-olmo-core/cfgs/OLMo3-1B-pretrain.py new file mode 100644 index 0000000..5aebe96 --- /dev/null +++ b/reproduce/train-olmo-core/cfgs/OLMo3-1B-pretrain.py @@ -0,0 +1,43 @@ +""" +OLMo 3 1B stage-1 pretraining configuration for the local 150B data sample. + +This is a 1B adaptation of the official OLMo-3-1025-7B stage-1 recipe in +``src/scripts/official/OLMo3/OLMo-3-1025-7B-pretrain-1.py``. OLMo 3 does not publish an +officially tuned 1B recipe. The local pipeline adds Muon and uses weight decay 0.033 instead of +the published AdamW value 0.1; its Python-default LR is 5e-3 instead of the published 3e-4. +""" + +import argparse +from typing import List + +from _olmo3_1b_base import build_pretrain_config, get_olmo3_1b_cli_parser + +from olmo_core.script_utils import ExperimentConfig, main + +# Default variables for this stage; supported overrides are shown inline. +lr = 5e-3 # Override with `--train_module.optim.lr=3e-3`. +sequence_length = 4096 # Override with `--sequence-length=8192`. +# Override with `--train_module.rank_microbatch_size=8192`. +rank_microbatch_size = 4 * sequence_length +# DO NOT override with `--data_loader.global_batch_size=...`; the ephemeral checkpoint interval +# depends on this value. Edit this default directly instead. +global_batch_size = 2**21 # 2M tokens + + +def build_config(opts: argparse.Namespace, overrides: List[str]) -> ExperimentConfig: + """Apply CLI overrides to the short-context defaults.""" + # This complete stage-1 recipe, including the padded vocabulary size, is also + # the baseline imported by stage 2. + # Merge CLI overrides only after the shared defaults have been assembled. + config = build_pretrain_config( + opts, + lr=lr, + sequence_length=sequence_length, + rank_microbatch_size=rank_microbatch_size, + global_batch_size=global_batch_size, + ) + return config.merge(overrides) + + +if __name__ == "__main__": + main(build_config, parser=get_olmo3_1b_cli_parser()) diff --git a/reproduce/train-olmo-core/cfgs/OLMo3-1B-sft.py b/reproduce/train-olmo-core/cfgs/OLMo3-1B-sft.py new file mode 100644 index 0000000..9661091 --- /dev/null +++ b/reproduce/train-olmo-core/cfgs/OLMo3-1B-sft.py @@ -0,0 +1,123 @@ +""" +OLMo 3 1B Think SFT (stage 4) and Instruct SFT (stage 5). + +OLMo 3 publishes these post-training hyperparameters for 7B, not 1B. This configuration keeps +the disclosed 7B sequence length, token batch, duration, packing, masking, and schedule while +substituting the local OLMo 3 1B long-context architecture and locally selected phase LRs. + +Sources: + +* OLMo 3 Appendix A.6.1 / Table 47: https://arxiv.org/html/2512.13961v2#A6.SS1 +* Open Instruct Think launcher: + https://github.com/allenai/open-instruct/blob/5fb2acc161b572201628d50cc7145d109edac140/scripts/train/olmo3/7b_think_sft.sh +* Open Instruct Instruct launcher: + https://github.com/allenai/open-instruct/blob/5fb2acc161b572201628d50cc7145d109edac140/scripts/train/olmo3/7b_instruct_sft.sh +* Published OLMo-core Think implementation: + https://github.com/allenai/OLMo-core/blob/38f66526c9d1ba6b97269ebfb429749a5feb528f/src/scripts/train/sft/OLMo2-7B-sft.py +* Published OLMo-core Instruct implementation: + https://github.com/allenai/OLMo-core/blob/9e97471057d7046f0ae7315e0225d117b54186f9/src/scripts/train/sft/OLMo-sft.py + +The optimizer choice is intentionally pipeline-wide. Published SFT uses SkipStep AdamW with +weight decay 0; this pipeline adds Muon and gives both optimizer choices weight decay 0.033, with +an AdamW no-decay embedding group, so optimizer state remains compatible across all five stages. +""" + +import argparse +from typing import List + +from _olmo3_1b_base import get_olmo3_1b_cli_parser +from _olmo3_1b_long import build_long_context_config, source_group_size + +from olmo_core.data import NumpyPackedFSLDatasetConfig, TokenizerConfig +from olmo_core.optim import LinearWithWarmup +from olmo_core.script_utils import ExperimentConfig, main +from olmo_core.train import Duration + +# These subdirectories match the local Open Instruct conversion layout. Each contains paired +# token_ids_part_*.npy and labels_mask_part_*.npy files. +THINK_DATASET_SUBDIR = "Dolci-Think-SFT-7B" +INSTRUCT_DATASET_SUBDIR = "Dolci-Instruct-SFT" + +# Default variables for these SFT stages; supported overrides are shown inline. +think_lr = 5e-5 # Override Think SFT with `--train_module.optim.lr=1e-4`. +instruct_lr = 8e-5 # Override Instruct SFT with `--train_module.optim.lr=1e-4`. +sequence_length = 32768 # Override with `--sequence-length=16384`. +rank_microbatch_size = ( + sequence_length # Override with `--train_module.rank_microbatch_size=16384`. +) +# DO NOT override with `--data_loader.global_batch_size=...`; the ephemeral checkpoint interval +# depends on this value. Edit this default directly instead. +global_batch_size = 2**20 # 1M tokens +# With FlashAttention 3 and CP disabled, B / rank_microbatch_size = 32, so DP must divide 32. + + +def get_sft_cli_parser() -> argparse.ArgumentParser: + """Build the CLI shared by the Think and Instruct SFT phases.""" + parser = get_olmo3_1b_cli_parser() + parser.add_argument( + "--sft-stage", + choices=("think", "instruct"), + required=True, + help="Select the Think or Instruct SFT recipe.", + ) + return parser + + +def build_config(opts: argparse.Namespace, overrides: List[str]) -> ExperimentConfig: + """Apply the selected SFT delta and CLI overrides to the long-context baseline.""" + if opts.sft_stage == "think": + lr = think_lr + dataset_subdir = THINK_DATASET_SUBDIR + elif opts.sft_stage == "instruct": + lr = instruct_lr + dataset_subdir = INSTRUCT_DATASET_SUBDIR + else: + raise ValueError(f"Unknown SFT stage '{opts.sft_stage}'") + + # The model, loader, optimizer, parallelism, clipping, and microbatch settings are inherited + # from the long-context baseline. + config = build_long_context_config( + opts, + lr=lr, + sequence_length=sequence_length, + rank_microbatch_size=rank_microbatch_size, + global_batch_size=global_batch_size, + ) + resolved_sequence_length = config.train_module.max_sequence_length + + # Open Instruct has already applied the chat template. Masked SFT packed data pairs token IDs + # with assistant-only trainable-token masks, uses the default truncate strategy for overlong + # conversations, and records document boundaries for isolated attention. Eight consecutive + # token/mask shard pairs in resolved path order share each OBFD packing pool. + dataset_path = f"{opts.data_root.rstrip('/')}/{dataset_subdir}" + config.dataset = NumpyPackedFSLDatasetConfig( + tokenizer=TokenizerConfig.dolma2(), + paths=[f"{dataset_path}/token_ids_part_*.npy"], + label_mask_paths=[f"{dataset_path}/labels_mask_part_*.npy"], + expand_glob=True, + sequence_length=resolved_sequence_length, + generate_doc_lengths=True, # enables intra-document masking + source_group_size=source_group_size, + work_dir=opts.work_dir, + ) + + config.train_module.scheduler = LinearWithWarmup( + warmup_fraction=0.03, + alpha_f=0.0, + ) + config.train_module.z_loss_multiplier = None + + # Keep the callback/checkpointer profile common to stages 1-3; only SFT duration differs. + # Two epochs follow OLMo 3 Appendix A.6.1 / Table 47. + config.trainer.max_duration = Duration.epochs(2) + config = config.merge(overrides) + if config.dataset.sequence_length != config.train_module.max_sequence_length: + raise ValueError( + "SFT dataset sequence_length must match train_module.max_sequence_length; " + "use --sequence-length to change both" + ) + return config + + +if __name__ == "__main__": + main(build_config, parser=get_sft_cli_parser()) diff --git a/reproduce/olmo-core-backend/cfgs/_olmo3_1b.py b/reproduce/train-olmo-core/cfgs/_olmo3_1b_base.py similarity index 62% rename from reproduce/olmo-core-backend/cfgs/_olmo3_1b.py rename to reproduce/train-olmo-core/cfgs/_olmo3_1b_base.py index 4881adc..a5a33bd 100644 --- a/reproduce/olmo-core-backend/cfgs/_olmo3_1b.py +++ b/reproduce/train-olmo-core/cfgs/_olmo3_1b_base.py @@ -1,4 +1,4 @@ -"""Shared configuration for the OLMo 3 1B training stages.""" +"""Shared base configuration for the OLMo 3 1B training stages.""" import argparse @@ -13,7 +13,6 @@ ) from olmo_core.distributed.parallel import DataParallelType from olmo_core.eval.task_groups import FAST_TASKS -from olmo_core.float8 import Float8Config from olmo_core.nn.attention import AttentionBackendName from olmo_core.nn.transformer import TransformerConfig from olmo_core.optim import ( @@ -24,7 +23,7 @@ SkipStepAdamWConfig, ) from olmo_core.script_utils import ExperimentConfig, get_cli_parser -from olmo_core.train import Duration, TrainerConfig +from olmo_core.train import Duration, LoadStrategy, TrainerConfig from olmo_core.train.callbacks import ( CheckpointerCallback, CometCallback, @@ -40,14 +39,13 @@ TransformerTrainModuleConfig, ) -DEFAULT_SEQUENCE_LENGTH = 4096 -GLOBAL_BATCH_SIZE = 2**21 # 2M tokens -SEED = 34521 -EVAL_LM_STEPS = 500 # 500 steps (~1B token) for 150B data, 2500 steps (~5B token) for 6T data. -EVAL_DOWN_STEPS = 12500 # 12.5K steps (25B tokens) for 150B data -# Keep the current Muon recipe and the official OLMo 3 AdamW recipe independent. -MUON_LR = 1e-3 -ADAM_LR = 1e-3 +# Shared local data-loader policy for all five stages. +seed = 34_521 +num_workers = 8 +prefetch_factor = 4 + +EVAL_LM_STEPS = 500 +EVAL_DOWN_STEPS = 12500 def get_olmo3_1b_cli_parser() -> argparse.ArgumentParser: @@ -62,34 +60,65 @@ def get_olmo3_1b_cli_parser() -> argparse.ArgumentParser: return parser +def build_short_context_model(tokenizer: TokenizerConfig) -> TransformerConfig: + """Build the complete short-context model used by pretraining and midtraining.""" + return TransformerConfig.olmo3_1B( + vocab_size=tokenizer.padded_vocab_size(), # pad to a multiple of 128 + attn_backend=AttentionBackendName.flash_3, + ) + + def build_optim_config( name: str, *, - muon_lr: float, - adam_lr: float, + lr: float, ) -> OptimConfig: - """Build the selected optimizer with its stage-specific learning rate.""" + """Build the pipeline-wide optimizer with the stage learning rate.""" + # All five stages use this local profile so adjacent checkpoints have compatible optimizer + # types and parameter groups. Published OLMo 3 pre/mid/long runs use SkipStep AdamW with + # weight decay 0.1, while published SFT uses weight decay 0. This pipeline instead uses 0.033 + # for both Muon and AdamW; the AdamW embedding group remains exempt from weight decay. + # Both branches inherit optim.compile=False: Muon does not support optimizer-step compilation, + # and the published OLMo 3 AdamW recipes also keep it disabled. Model compilation is separate. # Equivalent whole-object CLI override; define `lr` in the shell first: # "--train_module.optim={type: muon, lr: ${lr}, weight_decay: 0.033, betas: [0.9, 0.95]}" if name == "muon": return MuonConfig( - lr=muon_lr, + lr=lr, weight_decay=0.033, betas=(0.9, 0.95), ) # Equivalent whole-object CLI override; define `lr` in the shell first: # "--train_module.optim={type: skip_step_adamw, lr: ${lr}, weight_decay: 0.033, betas: [0.9, 0.95], group_overrides: [{params: [embeddings.weight], opts: {weight_decay: 0.0}}]}" if name == "adam": - # Match the official OLMo 3 AdamW recipe, including no decay on embeddings. return SkipStepAdamWConfig( - lr=adam_lr, + lr=lr, weight_decay=0.033, betas=(0.9, 0.95), - group_overrides=[OptimGroupOverride(params=["embeddings.weight"], opts={"weight_decay": 0.0})], + group_overrides=[ + OptimGroupOverride( + params=["embeddings.weight"], opts={"weight_decay": 0.0} + ) + ], ) raise ValueError(f"Unknown optimizer '{name}'") +def get_ephemeral_save_interval(global_batch_size: int) -> int: + """Derive a ten-step-aligned checkpoint interval of approximately one Gi tokens.""" + return round(2**30 / global_batch_size / 10) * 10 + + +def configure_stage_continuation(trainer: TrainerConfig) -> None: + """Configure checkpoint loading for a stage that continues from its parent.""" + # script_utils.main first probes save_folder, where these flags require a full same-stage + # trainer+optimizer resume. If none exists, ExperimentConfig.load_path loads the parent while + # explicitly skipping trainer state; optimizer state is retained through load_optim_state=True. + trainer.load_strategy = LoadStrategy.always + trainer.load_trainer_state = True + trainer.load_optim_state = True + + def build_common_config( opts: argparse.Namespace, *, @@ -99,8 +128,7 @@ def build_common_config( train_module: TransformerTrainModuleConfig, ) -> ExperimentConfig: """Build an experiment from required stage components and the shared trainer.""" - # Temporary checkpoint approximately every 1B tokens. - ephemeral_save_interval = round(2**30 / data_loader.global_batch_size / 10) * 10 + ephemeral_save_interval = get_ephemeral_save_interval(data_loader.global_batch_size) trainer = ( TrainerConfig( @@ -115,10 +143,14 @@ def build_common_config( .with_callback( "checkpointer", CheckpointerCallback( - save_interval=None, # Only save the final ckpt + save_interval=None, # No periodic permanent checkpoints; final save remains. ephemeral_save_interval=ephemeral_save_interval, max_checkpoints=1, + # Optional local override: skip the initial pre-training checkpoint. + # CLI: `--trainer.callbacks.checkpointer.pre_train_checkpoint=false`. # pre_train_checkpoint=False, + # Optional local override: force synchronous saves; None auto-selects by backend. + # CLI: `--trainer.callbacks.checkpointer.save_async=false`. # save_async=False, ), ) @@ -147,42 +179,47 @@ def build_common_config( data_loader=data_loader, train_module=train_module, trainer=trainer, + init_seed=seed, ) -def build_pretrain_config(opts: argparse.Namespace) -> ExperimentConfig: - """Build the OLMo 3 1B stage-1 pretraining configuration.""" - sequence_length = opts.sequence_length or DEFAULT_SEQUENCE_LENGTH +def build_pretrain_config( + opts: argparse.Namespace, + *, + lr: float, + sequence_length: int, + rank_microbatch_size: int, + global_batch_size: int, +) -> ExperimentConfig: + """Build the short-context baseline inherited by midtraining.""" + sequence_length = opts.sequence_length or sequence_length tokenizer = TokenizerConfig.dolma2() - model = TransformerConfig.olmo3_1B( - vocab_size=tokenizer.padded_vocab_size(), # pad to a multiple of 128 - attn_backend=AttentionBackendName.flash_3, - ) + model = build_short_context_model(tokenizer) + # Plain FSL concatenates token arrays into fixed contiguous windows; documents may be split, + # and it consumes neither packed-document boundaries nor assistant-label-mask sidecars. dataset = NumpyFSLDatasetConfig.from_data_mix( DataMix.OLMo_mix_0625_150Bsample, tokenizer=tokenizer, mix_base_dir=opts.data_root, sequence_length=sequence_length, - max_target_sequence_length=max(8192, sequence_length), work_dir=opts.work_dir, ) data_loader = NumpyDataLoaderConfig( - global_batch_size=GLOBAL_BATCH_SIZE, - seed=SEED, - num_workers=8, - prefetch_factor=2, + global_batch_size=global_batch_size, + seed=seed, + num_workers=num_workers, + prefetch_factor=prefetch_factor, ) train_module = TransformerTrainModuleConfig( - rank_microbatch_size=4 * DEFAULT_SEQUENCE_LENGTH, + rank_microbatch_size=rank_microbatch_size, max_sequence_length=sequence_length, optim=build_optim_config( opts.optim, - muon_lr=MUON_LR, - adam_lr=ADAM_LR, + lr=lr, ), scheduler=CosWithWarmup(warmup_steps=2000), compile_model=True, @@ -192,7 +229,7 @@ def build_pretrain_config(opts: argparse.Namespace) -> ExperimentConfig: reduce_dtype=DType.float32, wrapping_strategy=TransformerDataParallelWrappingStrategy.blocks, ), - float8_config=Float8Config(enabled=False), + float8_config=None, z_loss_multiplier=1e-5, max_grad_norm=1.0, ) @@ -209,13 +246,12 @@ def build_pretrain_config(opts: argparse.Namespace) -> ExperimentConfig: LMEvaluatorCallbackConfig( eval_dataset=NumpyPaddedFSLDatasetConfig.from_data_mix( DataMix.v3_small_ppl_validation, + tokenizer=tokenizer, mix_base_dir=opts.data_root, sequence_length=sequence_length, - tokenizer=tokenizer, work_dir=opts.work_dir, ), eval_interval=EVAL_LM_STEPS, - # eval_interval=50, ), ).with_callback( "downstream_evaluator", @@ -223,8 +259,6 @@ def build_pretrain_config(opts: argparse.Namespace) -> ExperimentConfig: tasks=sorted(FAST_TASKS), tokenizer=tokenizer, eval_interval=EVAL_DOWN_STEPS, - # eval_interval=50, ), ) - config.init_seed = SEED return config diff --git a/reproduce/train-olmo-core/cfgs/_olmo3_1b_long.py b/reproduce/train-olmo-core/cfgs/_olmo3_1b_long.py new file mode 100644 index 0000000..7304554 --- /dev/null +++ b/reproduce/train-olmo-core/cfgs/_olmo3_1b_long.py @@ -0,0 +1,192 @@ +""" +Shared configuration for the OLMo 3 1B long-context and SFT stages. + +Parallelism boundaries +---------------------- +DP=data-parallel world size; PP/CP/TP/EP are their degrees. +H_rep=HSDP replicas, H_shard=HSDP shard degree, L=seqlen, M=microbatch tokens, +B=global batch tokens; heads=16; n_layers=16. + +Mesh: world_size = PP*CP*TP*DP; world_size % (PP*CP*TP) = 0 +HSDP: DP = H_rep*H_shard; DP % H_shard = 0 +Batch: M % L = 0; B % (M*DP) = 0; grad_accum = B/(M*DP) +CP: local_L = L/CP; exact split requires L % CP = 0 +Ulysses CP: q_heads % CP = kv_heads % CP = 0 +TP: tensor_dim % TP = 0 for every sharded dimension +PP: world_size % PP = 0; num_stages % PP = 0; num_stages <= n_layers +EP: MoE and HSDP only; EP = H_shard; TP = 1 (off) + +Optimizer / parallelism matrix: +| Mode | AdamW | Muon | +|--------------|-----------------|--------------------------------------| +| FSDP | yes | yes: heads % (DP*CP) = 0 | +| HSDP, CP off | yes | yes: heads % H_shard = 0 | +| HSDP + CP | yes | no: dp_shard is flattened into dp_cp | +| TP | yes | no: hard error | +| PP | yes (beta) | beta; changes DP | +| EP | MoE + HSDP only | no: flattened/3D expert parameters | + +Other conflicts: flash_3 has no CP, use flash_2 for CP; Muon also has no CP. +TP + EP is forbidden. Multi-stage PP + tied embeddings is forbidden. +Adjacent stages must keep the same optimizer type and parameter groups while optimizer-state +loading is enabled. + +Stage-3 examples: world_size=64 (GPUs), PP=1 (off), TP=1 (off), heads=16 +B=2^21 (tokens), M=L=32,768 (tokens) +| Optim | DP layout | CP | Muon mesh | Result | +|-------|-------------------------|----|-----------|-------------------------------| +| AdamW | HSDP H_rep=16,H_shard=1 | 4 | - | valid | +| AdamW | HSDP H_rep=8,H_shard=1 | 8 | - | valid | +| Muon | HSDP H_rep=8,H_shard=8 | 1 | 8 | valid:16(heads)%8(mesh)=0 | +| Muon | FSDP DP=64 | 1 | 64 | invalid:16(heads)%64(mesh)!=0 | +| Muon | FSDP DP=16 | 4 | DP*CP=64 | invalid:16(heads)%64(mesh)!=0 | +""" + +import argparse + +from _olmo3_1b_base import ( + build_common_config, + build_optim_config, + configure_stage_continuation, + num_workers, + prefetch_factor, + seed, +) + +from olmo_core.config import DType +from olmo_core.data import ( + DataMix, + NumpyDataLoaderConfig, + NumpyPackedFSLDatasetConfig, + TokenizerConfig, +) +from olmo_core.data.types import LongDocStrategy +from olmo_core.distributed.parallel import DataParallelType +from olmo_core.float8 import ( # noqa: F401 - optional commented configuration below + AOFloat8LinearConfig, + Float8Config, +) +from olmo_core.nn.attention import AttentionBackendName +from olmo_core.nn.rope import YaRNRoPEScalingConfig +from olmo_core.nn.transformer import TransformerConfig +from olmo_core.optim import LinearWithWarmup +from olmo_core.script_utils import ExperimentConfig +from olmo_core.train.train_module import ( + TransformerActivationCheckpointingConfig, # noqa: F401 - optional commented config below + TransformerActivationCheckpointingMode, # noqa: F401 - optional commented config below + TransformerContextParallelConfig, # noqa: F401 - optional commented config below + TransformerDataParallelConfig, + TransformerDataParallelWrappingStrategy, + TransformerTrainModuleConfig, +) + +# The long-context model/YaRN pattern, packed-dataset structure, and linear scheduler are adapted from: +# src/scripts/official/OLMo3/OLMo-3-1025-7B-long-context.py. + +source_group_size = 8 + + +def build_long_context_model(tokenizer: TokenizerConfig) -> TransformerConfig: + """Build the complete checkpoint-compatible model used by stages 3-5.""" + return TransformerConfig.olmo3_1B( + vocab_size=tokenizer.padded_vocab_size(), # pad to a multiple of 128 + attn_backend=AttentionBackendName.flash_3, + ).with_rope_scaling( + YaRNRoPEScalingConfig( + factor=8, + beta_fast=32, + beta_slow=1, + old_context_len=4096, + ) + ) + + +def build_long_context_config( + opts: argparse.Namespace, + *, + lr: float, + sequence_length: int, + rank_microbatch_size: int, + global_batch_size: int, +) -> ExperimentConfig: + """Build the long-context baseline inherited by both SFT phases.""" + sequence_length = opts.sequence_length or sequence_length + tokenizer = TokenizerConfig.dolma2() + + model = build_long_context_model(tokenizer) + + # Longmino packed data records document lengths for isolated attention and jointly packs + # consecutive source files in resolved data-mix order, but has no label-mask sidecar. + dataset = NumpyPackedFSLDatasetConfig.from_data_mix( + DataMix.OLMo_longmino_mix_0625, + tokenizer=tokenizer, + mix_base_dir=opts.data_root, + sequence_length=sequence_length, + generate_doc_lengths=True, # enables intra-document masking + long_doc_strategy=LongDocStrategy.truncate, + source_group_size=source_group_size, + work_dir=opts.work_dir, + ) + + data_loader = NumpyDataLoaderConfig( + global_batch_size=global_batch_size, + seed=seed, + num_workers=num_workers, + prefetch_factor=prefetch_factor, + ) + + # Model, packed dataset, loader, and train module all differ materially from stage 2, so this + # baseline constructs those configs in full instead of overwriting short-context members. + train_module = TransformerTrainModuleConfig( + rank_microbatch_size=rank_microbatch_size, + max_sequence_length=sequence_length, + optim=build_optim_config(opts.optim, lr=lr), + scheduler=LinearWithWarmup(warmup=200, alpha_f=0.0), + compile_model=True, + dp_config=TransformerDataParallelConfig( + name=DataParallelType.hsdp, + param_dtype=DType.bfloat16, + reduce_dtype=DType.float32, + wrapping_strategy=TransformerDataParallelWrappingStrategy.full, + ), + # The API alternatives below affect stages 3-5; use CLI overrides for one stage only. + # Optional shared context-parallel profiles. CP cannot be used with FlashAttention 3: + # switch the model backend to FlashAttention 2 and select Adam pipeline-wide first. + # Common CLI: `--model.block.sequence_mixer.backend=flash_2 --optim=adam`. + # Stage-3 CLI: `'--train_module.cp_config={degree: 4, ring: {load_balancer: llama3, head_stride: 4}}'`. + # cp_config=TransformerContextParallelConfig.llama3(degree=4, head_stride=4), + # SFT CLI: `'--train_module.cp_config={degree: 2, ring: {load_balancer: llama3, head_stride: 4}}'`. + # cp_config=TransformerContextParallelConfig.llama3(degree=2, head_stride=4), + cp_config=None, + # Optional activation-checkpointing profiles. + # CLI: `'--train_module.ac_config={mode: budget, activation_memory_budget: 0.7}'`. + # ac_config=TransformerActivationCheckpointingConfig( + # mode=TransformerActivationCheckpointingMode.budget, + # activation_memory_budget=0.7, + # ) + # CLI: `'--train_module.ac_config={mode: selected_modules, modules: ["blocks.*.feed_forward"]}'`. + # ac_config=TransformerActivationCheckpointingConfig( + # mode=TransformerActivationCheckpointingMode.selected_modules, + # modules=["blocks.*.feed_forward"], + # ) + ac_config=None, + # Published 7B float8 profile; the local 1B baseline keeps float8 disabled. + # CLI: `'--train_module.float8_config={enabled: true, ao: {enable_fsdp_float8_all_gather: true, force_recompute_fp8_weight_in_bwd: true, round_scales_to_power_of_2: true}}'`. + # float8_config=Float8Config( + # enabled=True, + # ao=AOFloat8LinearConfig.recommended(), + # ) + float8_config=None, + z_loss_multiplier=1e-5, + max_grad_norm=1.0, + ) + + config = build_common_config( + opts, + model=model, + dataset=dataset, + data_loader=data_loader, + train_module=train_module, + ) + configure_stage_continuation(config.trainer) + return config diff --git a/reproduce/train-olmo-core/run/envs.sh.example b/reproduce/train-olmo-core/run/envs.sh.example new file mode 100644 index 0000000..56921c6 --- /dev/null +++ b/reproduce/train-olmo-core/run/envs.sh.example @@ -0,0 +1,25 @@ +#!/usr/bin/env bash + +# Copy this file to envs.sh and replace the placeholders with paths available on every node. +# envs.sh is intentionally ignored by Git because it contains machine-specific configuration. + +# Local paths used to construct training arguments. These variables do not need to be exported. +# This root contains the stage 1-3 data trees and validation tree referenced by the built-in +# OLMo 3 data-mix manifests, plus the stage 4/5 SFT directories documented in README.md. +olmo3_data_root=/path/to/olmo3-data +out_root=/path/to/training-output +# Used only by the downstream evaluator; LM eval data is resolved under olmo3_data_root. +tokenizer_json=/path/to/tokenizer.json + +# W&B destination. The launcher defaults to online mode unless WANDB_MODE is set to offline. +wandb_entity=YOUR_ENTITY +wandb_project=YOUR_PROJECT +export WANDB_MODE=offline + +# For online logging +# export WANDB_API_KEY=YOUR_SECRET + +# Optional cluster setup. Uncomment and adapt only what this machine needs. +# export CUDA_HOME=/path/to/cuda +# export PATH=${CUDA_HOME}/bin:${PATH} +# export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH:-} diff --git a/reproduce/olmo-core-backend/run/run.sh b/reproduce/train-olmo-core/run/run.sh similarity index 65% rename from reproduce/olmo-core-backend/run/run.sh rename to reproduce/train-olmo-core/run/run.sh index a9d3ec6..f05a021 100755 --- a/reproduce/olmo-core-backend/run/run.sh +++ b/reproduce/train-olmo-core/run/run.sh @@ -9,20 +9,37 @@ echo "Local environment file not found. Copy ${script_dir}/envs.sh.example to ${env_file}." >&2 exit 1 } - # shellcheck source=/dev/null + # shellcheck source=envs.sh.example source "${env_file}" - # All nodes in one distributed attempt must receive the same timestamp and - # base port. Give every resume attempt a new timestamp so its W&B ID differs. + # All nodes in one distributed attempt must receive the same timestamp and base port. + # Give every resume attempt a new timestamp so its W&B ID differs. timestamp=${1:-$(date +'%m%d_%H%M%S')} base_port=${2:-29500} pipeline_name=olmo3-1b config_basename=${reproduce_dir}/cfgs/OLMo3-1B - extra_args=( - # For a short smoke run, uncomment both overrides. Do not use them for - # the full reproduction. - # "--trainer.max_duration.value=10" - # "--trainer.max_duration.unit=steps" + + # Optional CLI overrides applied to every stage. + all_stage_args=( + # "--optim=adam" # switch to Adam for the entire five-stage pipeline + # "--trainer.max_duration.value=50" "--trainer.max_duration.unit=steps" # smoke run + ) + # Optional CLI overrides owned by one stage. + stage1_args=( + # "--train_module.optim.lr=1e-3" + # "--train_module.scheduler={type: wsd_sqrt_decay, warmup: 2000, decay_fraction: 0.2, decay_min_lr_ratio: 0.1}" + ) + stage2_args=( + # "--train_module.optim.lr=5e-4" + ) + stage3_args=( + # "--train_module.optim.lr=5e-4" + ) + stage4_args=( + # "--train_module.optim.lr=1e-4" + ) + stage5_args=( + # "--train_module.optim.lr=2e-4" ) olmo3_data_root=${olmo3_data_root:?Set olmo3_data_root in envs.sh} @@ -30,30 +47,43 @@ tokenizer_json=${tokenizer_json:?Set tokenizer_json in envs.sh} pipeline_root=${out_root}/runs/${pipeline_name} - for stage_index in 1 2 3; do + for stage_index in 1 2 3 4 5; do stage="stage${stage_index}" - previous_save_folder= stage_args=() + previous_save_folder= + # Set previous_save_folder in a stage branch to override the automatic parent checkpoint. case "${stage}" in stage1) config_file=${config_basename}-pretrain.py stage_args=( "--trainer.callbacks.downstream_evaluator.tokenizer.identifier=${tokenizer_json}" + "${stage1_args[@]}" ) ;; stage2) config_file=${config_basename}-midtraining.py stage_args=( "--trainer.callbacks.downstream_evaluator.tokenizer.identifier=${tokenizer_json}" + "${stage2_args[@]}" ) - previous_save_folder=${pipeline_root}/stage1/checkpoints ;; stage3) config_file=${config_basename}-long-context.py - previous_save_folder=${pipeline_root}/stage2/checkpoints + stage_args=("${stage3_args[@]}") + ;; + stage4) + config_file=${config_basename}-sft.py + stage_args=("--sft-stage=think" "${stage4_args[@]}") + ;; + stage5) + config_file=${config_basename}-sft.py + stage_args=("--sft-stage=instruct" "${stage5_args[@]}") ;; esac + if ((stage_index > 1)) && [[ -z "${previous_save_folder}" ]]; then + previous_save_folder=${pipeline_root}/stage$((stage_index - 1))/checkpoints + fi stage_port=$((base_port + stage_index)) run_name=${pipeline_name}-${stage} @@ -74,9 +104,8 @@ "--trainer.work_dir=${run_root}/trainer" ) if [[ -n "${previous_save_folder}" ]]; then - # On a fresh stage this initializes model and optimizer state from the - # parent stage without loading parent trainer progress. If the current - # stage already has a checkpoint, OLMo-core resumes full local state. + # script_utils.main loads this top-level path without trainer state only when the current + # stage has no checkpoint. TrainerConfig requires trainer state for same-stage resume. train_args+=("--load_path=${previous_save_folder}") fi @@ -85,7 +114,7 @@ "--trainer.callbacks.wandb.enabled=true" "--trainer.callbacks.wandb.entity=${wandb_entity:?Set wandb_entity in envs.sh}" "--trainer.callbacks.wandb.project=${wandb_project:?Set wandb_project in envs.sh}" - "--trainer.callbacks.wandb.group=${pipeline_name}" + "--trainer.callbacks.wandb.group=${run_name}" "--trainer.callbacks.wandb.name=${run_name}_${timestamp}" # wandb.id = wandb.name ) if [[ "${WANDB_MODE:-online}" != "offline" ]]; then @@ -96,10 +125,7 @@ train_args+=("--trainer.callbacks.wandb.cancel_tags=null") fi fi - train_args+=( - "${stage_args[@]}" - "${extra_args[@]}" - ) + train_args+=("${stage_args[@]}" "${all_stage_args[@]}") echo "Starting ${stage} for pipeline '${pipeline_name}' on port ${stage_port}" if [[ "${DRY_RUN:-0}" == "1" ]]; then @@ -119,11 +145,11 @@ ) fi mkdir -p "${run_root}" - torchrun "${torchrun_args[@]}" "${config_file}" -- "${train_args[@]}" + torchrun "${torchrun_args[@]}" "${config_file}" "${train_args[@]}" [[ "${NODE_RANK:-0}" == "0" ]] && touch "${success_marker}" done - echo "Pipeline '${pipeline_name}' completed all three stages" + echo "Pipeline '${pipeline_name}' completed all five stages" exit } diff --git a/reproduce/train-olmo-core/third_party/OLMo-core b/reproduce/train-olmo-core/third_party/OLMo-core new file mode 160000 index 0000000..45f248d --- /dev/null +++ b/reproduce/train-olmo-core/third_party/OLMo-core @@ -0,0 +1 @@ +Subproject commit 45f248d361f0e292c39298d6fba4b4450aed3cc5