modeling_fuse_glm.py

"""LFM2 with sparse, folded GLM-5.3-Flash coding experts.

The implementation intentionally depends only on upstream ``transformers``
LFM2 classes.  It does not import GLM modeling code and it never downloads a
checkpoint at import or construction time.  Folded expert tensors can be
copied into the exposed ``gate_proj``, ``up_proj`` and ``down_proj`` modules
after a separate extraction/folding step.
"""

from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Iterator, Literal

import torch
import torch.nn.functional as F
from torch import nn
from transformers import Lfm2Config, Lfm2ForCausalLM
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_fuse_glm import FuseGlmConfig


class _AddAuxiliaryLoss(torch.autograd.Function):
    """Attach an auxiliary scalar gradient without changing forward values.

    This is useful here because LFM2 decoder layers return only hidden states.
    Returning the router loss through that API would require replacing the
    entire decoder stack and would also interfere with generation caches.  The
    identity operation keeps the native forward contract and remains valid
    when gradient checkpointing recomputes a decoder layer.
    """

    @staticmethod
    def forward(ctx: Any, hidden_states: torch.Tensor, auxiliary_loss: torch.Tensor) -> torch.Tensor:
        ctx.auxiliary_dtype = auxiliary_loss.dtype
        ctx.auxiliary_device = auxiliary_loss.device
        return hidden_states

    @staticmethod
    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        auxiliary_grad = torch.ones((), dtype=ctx.auxiliary_dtype, device=ctx.auxiliary_device)
        return grad_output, auxiliary_grad


@dataclass(frozen=True)
class RouterState:
    """Differentiable routing state from one fused decoder layer."""

    layer_index: int
    token_indices: torch.Tensor
    router_logits: torch.Tensor
    topk_indices: torch.Tensor
    topk_weights: torch.Tensor
    token_gate: torch.Tensor
    auxiliary_loss: torch.Tensor

    def detached(self) -> "RouterState":
        return RouterState(
            layer_index=self.layer_index,
            token_indices=self.token_indices.detach(),
            router_logits=self.router_logits.detach(),
            topk_indices=self.topk_indices.detach(),
            topk_weights=self.topk_weights.detach(),
            token_gate=self.token_gate.detach(),
            auxiliary_loss=self.auxiliary_loss.detach(),
        )


@dataclass(frozen=True)
class RouterDiagnostics:
    """Compact, detached-by-default statistics for monitoring routing."""

    layer_index: int
    token_count: int
    active_token_count: int
    expert_counts: torch.Tensor
    mean_selected_weights: torch.Tensor
    router_entropy: torch.Tensor
    token_gate_mean: torch.Tensor
    token_gate_active_fraction: torch.Tensor
    residual_scale: torch.Tensor
    auxiliary_loss: torch.Tensor
    coding_enabled: bool

    def detached(self) -> "RouterDiagnostics":
        return RouterDiagnostics(
            layer_index=self.layer_index,
            token_count=self.token_count,
            active_token_count=self.active_token_count,
            expert_counts=self.expert_counts.detach(),
            mean_selected_weights=self.mean_selected_weights.detach(),
            router_entropy=self.router_entropy.detach(),
            token_gate_mean=self.token_gate_mean.detach(),
            token_gate_active_fraction=self.token_gate_active_fraction.detach(),
            residual_scale=self.residual_scale.detach(),
            auxiliary_loss=self.auxiliary_loss.detach(),
            coding_enabled=self.coding_enabled,
        )


@dataclass
class FuseGlmCausalLMOutputWithPast(CausalLMOutputWithPast):
    """Causal LM output augmented with sparse-router training information."""

    router_aux_loss: torch.FloatTensor | None = None
    router_diagnostics: tuple[RouterDiagnostics, ...] | None = None


