"""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 def forward_for_serving(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Route without training-only softmax, one-hot, or auxiliary loss work.""" logits = F.linear(hidden_states.float(), self.proj.weight.float()) scores = torch.sigmoid(logits) selected_indices = torch.topk( scores + self.e_score_correction_bias.float(), 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) return selected_indices, selected_weights 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 self.fast_expert_bank: nn.Module | None = None self.serving_mode = False @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: if self.fast_expert_bank is not None: return self.fast_expert_bank(hidden_states, selected_indices, selected_weights) 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 _forward_for_serving( self, hidden_states: torch.Tensor, base_output: torch.Tensor ) -> torch.Tensor: """Inference fast path without router state or diagnostic construction.""" 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(), ) ).to(hidden_states.dtype) selected_indices, selected_weights = self.router.forward_for_serving(flat_hidden) expert_delta = self._dispatch(flat_hidden, selected_indices, selected_weights) expert_delta = expert_delta * token_gate return base_output + self.residual_scale.to(expert_delta.dtype) * expert_delta.reshape( original_shape ) 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 if self.serving_mode and not self.training: return self._forward_for_serving(hidden_states, 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 def enable_fast_fp8_serving(self) -> dict[str, Any]: """Pack TorchAO experts into the Triton grouped-FP8 serving runtime.""" from .fast_fp8_runtime import install_fast_fp8_runtime self.eval() return install_fast_fp8_runtime(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", ]