From 9b22b2b9002d81e813f9006214eb652522ea714b Mon Sep 17 00:00:00 2001 From: DanieleS <59105724+danielesalpietro@users.noreply.github.com> Date: Sat, 18 Jul 2026 20:08:59 +0200 Subject: [PATCH 1/3] Fix OLMoE support and silent model-init failures - moe_infinity/utils/hf_config.py: parse_moe_param()/parse_expert_id() were missing an "olmoe" branch, even though olmoe is registered in moe_infinity/common/constants.py's MODEL_MAPPING_NAMES/ MODEL_MAPPING_TYPES and has a working monkey-patch class (SyncOlmoeMoEBlock in moe_infinity/models/olmoe.py). Loading any OLMoE checkpoint raised RuntimeError("Unsupported architecture ...forcausallm"). OLMoE's config fields (num_experts, num_experts_per_tok) and expert parameter naming (layers.N.mlp.experts.M....) match the already-handled qwen3 branch, so "olmoe" was added there in both functions. The same registry/parser gap still exists for dbrx, jamba, and opt. - moe_infinity/entrypoints/openai/api_server_v2.py: _initialize_model() had no except clause, and runs as a fire-and-forget asyncio.create_task that nothing ever awaits or inspects (startup_event()). Any exception raised while constructing the model (unsupported architecture, CUDA allocation failure, etc.) was silently dropped: the process kept running, /health stayed on "starting" forever, and every request got a 503 indistinguishable from a genuine hang -- this is what made the OLMoE bug above so hard to diagnose in the first place. The exception is now logged (logger.exception) and surfaced through /health as {"status": "unhealthy", "reason": ""}. Verified: allenai/OLMoE-1B-7B-0924-Instruct now loads and serves correctly via /v1/chat/completions with these fixes. Co-Authored-By: Claude Sonnet 5 --- moe_infinity/entrypoints/openai/api_server_v2.py | 12 ++++++++++-- moe_infinity/utils/hf_config.py | 4 ++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/moe_infinity/entrypoints/openai/api_server_v2.py b/moe_infinity/entrypoints/openai/api_server_v2.py index bc04631d..6a846043 100644 --- a/moe_infinity/entrypoints/openai/api_server_v2.py +++ b/moe_infinity/entrypoints/openai/api_server_v2.py @@ -142,6 +142,7 @@ def _decorator(func: Any) -> Any: _contextpilot_fallback_count: int = 0 _contextpilot_last_fallback_count: int = 0 +logger = logging.getLogger(__name__) _cp_logger = logging.getLogger("moe_infinity.contextpilot") _cp_middleware: Optional[Any] = None _eviction_sync: Optional[Any] = None @@ -1093,12 +1094,19 @@ async def _initialize_model() -> None: ) _ensure_engine_loop_running() + _health_state.set_healthy() + except Exception as exc: + # _model_init_task is fire-and-forget (asyncio.create_task, never + # awaited) -- without this, an exception here is silently dropped: + # the process keeps running, /health stays stuck on "starting" + # forever, and every request gets a 503 indistinguishable from a + # genuine hang. Log it and surface it through /health instead. + logger.exception("Model initialization failed") + _health_state.set_unhealthy(f"{type(exc).__name__}: {exc}") finally: if _startup_watchdog is not None: _startup_watchdog.cancel() - _health_state.set_healthy() - @app.on_event("startup") async def startup_event() -> None: diff --git a/moe_infinity/utils/hf_config.py b/moe_infinity/utils/hf_config.py index bed2e19c..b1a3b7b0 100644 --- a/moe_infinity/utils/hf_config.py +++ b/moe_infinity/utils/hf_config.py @@ -78,7 +78,7 @@ def parse_moe_param(config: PretrainedConfig) -> Tuple[int, int, int]: num_decoder_layers = config.num_hidden_layers num_layers = config.num_hidden_layers num_experts = config.num_local_experts - elif "qwen3" in arch: + elif "qwen3" in arch or "olmoe" in arch: num_encoder_layers = 0 num_decoder_layers = config.num_hidden_layers num_layers = config.num_hidden_layers @@ -147,7 +147,7 @@ def parse_expert_id( layer_id, expert_id = result[0] layer_id = int(layer_id) expert_id = int(expert_id) - elif "deepseek" in arch or "qwen3" in arch: + elif "deepseek" in arch or "qwen3" in arch or "olmoe" in arch: decoder_sparse_step = 1 layer_type = "decoder" From dd2779f5a8e6ef4af0eb8f95229ea9cf94a843c7 Mon Sep 17 00:00:00 2001 From: DanieleS <59105724+danielesalpietro@users.noreply.github.com> Date: Sun, 19 Jul 2026 17:10:37 +0200 Subject: [PATCH 2/3] ci: drop exact-match setuptools pin in build-system requires Same root cause as EfficientMoE/MoE-Infinity#121: installing torch pulls in setuptools>=77.0.3 as a transitive dependency, upgrading the environment's setuptools past the exact pin. Under --no-isolation, python -m build then fails its build-dependency check because the installed version no longer matches ==75.3.2. Co-Authored-By: Claude Sonnet 5 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index b11a21c0..c31eb89c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools==75.3.2", "wheel", "torch"] +requires = ["setuptools>=75.3.2", "wheel", "torch"] build-backend = "setuptools.build_meta" From bbde96c521aade2d6d733ab5ffc6d4fa71e3f7c7 Mon Sep 17 00:00:00 2001 From: DanieleS <59105724+danielesalpietro@users.noreply.github.com> Date: Mon, 20 Jul 2026 02:20:44 +0200 Subject: [PATCH 3/3] fix(olmoe): correct expert-type mapping to fix fused_moe_ffn_into hidden dim mismatch MODEL_MAPPING_TYPES["olmoe"] was set to 4 (MIXTRAL_MOE_DENSE_ACT_DENSE), which makes MoEMLP::ForwardHelper read expert weight buffers in Mixtral's gate/down/up order (matching w1/w2/w3). But OlmoeMLP registers weights as gate_proj/up_proj/down_proj (same convention as qwen3, already mapped to 5/DEEPSEEK_MOE_DENSE_ACT_DENSE). This swapped up_proj and down_proj for every OLMoE expert, crashing fused_moe_ffn_into's H == H_out check whenever hidden_size != intermediate_size (e.g. allenai/OLMoE-1B-7B-0924-Instruct: hidden_size=2048, intermediate_size=1024). Fixes #123 Co-Authored-By: Claude Sonnet 5 --- moe_infinity/common/constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/moe_infinity/common/constants.py b/moe_infinity/common/constants.py index fd6420af..41d72bd8 100644 --- a/moe_infinity/common/constants.py +++ b/moe_infinity/common/constants.py @@ -39,7 +39,7 @@ "gptoss": 4, "qwen3": 5, "dbrx": 4, - "olmoe": 4, + "olmoe": 5, "jamba": 4, }