class FoldedGlmExpert(nn.Module):
    """A folded GLM-clamped SwiGLU expert in the LFM hidden space.

    For the production configuration all three dimensions are 2048.  The
    separate ``intermediate_size`` argument exists to enable inexpensive unit
    tests and later structured compression experiments.
    """

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        *,
        gate_clamp_max: float = 10.0,
        up_clamp_min: float = -10.0,
        up_clamp_max: float = 10.0,
        initializer_range: float = 0.02,
    ) -> None:
        super().__init__()
        self.hidden_size = int(hidden_size)
        self.intermediate_size = int(intermediate_size)
        self.gate_clamp_max = float(gate_clamp_max)
        self.up_clamp_min = float(up_clamp_min)
        self.up_clamp_max = float(up_clamp_max)

        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.reset_parameters(initializer_range)

    def reset_parameters(self, initializer_range: float) -> None:
        for projection in (self.gate_proj, self.up_proj, self.down_proj):
            nn.init.normal_(projection.weight, mean=0.0, std=initializer_range)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        gate = self.gate_proj(hidden_states).clamp(max=self.gate_clamp_max)
        up = self.up_proj(hidden_states).clamp(min=self.up_clamp_min, max=self.up_clamp_max)
        return self.down_proj(F.silu(gate) * up)

    @torch.no_grad()
    def load_folded_weights(
        self,
        *,
        gate_proj: torch.Tensor,
        up_proj: torch.Tensor,
        down_proj: torch.Tensor,
    ) -> None:
        """Validate and copy three already-folded expert matrices."""

        supplied = {
            "gate_proj": gate_proj,
            "up_proj": up_proj,
            "down_proj": down_proj,
        }
        modules = {
            "gate_proj": self.gate_proj,
            "up_proj": self.up_proj,
            "down_proj": self.down_proj,
        }
        for name, tensor in supplied.items():
            expected_shape = tuple(modules[name].weight.shape)
            if tuple(tensor.shape) != expected_shape:
                raise ValueError(f"{name} has shape {tuple(tensor.shape)}; expected {expected_shape}")
            modules[name].weight.copy_(tensor.to(device=modules[name].weight.device, dtype=modules[name].weight.dtype))


