Skip to content

feat(moe): add elastic expert fallback executor with weight redistribution and residual bypass - #2457

Open
presentyourlove wants to merge 1 commit into
exo-explore:mainfrom
presentyourlove:feature/moe-elastic-fallback-executor
Open

presentyourlove wants to merge 1 commit into
exo-explore:mainfrom
presentyourlove:feature/moe-elastic-fallback-executor

Conversation

@presentyourlove

Copy link
Copy Markdown

Description

This PR implements the Elastic MoE Fallback Executor (ElasticMoEFallbackExecutor), providing resilient dynamic expert dropout handling, weight redistribution, and residual highway bypass to prevent inference failures caused by edge node timeouts or network disconnections.

Background & Motivation

In edge decentralized MoE clusters operating over volatile wireless/consumer networks, remote worker nodes hosting routed experts can experience transient latency spikes (>100ms), packet loss, thermal sleep, or abrupt disconnections. Traditional distributed MoE runtimes crash or stall indefinitely when a selected expert fails to return activations. This PR guarantees continuous generation by gracefully redistributing lost weights among successful experts, or falling back to a zero-delta residual highway bypass if all selected experts drop.

Key Changes

  1. Elastic MoE Fallback Data Models (src/exo/shared/types/moe_fallback.py):

    • ElasticMoEConfig: Immutable configuration specifying expert_timeout_ms, enable_residual_bypass, and min_successful_experts.
    • ExpertExecutionResult: Tracks expert dispatch status (SUCCESS, TIMEOUT, DISCONNECTED), original gating weights, and latency.
    • MoEFallbackDecision: Captures successful vs skipped experts, redistributed weights, and bypass state.
    • MoEFallbackTelemetry: Computes lost weight fractions, timeout/disconnection counts, and bypass occurrences.
  2. Elastic Fallback Executor Engine (src/exo/worker/engines/moe/elastic_fallback_executor.py):

    • Dynamic Weight Redistribution: Re-normalizes the weights of surviving experts to preserve output magnitude:
      $$\hat{w}{e} = \frac{w{e}}{\sum_{k \in \text{successful}} w_{k}}$$
    • Residual Highway Bypass: If all experts time out or disconnect, returns zero-delta activations so that $\text{Activation} + \text{MoE}(0) = \text{Activation}$, allowing the Transformer forward pass to continue uninterrupted without crashing.
  3. Module Exports (src/exo/worker/engines/moe/__init__.py):

    • Exported ElasticMoEFallbackExecutor.
  4. Unit Tests (src/exo/shared/tests/test_elastic_moe_fallback.py):

    • Verified standard multi-expert aggregation when all succeed.
    • Verified weight redistribution when a single expert times out (surviving expert scales to 1.0).
    • Verified re-normalization across partial disconnections.
    • Verified zero-delta residual highway bypass when all experts fail.

Verification

  • Ran unit tests:
    .venv/Scripts/pytest.exe src/exo/shared/tests/test_elastic_moe_fallback.py
    # 4 passed in 0.17s
  • Linting and type checking:
    .venv/Scripts/ruff.exe check src/exo/shared/types/moe_fallback.py src/exo/worker/engines/moe/elastic_fallback_executor.py src/exo/worker/engines/moe/__init__.py src/exo/shared/tests/test_elastic_moe_fallback.py
    # All checks passed!

Related Issues

  • Resolves: Frontier Dimension 3.2 in 人人都能跑llm.md (Dynamic Elastic Expert Pruning & Early-Exit Fallback)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant