moe.py

# MoE (Mixture of Experts) implementation for Falcon Vision
# Simplified from torchtitan's MoE for standalone use

from dataclasses import dataclass
from typing import Literal

import torch
import torch.nn.functional as F
from torch import nn


@dataclass
class MoEArgs:
    num_experts: int = 8
    num_shared_experts: int = 1
    score_func: Literal["softmax", "sigmoid"] = "sigmoid"
    route_norm: bool = False
    route_scale: float = 1.0
    score_before_experts: bool = True
    top_k: int = 1
    use_grouped_mm: bool = False  # disabled by default for compatibility
    activation: Literal["silu", "relu2"] = "silu"


class FeedForward(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, activation: str = "silu"):
        super().__init__()
        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
        self.act = activation

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.act == "relu2":
            return self.w2(2 * F.relu(self.w1(x)).square() * self.w3(x))
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

    def init_weights(self, init_std: float = 0.02):
        nn.init.trunc_normal_(self.w1.weight, mean=0.0, std=0.02)
        nn.init.trunc_normal_(self.w3.weight, mean=0.0, std=init_std)
        nn.init.zeros_(self.w2.weight)


def _run_experts_for_loop(
    w1: torch.Tensor,
    w2: torch.Tensor,
    w3: torch.Tensor,
    x: torch.Tensor,
    num_tokens_per_expert: torch.Tensor,
    act: str = "silu",
) -> torch.Tensor:
    num_tokens_list = num_tokens_per_expert.to(torch.int32).tolist()
    total_tokens = sum(num_tokens_list)    
    num_padding = x.shape[0] - total_tokens
    x_splits = torch.split(x[:total_tokens], split_size_or_sections=num_tokens_list, dim=0)
    out_splits = []
    for expert_idx, x_expert in enumerate(x_splits):
        if act == "relu2":
            h = 2 * F.relu(torch.matmul(x_expert, w1[expert_idx].T)).square()
        else:
            h = F.silu(torch.matmul(x_expert, w1[expert_idx].T))
        h = h * torch.matmul(x_expert, w3[expert_idx].T)
        h = torch.matmul(h, w2[expert_idx].T)
        out_splits.append(h)
    
    out = torch.cat(out_splits, dim=0)
    if num_padding > 0:
        out = torch.vstack((out, out.new_zeros((num_padding, out.shape[-1]))))
    return out


class GroupedExperts(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, num_experts: int, activation: str = "silu"):
        super().__init__()
        self.num_experts = num_experts
        self.w1 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim))
        self.w2 = nn.Parameter(torch.empty(num_experts, dim, hidden_dim))
        self.w3 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim))
        self.activation = activation

    def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor:
        return _run_experts_for_loop(
            self.w1, self.w2, self.w3, x, num_tokens_per_expert, self.activation
        )

    def init_weights(self, init_std: float):
        nn.init.trunc_normal_(self.w1, mean=0.0, std=0.02)
        nn.init.zeros_(self.w2)
        nn.init.trunc_normal_(self.w3, mean=0.0, std=init_std)


class TokenChoiceTopKRouter(nn.Module):
    def __init__(
        self,
        dim: int,
        num_experts: int,
        top_k: int,
        score_func: str = "sigmoid",
        route_norm: bool = False,
        route_scale: float = 1.0,
    ):
        super().__init__()
        self.gate = nn.Linear(dim, num_experts, bias=False)
        self.num_experts = num_experts
        self.top_k = top_k
        self.score_func = score_func
        self.route_norm = route_norm
        self.route_scale = route_scale

    def forward(self, x: torch.Tensor, expert_bias: torch.Tensor | None = None):
        scores = self.gate(x)
        if self.score_func == "sigmoid":
            scores = torch.sigmoid(scores.float())
        else:
            scores = F.softmax(scores.float(), dim=1)

        if expert_bias is not None:
            _, selected_experts_indices = torch.topk(scores + expert_bias, k=self.top_k, dim=1)
        else:
            _, selected_experts_indices = torch.topk(scores, k=self.top_k, dim=1)

        top_scores = scores.gather(dim=1, index=selected_experts_indices)
        if self.route_norm:
            top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20)
        top_scores = top_scores * self.route_scale

        num_tokens_per_expert = torch.histc(
            selected_experts_indices.view(-1).float(),
            bins=self.num_experts,
            min=0,
            max=self.num_experts,
        )
        return top_scores, selected_experts_indices, num_tokens_per_expert

    def init_weights(self, init_std: float):
        nn.init.trunc_normal_(self.gate.weight, mean=0.0, std=init_std)


class MoE(nn.Module):
    def __init__(self, moe_args: MoEArgs, dim: int, hidden_dim: int):
        super().__init__()
        num_experts = moe_args.num_experts

        self.experts = GroupedExperts(
            dim=dim,
            hidden_dim=hidden_dim,
            num_experts=num_experts,
            activation=moe_args.activation,
        )
        self.router = TokenChoiceTopKRouter(
            dim=dim,
            num_experts=num_experts,
            top_k=moe_args.top_k,
            score_func=moe_args.score_func,
            route_norm=moe_args.route_norm,
            route_scale=moe_args.route_scale,
        )
        self.shared_experts = (
            FeedForward(dim=dim, hidden_dim=hidden_dim * moe_args.num_shared_experts, activation=moe_args.activation)
            if moe_args.num_shared_experts > 0
            else None
        )
        self.score_before_experts = moe_args.score_before_experts
        self.top_k = moe_args.top_k
        
        # Register buffer for load balancing (matches torchtitan checkpoint)
        self.register_buffer(
            "expert_bias",
            torch.zeros(num_experts, dtype=torch.float32),
            persistent=True,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        bs, slen, dim = x.shape
        x = x.view(-1, dim)

        top_scores, selected_experts_indices, num_tokens_per_expert = self.router(x, expert_bias=self.expert_bias)

        # Reorder tokens by expert
        token_indices_sorted = torch.argsort(selected_experts_indices.view(-1), stable=True)
        top_scores_sorted = top_scores.view(-1)[token_indices_sorted]
        token_indices_sorted = token_indices_sorted // self.top_k

        token_indices_expanded = token_indices_sorted.view(-1, 1).expand(-1, dim)
        routed_input = torch.gather(x, dim=0, index=token_indices_expanded)

        if self.score_before_experts:
            routed_input = (routed_input.float() * top_scores_sorted.view(-1, 1)).to(x.dtype)

        routed_output = self.experts(routed_input, num_tokens_per_expert)

        if self.shared_experts is not None:
            out = self.shared_experts(x)
        else:
            out = torch.zeros_like(x)

        routed_output = (
            routed_output.to(torch.float32)
            * top_scores_sorted.view(-1, 1)
        ).to(x.dtype)

        out = out.scatter_add(dim=0, index=token_indices_expanded, src=routed_output)
        return out.view(bs, slen, dim)

    def init_weights(self, init_std: float, buffer_device: torch.device = None):
        self.experts.init_weights(init_std)
        self.router.init_weights(init_std)
        if self.shared_experts is not None:
            self.shared_experts.init_weights(init_std)