class TopKFoldedExpertRouter(nn.Module):
    """GLM-style sigmoid top-k router with post-sigmoid choice correction."""

    def __init__(self, hidden_size: int, num_experts: int, top_k: int, initializer_range: float) -> None:
        super().__init__()
        self.num_experts = int(num_experts)
        self.top_k = int(top_k)
        self.proj = nn.Linear(hidden_size, num_experts, bias=False)
        nn.init.normal_(self.proj.weight, mean=0.0, std=initializer_range)
        # GLM adds this value only while choosing experts.  The routed mixture
        # weights are gathered from the uncorrected sigmoid scores, so this
        # must not be represented as a Linear bias.
        self.register_buffer(
            "e_score_correction_bias",
            torch.zeros(self.num_experts, dtype=torch.float32),
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        # GLM computes router logits in FP32 even when expert weights are FP8
        # or BF16.  Keeping that behavior also stabilizes small synthetic tests.
        logits = F.linear(hidden_states.float(), self.proj.weight.float())
        scores = torch.sigmoid(logits)
        choice_scores = scores + self.e_score_correction_bias.float()
        selected_indices = torch.topk(choice_scores, self.top_k, dim=-1).indices
        selected_scores = scores.gather(-1, selected_indices)
        selected_weights = selected_scores / selected_scores.sum(dim=-1, keepdim=True).clamp_min(1e-12)

        # Switch-style balancing objective.  The top-k assignment fraction is
        # normalized by k, making a perfectly balanced value equal to 1.
        probabilities = torch.softmax(logits, dim=-1)
        assignment = F.one_hot(selected_indices, num_classes=self.num_experts).float().sum(dim=-2)
        assignment = assignment / float(self.top_k)
        probability_fraction = probabilities.mean(dim=0)
        token_fraction = assignment.mean(dim=0)
        auxiliary_loss = self.num_experts * torch.sum(probability_fraction * token_fraction)
        return logits, selected_indices, selected_weights, auxiliary_loss


class FuseGlmFeedForward(nn.Module):
    """Preserve the native LFM FFN and add a sparse expert sidecar in parallel.

    The native projections remain registered directly as ``w1``, ``w2`` and
    ``w3``.  Consequently their state-dict keys are identical to an ordinary
    ``Lfm2ForCausalLM`` checkpoint even after this wrapper is installed.
    """

    def __init__(
        self,
        base_ffn: nn.Module,
        config: FuseGlmConfig,
        layer_index: int,
        *,
        auxiliary_loss_scale: float,
    ) -> None:
        super().__init__()
        for name in ("w1", "w2", "w3"):
            if not hasattr(base_ffn, name):
                raise TypeError(f"Unsupported LFM2 feed_forward module: missing {name}")

        # Preserve original checkpoint paths: feed_forward.w1/w2/w3.
        self.w1 = base_ffn.w1
        self.w2 = base_ffn.w2
        self.w3 = base_ffn.w3

        self.layer_index = int(layer_index)
        self.num_experts = config.fuse_glm_num_experts
        self.top_k = config.fuse_glm_top_k
        self.residual_scale_max = config.fuse_glm_residual_scale_max
        self.token_gate_threshold = config.fuse_glm_token_gate_threshold
        self.hard_token_gate_at_eval = config.fuse_glm_hard_token_gate_at_eval
        self.coding_enabled = config.fuse_glm_coding_enabled
        self.auxiliary_loss_scale = float(auxiliary_loss_scale)

        self.experts = nn.ModuleList(
            [
                FoldedGlmExpert(
                    config.hidden_size,
                    config.fuse_glm_expert_intermediate_size,
                    gate_clamp_max=config.fuse_glm_gate_clamp_max,
                    up_clamp_min=config.fuse_glm_up_clamp_min,
                    up_clamp_max=config.fuse_glm_up_clamp_max,
                    initializer_range=config.initializer_range,
                )
                for _ in range(self.num_experts)
            ]
        )
        self.router = TopKFoldedExpertRouter(
            config.hidden_size,
            self.num_experts,
            self.top_k,
            config.initializer_range,
        )
        self.token_gate = nn.Linear(config.hidden_size, 1, bias=True)
        nn.init.zeros_(self.token_gate.weight)
        nn.init.constant_(self.token_gate.bias, config.fuse_glm_token_gate_bias)

        # tanh(0) is exactly zero, so a newly constructed fused model computes
        # the same function as its LFM host while retaining nonzero expert
        # activations from which this scale can learn.
        self.raw_residual_scale = nn.Parameter(torch.zeros(()))
        self.last_router_state: RouterState | None = None
        self.last_router_diagnostics: RouterDiagnostics | None = None

    @property
    def residual_scale(self) -> torch.Tensor:
        return self.residual_scale_max * torch.tanh(self.raw_residual_scale)

    def base_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.w2(F.silu(self.w1(hidden_states)) * self.w3(hidden_states))

    def _disabled_diagnostics(self, hidden_states: torch.Tensor) -> RouterDiagnostics:
        scalar_zero = hidden_states.new_zeros(())
        return RouterDiagnostics(
            layer_index=self.layer_index,
            token_count=hidden_states.numel() // hidden_states.shape[-1],
            active_token_count=0,
            expert_counts=torch.zeros(self.num_experts, dtype=torch.long, device=hidden_states.device),
            mean_selected_weights=hidden_states.new_zeros(self.num_experts),
            router_entropy=scalar_zero,
            token_gate_mean=scalar_zero,
            token_gate_active_fraction=scalar_zero,
            residual_scale=self.residual_scale,
            auxiliary_loss=scalar_zero,
            coding_enabled=False,
        )

    def _dispatch(
        self,
        hidden_states: torch.Tensor,
        selected_indices: torch.Tensor,
        selected_weights: torch.Tensor,
    ) -> torch.Tensor:
        output = torch.zeros_like(hidden_states)
        for expert_index, expert in enumerate(self.experts):
            token_positions, route_slots = torch.where(selected_indices == expert_index)
            if token_positions.numel() == 0:
                continue
            expert_input = hidden_states.index_select(0, token_positions)
            expert_output = expert(expert_input)
            route_weight = selected_weights[token_positions, route_slots].to(expert_output.dtype).unsqueeze(-1)
            output = output.index_add(0, token_positions, expert_output * route_weight)
        return output

    def _make_diagnostics(
        self,
        state: RouterState,
        token_count: int,
        active_mask: torch.Tensor,
    ) -> RouterDiagnostics:
        expert_counts = F.one_hot(state.topk_indices, num_classes=self.num_experts).sum(dim=(0, 1))
        selected_weight_sums = torch.zeros(
            self.num_experts,
            device=state.topk_weights.device,
            dtype=state.topk_weights.dtype,
        )
        selected_weight_sums.scatter_add_(0, state.topk_indices.reshape(-1), state.topk_weights.reshape(-1))
        mean_selected_weights = selected_weight_sums / expert_counts.clamp_min(1).to(selected_weight_sums.dtype)

        if state.router_logits.shape[0] == 0:
            entropy = state.router_logits.new_zeros(())
        else:
            probabilities = torch.softmax(state.router_logits, dim=-1)
            entropy = -(probabilities * probabilities.clamp_min(1e-12).log()).sum(dim=-1).mean()
        gate = state.token_gate
        return RouterDiagnostics(
            layer_index=self.layer_index,
            token_count=token_count,
            active_token_count=int(state.token_indices.numel()),
            expert_counts=expert_counts,
            mean_selected_weights=mean_selected_weights,
            router_entropy=entropy,
            token_gate_mean=gate.mean(),
            token_gate_active_fraction=active_mask.float().mean(),
            residual_scale=self.residual_scale,
            auxiliary_loss=state.auxiliary_loss,
            coding_enabled=True,
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        base_output = self.base_forward(hidden_states)
        if not self.coding_enabled:
            self.last_router_state = None
            self.last_router_diagnostics = self._disabled_diagnostics(hidden_states)
            return base_output

        original_shape = hidden_states.shape
        flat_hidden = hidden_states.reshape(-1, original_shape[-1])
        token_gate = torch.sigmoid(F.linear(flat_hidden.float(), self.token_gate.weight.float(), self.token_gate.bias.float()))
        token_gate = token_gate.squeeze(-1)
        active_mask = token_gate >= self.token_gate_threshold

        # Soft gating is used while training so the token gate itself receives
        # dense gradients.  Optional hard gating at evaluation actually skips
        # expert computation for low-confidence non-coding tokens.
        if self.hard_token_gate_at_eval and not self.training:
            token_indices = torch.where(active_mask)[0]
        else:
            token_indices = torch.arange(flat_hidden.shape[0], device=flat_hidden.device)

        if token_indices.numel() == 0:
            expert_delta = torch.zeros_like(flat_hidden)
            auxiliary_loss = flat_hidden.sum() * 0.0
            state = RouterState(
                layer_index=self.layer_index,
                token_indices=token_indices,
                router_logits=flat_hidden.new_empty((0, self.num_experts), dtype=torch.float32),
                topk_indices=torch.empty((0, self.top_k), dtype=torch.long, device=flat_hidden.device),
                topk_weights=flat_hidden.new_empty((0, self.top_k), dtype=torch.float32),
                token_gate=token_gate,
                auxiliary_loss=auxiliary_loss,
            )
        else:
            active_hidden = flat_hidden.index_select(0, token_indices)
            logits, selected_indices, selected_weights, auxiliary_loss = self.router(active_hidden)
            active_delta = self._dispatch(active_hidden, selected_indices, selected_weights)
            active_gate = token_gate.index_select(0, token_indices).to(active_delta.dtype).unsqueeze(-1)
            expert_delta = torch.zeros_like(flat_hidden).index_add(
                0,
                token_indices,
                active_delta * active_gate,
            )
            state = RouterState(
                layer_index=self.layer_index,
                token_indices=token_indices,
                router_logits=logits,
                topk_indices=selected_indices,
                topk_weights=selected_weights,
                token_gate=token_gate,
                auxiliary_loss=auxiliary_loss,
            )

        self.last_router_state = state
        self.last_router_diagnostics = self._make_diagnostics(state, flat_hidden.shape[0], active_mask)

        output = base_output + self.residual_scale.to(expert_delta.dtype) * expert_delta.reshape(original_shape)
        if self.training and self.auxiliary_loss_scale > 0.0:
            output = _AddAuxiliaryLoss.apply(output, auxiliary_loss * self.auxiliary_loss_scale)
        return output


class FuseGlmForCausalLM(Lfm2ForCausalLM):
    """LFM2 causal LM augmented with folded GLM coding experts."""

    config_class = FuseGlmConfig
    _no_split_modules = ["Lfm2DecoderLayer", "FuseGlmFeedForward", "FoldedGlmExpert"]
    _keep_in_fp32_modules_strict = [
        "router.proj.weight",
        "e_score_correction_bias",
        "token_gate.weight",
        "token_gate.bias",
        "raw_residual_scale",
    ]

    def __init__(self, config: FuseGlmConfig) -> None:
        if not isinstance(config, FuseGlmConfig):
            if isinstance(config, Lfm2Config):
                config = FuseGlmConfig.from_lfm_config(config)
            else:
                raise TypeError("config must be FuseGlmConfig or Lfm2Config")
        super().__init__(config)
        self._install_fusion_wrappers()

    @torch.no_grad()
    def _initialize_missing_keys(self, is_quantized: bool) -> None:
        """Initialize fusion-only keys correctly when loading a base LFM.

        ``from_pretrained`` constructs models under an empty/meta-parameter
        context.  The generic Transformers initializer handles Linear weights
        but does not know that our scalar scale must be zero, nor that the
        token gate needs a zero weight and negative bias.  Capture the names
        that were absent from the checkpoint, let upstream initialize all
        ordinary parameters, then repair only those absent fusion parameters.
        Existing values from a fused checkpoint are never overwritten.
        """

        missing_names = {
            name
            for name, parameter in self.named_parameters()
            if not getattr(parameter, "_is_hf_initialized", False)
        }
        super()._initialize_missing_keys(is_quantized)

        parameters = dict(self.named_parameters())
        for name in missing_names:
            parameter = parameters.get(name)
            if parameter is None:
                continue
            if name.endswith("feed_forward.raw_residual_scale"):
                nn.init.zeros_(parameter)
            elif name.endswith("feed_forward.token_gate.weight"):
                nn.init.zeros_(parameter)
            elif name.endswith("feed_forward.token_gate.bias"):
                nn.init.constant_(parameter, self.config.fuse_glm_token_gate_bias)

    def _install_fusion_wrappers(self) -> None:
        layer_indices = self.config.resolved_fuse_glm_layer_indices
        auxiliary_scale = (
            self.config.fuse_glm_router_aux_loss_coef / len(layer_indices) if layer_indices else 0.0
        )
        selected = set(layer_indices)
        for layer_index, decoder_layer in enumerate(self.model.layers):
            if layer_index not in selected:
                continue
            if isinstance(decoder_layer.feed_forward, FuseGlmFeedForward):
                continue
            decoder_layer.feed_forward = FuseGlmFeedForward(
                decoder_layer.feed_forward,
                self.config,
                layer_index,
                auxiliary_loss_scale=auxiliary_scale,
            )

    def fusion_layers(self) -> tuple[FuseGlmFeedForward, ...]:
        return tuple(
            layer.feed_forward
            for layer in self.model.layers
            if isinstance(layer.feed_forward, FuseGlmFeedForward)
        )

    def iter_folded_experts(self) -> Iterator[tuple[int, int, FoldedGlmExpert]]:
        for wrapper in self.fusion_layers():
            for expert_index, expert in enumerate(wrapper.experts):
                yield wrapper.layer_index, expert_index, expert

    def get_router_states(self, *, detach: bool = False) -> tuple[RouterState, ...]:
        states = tuple(
            wrapper.last_router_state
            for wrapper in self.fusion_layers()
            if wrapper.last_router_state is not None
        )
        if detach:
            return tuple(state.detached() for state in states)
        return states

    def get_router_diagnostics(self, *, detach: bool = True) -> tuple[RouterDiagnostics, ...]:
        diagnostics = tuple(
            wrapper.last_router_diagnostics
            for wrapper in self.fusion_layers()
            if wrapper.last_router_diagnostics is not None
        )
        if detach:
            return tuple(item.detached() for item in diagnostics)
        return diagnostics

    def get_router_aux_loss(
        self,
        reduction: Literal["mean", "sum", "none"] = "mean",
        *,
        detach: bool = False,
    ) -> torch.Tensor:
        losses = [state.auxiliary_loss for state in self.get_router_states(detach=detach)]
        if not losses:
            return next(self.parameters()).new_zeros(())
        stacked = torch.stack(losses)
        if reduction == "none":
            return stacked
        if reduction == "sum":
            return stacked.sum()
        if reduction == "mean":
            return stacked.mean()
        raise ValueError("reduction must be 'mean', 'sum', or 'none'")

    def clear_router_state(self) -> None:
        for wrapper in self.fusion_layers():
            wrapper.last_router_state = None
            wrapper.last_router_diagnostics = None

    def set_coding_enabled(self, enabled: bool = True) -> "FuseGlmForCausalLM":
        """Enable or bypass every coding-expert branch."""

        for wrapper in self.fusion_layers():
            wrapper.coding_enabled = bool(enabled)
        return self

    @property
    def coding_enabled(self) -> bool:
        layers = self.fusion_layers()
        return bool(layers) and all(wrapper.coding_enabled for wrapper in layers)

    @contextmanager
    def coding_experts(self, enabled: bool = True) -> Iterator["FuseGlmForCausalLM"]:
        """Temporarily enable/disable coding experts for one local operation."""

        wrappers = self.fusion_layers()
        previous = tuple(wrapper.coding_enabled for wrapper in wrappers)
        self.set_coding_enabled(enabled)
        try:
            yield self
        finally:
            for wrapper, old_value in zip(wrappers, previous):
                wrapper.coding_enabled = old_value

    @torch.no_grad()
    def load_folded_expert(
        self,
        layer_index: int,
        expert_index: int,
        *,
        gate_proj: torch.Tensor,
        up_proj: torch.Tensor,
        down_proj: torch.Tensor,
    ) -> None:
        """Copy one folded expert into a concrete decoder/expert slot."""

        if layer_index < 0 or layer_index >= len(self.model.layers):
            raise IndexError(f"layer_index {layer_index} is outside the decoder")
        wrapper = self.model.layers[layer_index].feed_forward
        if not isinstance(wrapper, FuseGlmFeedForward):
            raise ValueError(f"decoder layer {layer_index} has no fused expert branch")
        if expert_index < 0 or expert_index >= len(wrapper.experts):
            raise IndexError(f"expert_index {expert_index} is outside layer {layer_index}")
        wrapper.experts[expert_index].load_folded_weights(
            gate_proj=gate_proj,
            up_proj=up_proj,
            down_proj=down_proj,
        )

    @torch.no_grad()
    def load_router_initializer(
        self,
        layer_index: int,
        *,
        proj_weight: torch.Tensor,
        e_score_correction_bias: torch.Tensor,
    ) -> None:
        """Load one folded GLM router while preserving post-sigmoid bias semantics."""

        wrappers = {wrapper.layer_index: wrapper for wrapper in self.fusion_layers()}
        if layer_index not in wrappers:
            raise KeyError(f"decoder layer {layer_index} has no fused expert branch")
        router = wrappers[layer_index].router
        expected_weight_shape = tuple(router.proj.weight.shape)
        expected_bias_shape = tuple(router.e_score_correction_bias.shape)
        if tuple(proj_weight.shape) != expected_weight_shape:
            raise ValueError(
                f"router proj_weight has shape {tuple(proj_weight.shape)}; "
                f"expected {expected_weight_shape}"
            )
        if tuple(e_score_correction_bias.shape) != expected_bias_shape:
            raise ValueError(
                "router e_score_correction_bias has shape "
                f"{tuple(e_score_correction_bias.shape)}; expected {expected_bias_shape}"
            )
        if not bool(torch.isfinite(proj_weight).all()):
            raise ValueError("router proj_weight contains NaN or infinity")
        if not bool(torch.isfinite(e_score_correction_bias).all()):
            raise ValueError("router e_score_correction_bias contains NaN or infinity")
        router.proj.weight.copy_(
            proj_weight.to(device=router.proj.weight.device, dtype=router.proj.weight.dtype)
        )
        router.e_score_correction_bias.copy_(
            e_score_correction_bias.to(
                device=router.e_score_correction_bias.device,
                dtype=router.e_score_correction_bias.dtype,
            )
        )

    @classmethod
    def from_lfm_pretrained(
        cls,
        pretrained_model_name_or_path: str,
        *model_args: Any,
        config: Lfm2Config | FuseGlmConfig | None = None,
        fuse_overrides: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> "FuseGlmForCausalLM":
        """Load native LFM weights and initialize only the fusion branch.

        The method accepts local directories and normal Hugging Face loading
        arguments.  It does not set ``trust_remote_code`` or force a network
        lookup.  Use ``local_files_only=True`` when an offline-only guarantee is
        desired.
        """

        overrides = dict(fuse_overrides or {})
        if config is None:
            config_keys = (
                "cache_dir",
                "force_download",
                "local_files_only",
                "revision",
                "subfolder",
                "token",
            )
            config_kwargs = {key: kwargs[key] for key in config_keys if key in kwargs}
            config = Lfm2Config.from_pretrained(pretrained_model_name_or_path, **config_kwargs)
        if not isinstance(config, FuseGlmConfig):
            config = FuseGlmConfig.from_lfm_config(config, **overrides)
        elif overrides:
            config = FuseGlmConfig.from_lfm_config(config, **overrides)
        return cls.from_pretrained(
            pretrained_model_name_or_path,
            *model_args,
            config=config,
            **kwargs,
        )

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Any | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        labels: torch.LongTensor | None = None,
        use_cache: bool | None = None,
        logits_to_keep: int | torch.Tensor = 0,
        output_router_diagnostics: bool | None = None,
        coding_enabled: bool | None = None,
        **kwargs: Any,
    ) -> FuseGlmCausalLMOutputWithPast | tuple[torch.Tensor, ...]:
        """Run the native LFM forward plus the configured coding sidecars.

        Passing ``coding_enabled`` provides a convenient per-call override.
        For concurrent callers, prefer separate model instances because this
        override temporarily changes local module flags.
        """

        wrappers = self.fusion_layers()
        previous = tuple(wrapper.coding_enabled for wrapper in wrappers)
        if coding_enabled is not None:
            self.set_coding_enabled(coding_enabled)
        try:
            outputs = super().forward(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                inputs_embeds=inputs_embeds,
                labels=labels,
                use_cache=use_cache,
                logits_to_keep=logits_to_keep,
                **kwargs,
            )
        finally:
            if coding_enabled is not None:
                for wrapper, old_value in zip(wrappers, previous):
                    wrapper.coding_enabled = old_value

        # ``return_dict=False`` is retained for compatibility with upstream.
        # Router information remains accessible through the getter methods.
        if not isinstance(outputs, CausalLMOutputWithPast):
            return outputs

        router_aux_loss = self.get_router_aux_loss(reduction="mean")
        loss = outputs.loss
        if loss is not None and self.training and self.config.fuse_glm_router_aux_loss_coef > 0:
            # Gradients are attached inside each wrapper (checkpoint-safe).
            # This detached term makes the scalar reported to users equal to
            # CE + coefficient * mean(auxiliary loss) without double-counting.
            loss = loss + self.config.fuse_glm_router_aux_loss_coef * router_aux_loss.detach()

        include_diagnostics = (
            self.config.fuse_glm_output_router_diagnostics
            if output_router_diagnostics is None
            else bool(output_router_diagnostics)
        )
        diagnostics = self.get_router_diagnostics() if include_diagnostics else None
        return FuseGlmCausalLMOutputWithPast(
            loss=loss,
            logits=outputs.logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            router_aux_loss=router_aux_loss,
            router_diagnostics=diagnostics,
        )


__all__ = [
    "FoldedGlmExpert",
    "FuseGlmCausalLMOutputWithPast",
    "FuseGlmFeedForward",
    "FuseGlmForCausalLM",
    "RouterDiagnostics",
    "RouterState",
    "TopKFoldedExpertRouter",
]