allenai / allenai/open-instruct

Drop OLMoE support

Offen
#1,837 0 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
Vorherrschende Sprache
Python
Sterne
3.9k
Forks
585
Ø Merge
5 T. 17 Std.
Gemergte PRs (30 T.)
16

Beschreibung

Per discussion on #1820: OLMoE is an old architecture and we no longer need to support it. Dropping it also removes the only model type reachable through the OLMo-core export path whose HF weight naming requires transformers' `conversion_mapping` translation, which lets `save_state_dict_as_hf` stop reasoning about two naming schemes.

## Where OLMoE is supported

**OLMo-core SFT path** (the part implicated by #1820):

- `open_instruct/olmo_core_utils.py:355` — `OLMO_MODEL_CONFIG_MAP` entry `"allenai/OLMoE-1B-7B-0924": "olmoe_1B_7B"`
- `README.md:88` — lists OLMoE among supported OLMo-core SFT models

**HF trainer paths** (separate code, not affected by #1820, but part of "OLMoE support"):

- `open_instruct/olmo_adapter/__init__.py:6,113-117` — `OlmoeConfig` / `OlmoeModel` imports and `OlmoeForSequenceClassification`
- `open_instruct/reward_modeling.py:167` — registers `OlmoeForSequenceClassification`
- `open_instruct/dpo_utils.py:108`, `open_instruct/finetune.py:268` — `--load_balancing_loss`, documented as being for OLMoE
- `open_instruct/context_window_checker.py:60,63` — `"olmoe"`, `"allenai/olmoe"`
- `configs/train_configs/sft/olmoe_preview_mix_v3.1.yaml` — OLMoE SFT config
- `docs/algorithms/dpo.md:47,245` — load-balancing-loss documentation

**Incidental mentions** (comments, archived scripts — probably leave alone):

- `open_instruct/dataset_transformation.py:784`, `open_instruct/mix_data.py:31`, `open_instruct/mix_data_preferences.py:23` (comments)
- `scripts/submit_eval_jobs_old.py:618`, `docs/archived_dev_scripts/olmoe_0125.sh`, `docs/get_started/ai2_internal_setup.md:102`

The two groups can be done independently; only the first is needed to close out #1820.

## Note on the transformers 5.5 interaction

While inventorying this: OLMo-core's converter (`src/olmo_core/nn/hf/convert.py`) maps MoE weights to the **per-expert** HF names (`model.layers.N.mlp.experts.{i}.gate_proj.weight`) and has no fused `gate_up_proj` mapping. If transformers ≥5.5 represents OLMoE with fused expert tensors in memory — as reported in the review on #1820 — then `hf_model.load_state_dict(converted)` in `save_state_dict_as_hf`, and `verify_can_save_as_hf` before it, would already fail for OLMoE on that version, independent of #1820.

That would mean OLMoE export is broken on transformers ≥5.5 today rather than something #1820 introduces. Not verified against 5.5.3 — worth confirming before deciding whether this is a "drop it" or a "fix the converter" for anyone who still needs OLMoE.

Companion issue for the OLMo-core side: allenai/OLMo-core (filed separately).

Beitragsleitfaden

Beitragsleitfaden öffnen

Bewertung

Dieses Issue wurde noch nicht bewertet.

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.