allenai / allenai/OLMo-core

Drop OLMoE-1B-7B support

未關閉
#838 2 則留言 0 個 reaction 已指派 0 人 在 GitHub 檢視
主要語言
Python
星號
1.5k
分支
315
平均合併
1 天 9 小時
30 天內合併 PR
11

描述

Companion to allenai/open-instruct#1837, which drops OLMoE from open-instruct. Filing the OLMo-core half so the two can be coordinated.

## Where OLMoE-1B-7B is supported

- `src/olmo_core/nn/transformer/config.py:1080` — `TransformerConfig.olmoe_1B_7B`
- `src/scripts/train/OLMoE-1B-7B.py` — training script (calls the above at line 26)
- `src/examples/huggingface/convert_checkpoint_from_hf.py:69` — `"olmoe_1b_7b": TransformerConfig.olmoe_1B_7B` in the config lookup

## What should NOT be removed

Several things match a naive `grep -i olmoe` but are unrelated, and removing them would break other models:

- **`DataMix.OLMoE_mix_0824`** (`src/olmo_core/data/mixes/__init__.py:39`) — a data mix, not a model. It is used by `src/scripts/official/OLMo2/OLMo-2-0325-32B-train.py:52` and `src/olmo_core/internal/experiment.py:236`, and covered by `src/test/data/mixes_test.py`.
- **`src/scripts/train/OLMoE2-medium.py` / `OLMoE2-large.py`** — OLMoE2, a different and newer MoE architecture.
- **The MoE machinery itself** (`feed_forward_moe`, the per-expert mappings in `src/olmo_core/nn/hf/convert.py:51-53,149-165`) — shared with OLMoE2 and any future MoE model.
- `OLMoEnvironmentError` and friends in `src/olmo_core/exceptions.py` — substring collision only.

## Background

The prompt for this was a checkpoint-export bug in open-instruct (allenai/open-instruct#1820), where transformers' `conversion_mapping` rewrote weight names on save and produced checkpoints vLLM could not load. OLMoE is the one model in open-instruct's export path that legitimately needs that rewriting, so it is the one case that keeps the export logic branching on model type.

Worth confirming before acting: OLMo-core's converter maps MoE weights to the per-expert HF names and has no fused `gate_up_proj` mapping. If transformers ≥5.5 represents OLMoE with fused expert tensors in memory, the HF conversion path would already fail for OLMoE on that version. If OLMoE-1B-7B is still wanted here, that may be a converter fix rather than a removal.

貢獻指南

開啟貢獻指南

研究方向

Read the OLMoE-1B-7B entry in src/olmo_core/nn/transformer/config.py, then inspect src/scripts/train/OLMoE-1B-7B.py and the config lookup in src/examples/huggingface/convert_checkpoint_from_hf.py. Review the converter’s MoE mappings before deciding whether removal or a converter fix is appropriate; done means the intended support path is resolved without removing OLMoE2, shared MoE machinery, or DataMix.OLMoE_mix_0824.

由索引模型根據 Issue 內容生成。

評估

技術堆疊
python
領域
machine-learning
Issue 類型
重構
難度
3/5
預估耗時
1-2 天
活躍度
活躍
描述清晰度
基本清楚
新手友好度
64/100

把新 issue 寄到你的電子郵件信箱

精選適合新手參與的 GitHub issue 摘要。