configuration_fuse_glm.py

"""Configuration for the LFM2 + folded GLM expert model.

The defaults describe the intended production architecture: sixteen folded
GLM experts in every LFM2 layer, with four experts selected per token.
All dimensions remain configurable so the implementation can be exercised
with very small, synthetic models without downloading either checkpoint.
"""

from __future__ import annotations

from collections.abc import Mapping
from typing import Any

from transformers import Lfm2Config


class FuseGlmConfig(Lfm2Config):
    """Extend :class:`~transformers.Lfm2Config` with sparse GLM sidecars.

    Parameters prefixed with ``fuse_glm_`` only control the added expert
    branch.  All ordinary LFM2 configuration arguments are forwarded to
    ``Lfm2Config`` unchanged.

    ``fuse_glm_layer_indices=None`` means that every decoder layer receives a
    sidecar.  A concrete list can be supplied for ablations or staged builds.
    The production model uses hidden/intermediate size 2048; tests may select
    smaller values.
    """

    model_type = "fuse_glm"

    def __init__(
        self,
        *,
        fuse_glm_num_experts: int = 16,
        fuse_glm_top_k: int = 4,
        fuse_glm_expert_intermediate_size: int = 2048,
        fuse_glm_layer_indices: list[int] | tuple[int, ...] | None = None,
        fuse_glm_gate_clamp_max: float = 10.0,
        fuse_glm_up_clamp_min: float = -10.0,
        fuse_glm_up_clamp_max: float = 10.0,
        fuse_glm_residual_scale_max: float = 0.1,
        fuse_glm_router_aux_loss_coef: float = 0.01,
        fuse_glm_token_gate_bias: float = -4.0,
        fuse_glm_token_gate_threshold: float = 0.5,
        fuse_glm_hard_token_gate_at_eval: bool = False,
        fuse_glm_coding_enabled: bool = True,
        fuse_glm_output_router_diagnostics: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)

        # Transformers 5.16 remaps legacy attention names for configurations
        # it classifies as custom code.  LFM2's implementation still indexes
        # masks with the literal key ``conv``, however, so a FuseGlmConfig
        # must retain the host model's native spelling.  Without this reversal
        # a real LFM checkpoint fails in Lfm2Model.forward with
        # ``KeyError: 'linear_attention'`` even though an Lfm2Config loaded
        # from the same file works correctly.
        self.layer_types = [
            "conv" if layer_type == "linear_attention" else layer_type
            for layer_type in self.layer_types
        ]

        self.fuse_glm_num_experts = int(fuse_glm_num_experts)
        self.fuse_glm_top_k = int(fuse_glm_top_k)
        self.fuse_glm_expert_intermediate_size = int(fuse_glm_expert_intermediate_size)
        self.fuse_glm_layer_indices = (
            None if fuse_glm_layer_indices is None else [int(index) for index in fuse_glm_layer_indices]
        )
        self.fuse_glm_gate_clamp_max = float(fuse_glm_gate_clamp_max)
        self.fuse_glm_up_clamp_min = float(fuse_glm_up_clamp_min)
        self.fuse_glm_up_clamp_max = float(fuse_glm_up_clamp_max)
        self.fuse_glm_residual_scale_max = float(fuse_glm_residual_scale_max)
        self.fuse_glm_router_aux_loss_coef = float(fuse_glm_router_aux_loss_coef)
        self.fuse_glm_token_gate_bias = float(fuse_glm_token_gate_bias)
        self.fuse_glm_token_gate_threshold = float(fuse_glm_token_gate_threshold)
        self.fuse_glm_hard_token_gate_at_eval = bool(fuse_glm_hard_token_gate_at_eval)
        self.fuse_glm_coding_enabled = bool(fuse_glm_coding_enabled)
        self.fuse_glm_output_router_diagnostics = bool(fuse_glm_output_router_diagnostics)

        self._validate_fuse_glm_fields()

        # Saved checkpoints should resolve to the fused class rather than the
        # base LFM2 class.  This does not affect loading an ordinary LFM config
        # through ``from_lfm_config`` below.
        self.architectures = ["FuseGlmForCausalLM"]

    @property
    def resolved_fuse_glm_layer_indices(self) -> tuple[int, ...]:
        """Return the validated decoder layer indices receiving experts."""

        if self.fuse_glm_layer_indices is None:
            return tuple(range(self.num_hidden_layers))
        return tuple(self.fuse_glm_layer_indices)

    def _validate_fuse_glm_fields(self) -> None:
        if self.fuse_glm_num_experts < 1:
            raise ValueError("fuse_glm_num_experts must be at least 1")
        if not 1 <= self.fuse_glm_top_k <= self.fuse_glm_num_experts:
            raise ValueError("fuse_glm_top_k must be between 1 and fuse_glm_num_experts")
        if self.fuse_glm_expert_intermediate_size < 1:
            raise ValueError("fuse_glm_expert_intermediate_size must be at least 1")
        if self.fuse_glm_up_clamp_min >= self.fuse_glm_up_clamp_max:
            raise ValueError("fuse_glm_up_clamp_min must be smaller than fuse_glm_up_clamp_max")
        if self.fuse_glm_residual_scale_max <= 0:
            raise ValueError("fuse_glm_residual_scale_max must be positive")
        if self.fuse_glm_router_aux_loss_coef < 0:
            raise ValueError("fuse_glm_router_aux_loss_coef cannot be negative")
        if not 0.0 <= self.fuse_glm_token_gate_threshold <= 1.0:
            raise ValueError("fuse_glm_token_gate_threshold must be in [0, 1]")

        indices = self.resolved_fuse_glm_layer_indices
        if len(set(indices)) != len(indices):
            raise ValueError("fuse_glm_layer_indices cannot contain duplicates")
        invalid = [index for index in indices if index < 0 or index >= self.num_hidden_layers]
        if invalid:
            raise ValueError(
                "fuse_glm_layer_indices contains indices outside the LFM2 decoder: "
                f"{invalid} (num_hidden_layers={self.num_hidden_layers})"
            )

    @classmethod
    def from_lfm_config(
        cls,
        config: Lfm2Config | Mapping[str, Any],
        **fuse_overrides: Any,
    ) -> "FuseGlmConfig":
        """Create a fused config from an existing LFM2 config.

        This method is purely local and never resolves or downloads a model.
        ``config`` may be an instantiated ``Lfm2Config`` or its dictionary.
        """

        if isinstance(config, Lfm2Config):
            config_dict = config.to_dict()
        elif isinstance(config, Mapping):
            config_dict = dict(config)
        else:
            raise TypeError("config must be an Lfm2Config or a mapping")

        # These values describe the source class/revision, not constructor
        # fields of the new local class.
        for key in ("model_type", "_commit_hash"):
            config_dict.pop(key, None)
        config_dict.pop("architectures", None)
        config_dict.update(fuse_overrides)
        return cls(**config_dict)


__all__ = ["FuseGlmConfig"]