feat: Add OLMoE-1B-7B-0924 support

[prabhashj07]

Closes #2133

Adds full support for Allen AI's OLMoE-1B-7B-0924 Mixture-of-Experts model — config, checkpoint conversion (both directions), prompt style, and numerical equivalence tests. Also fixes a latent bug in MoE top-k routing probability computation that affected all MoE architectures.

Background

OLMoE is a 7B parameter MoE model with 1B active parameters, using 64 experts with 8 active per token. It is well-suited for academic research. LitGPT already supports other OLMo models (OLMo, OLMo2), so adding OLMoE is a natural extension. See #2133.

Changes

Bug Fix: MoE router probability ordering (model.py)

The previous implementation applied topk on raw router logits and then softmax on the top-k subset. This is mathematically incorrect — softmax over a subset of logits produces different values than softmax over all logits followed by selection.

This matches HuggingFace transformers behavior and is required for numerical equivalence with HF checkpoints. The change is safe for existing MoE models (Mixtral, Qwen3-MoE) since softmax-then-topk and topk-then-softmax produce identical indices; only probability values change, and those were already renormalized downstream.

Added norm_topk_prob config flag to re-normalize top-k probabilities to sum to 1 (used by Mixtral and Qwen3-MoE).

OLMoE-1B-7B-0924 Model Support

• litgpt/config.py: Config entries for OLMoE-1B-7B-0924 base, Instruct, and SFT variants. 64 experts, 8 active per token, GQA 16/8, RMSNorm with per-head QK norm.
• litgpt/prompts.py: OLMoE ChatML prompt style for instruct variants.
• litgpt/scripts/convert_hf_checkpoint.py: copy_weights_olmoe — handles HF per-expert weight layout and per-head q_norm/k_norm.
• litgpt/scripts/convert_lit_checkpoint.py: copy_weights_olmoe — reverse conversion, splits fused qkv back into separate q/k/v projections.
• tests/test_model.py: test_against_hf_olmoe — end-to-end numerical equivalence test.

Config additions for existing models (config.py)

• Added norm_topk_prob=True to all Mixtral-8x7B and Mixtral-8x22B variants
• Added norm_topk_prob=True to all Qwen3-MoE and Qwen3-2507 MoE variants

LoRA / Adapter Support

OLMoE uses the existing LLaMAMoE MLP class and standard attention, so LoRA and adapter_v2 work out of the box — no additional changes needed in lora.py or adapter_v2.py.

Notes

• OLMoE config sets norm_topk_prob=False matching the reference implementation's behavior.
• The intermediate_size field is set to the full n_embd value as a dense fallback required by Config.post_init, but is never used — MoE blocks use moe_intermediate_size (1024 per expert).
• No init.py exports were changed since copy_weights_olmoe is only referenced via partial() in the conversion scripts.