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] 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 +}