# SPDX-FileCopyrightText: 2026 ModelCloud.ai # SPDX-FileCopyrightText: 2026 qubitium@modelcloud.ai # SPDX-License-Identifier: Apache-2.0 # Contact: qubitium@modelcloud.ai, x.com/qubitium # Adapted from intel/auto-round # at https://github.com/intel/auto-round/blob/main/auto_round/modeling/fused_moe/moe_experts_interface.py """ Custom experts implementation for transformers' MOE integration. This module provides a `linear_loop` experts implementation that uses individual nn.Linear layers per expert instead of fused 3D Parameters. This enables proper quantization of MOE expert weights. The implementation integrates with transformers' `use_experts_implementation` decorator and `ALL_EXPERTS_FUNCTIONS` registry. Usage: from auto_round.modeling.fused_moe.moe_experts_interface import prepare_model_for_moe_quantization # Before quantization prepare_model_for_moe_quantization(model) # Now the model uses linear_loop forward which supports quantized nn.Linear layers """ from types import MethodType import torch from logbar import LogBar from torch import nn from defuser.model_registry import MODEL_CONFIG, PATCH from defuser.utils.common import compile_module_name_filter, matches_module_name_filter from defuser.utils.device import clear_memory, to_meta from defuser import DEBUG_ON logger = LogBar(__name__) try: from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS HAS_EXPERTS_INTERFACE = True except ImportError: HAS_EXPERTS_INTERFACE = False ALL_EXPERTS_FUNCTIONS = None # Expert implementation name - change this if transformers want to use a different name LINEAR_LOOP_IMPL = "linear_loop" BATCHED_INPUT_IMPL = "batched_input" # Known expert projection patterns for reference # These are used as hints when auto-detection needs to infer projection properties # Format: proj_name -> {"is_input_proj": bool, "output_multiplier": int} # is_input_proj: True if takes hidden_dim as input, False if takes intermediate_dim # output_multiplier: output dimension multiplier (e.g., 2 for fused gate+up projection) KNOWN_PROJECTION_PATTERNS = { # Transformers 5.0+ standard (Qwen3-MoE, etc.) # gate_up_proj is auto-split into gate_proj + up_proj during unfusing "gate_up_proj": {"is_input_proj": True, "output_multiplier": 2, "split_into": ["gate_proj", "up_proj"]}, "gate_proj": {"is_input_proj": True, "output_multiplier": 1}, # hidden -> intermediate (gate) "up_proj": {"is_input_proj": True, "output_multiplier": 1}, # hidden -> intermediate (up) "down_proj": {"is_input_proj": False, "output_multiplier": 1}, # intermediate -> hidden } class _ExpertContainer(nn.Module): """Lightweight container for a single expert's projection layers. Each expert has its projections (e.g., gate_proj, up_proj, down_proj) as direct attributes. When attached as numbered children of the experts module (e.g., module.add_module("0", container)), PyTorch naturally produces state_dict keys like: {prefix}0.gate_proj.weight which matches the standard checkpoint format without any hooks. """ pass def is_linear_loop_available() -> bool: """Check if linear_loop experts implementation is available.""" return HAS_EXPERTS_INTERFACE def _apply_expert_gate( module: nn.Module, gate_out: torch.Tensor | None, up_out: torch.Tensor, ) -> torch.Tensor: """Apply the expert activation path for gated and non-gated expert MLPs.""" if gate_out is None: act_fn = getattr(module, "act_fn", None) if act_fn is None: raise AttributeError(f"{module.__class__.__name__} must define `act_fn` for non-gated experts.") return act_fn(up_out) if hasattr(module, "_apply_gate"): return module._apply_gate(torch.cat([gate_out, up_out], dim=-1)) act_fn = getattr(module, "act_fn", None) if act_fn is None: raise AttributeError(f"{module.__class__.__name__} must define either `_apply_gate` or `act_fn`.") return act_fn(gate_out) * up_out def linear_loop_experts_forward( self: nn.Module, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor, ) -> torch.Tensor: """Forward using individual nn.Linear layers per expert. This implementation loops over experts and accesses per-expert containers (self._modules["0"], self._modules["1"], ...) each with gate_proj, up_proj, and down_proj as nn.Linear layers (or quantized equivalents), enabling proper quantization support. Expected module structure: - Numbered children (0, 1, ..., num_experts-1), each an _ExpertContainer with: - gate_proj: nn.Linear (in_features=hidden_dim, out_features=intermediate_dim) - up_proj: nn.Linear (in_features=hidden_dim, out_features=intermediate_dim) - down_proj: nn.Linear (in_features=intermediate_dim, out_features=hidden_dim) - act_fn: activation function - num_experts: number of experts - _apply_gate: optional custom gating function Args: self: The experts module hidden_states: Input tensor of shape (num_tokens, hidden_dim) top_k_index: Selected expert indices of shape (num_tokens, top_k) top_k_weights: Expert weights of shape (num_tokens, top_k) Returns: final_hidden_states: Output tensor of shape (num_tokens, hidden_dim) """ if DEBUG_ON: logger.debug(f"Using {LINEAR_LOOP_IMPL} experts forward for {self.__class__.__name__}") # Handle [batch_size, seq_len, hidden_dim] input format if hidden_states.dim() == 3: batch_size, seq_len, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) # [bs * seq_len, hidden_dim] top_k_index = top_k_index.view(-1, top_k_index.size(-1)) # [bs * seq_len, top_k] top_k_weights = top_k_weights.view(-1, top_k_weights.size(-1)) # [bs * seq_len, top_k] else: batch_size, seq_len = None, None hidden_dim = hidden_states.size(-1) device = hidden_states.device num_top_k = top_k_index.size(-1) num_tokens = hidden_states.size(0) num_experts = self.num_experts # Reshape for easier indexing # S is the number of selected token-expert pairs (S = num_tokens * num_top_k) token_idx = torch.arange(num_tokens, device=device).unsqueeze(1).expand(-1, num_top_k).reshape(-1) # (S,) sample_weights = top_k_weights.reshape(-1).to(hidden_states.dtype) # (S,) expert_ids = top_k_index.reshape(-1) # (S,) # Get current hidden states for selected samples selected_hidden_states = hidden_states[token_idx] # (S, hidden_dim) # Allocate output tensor out_per_sample = torch.zeros(token_idx.size(0), hidden_dim, device=device, dtype=hidden_states.dtype) # Process each expert for expert_idx in range(num_experts): # Find samples routed to this expert mask = expert_ids == expert_idx if not mask.any(): continue expert_input = selected_hidden_states[mask] # (num_samples_for_expert, hidden_dim) # Get this expert's container with its projection layers expert = getattr(self, str(expert_idx)) if hasattr(expert, "gate_proj"): gate_out = expert.gate_proj(expert_input) # (num_samples, intermediate_dim) up_out = expert.up_proj(expert_input) # (num_samples, intermediate_dim) else: gate_out = None up_out = expert.up_proj(expert_input) gated_out = _apply_expert_gate(self, gate_out, up_out) # Down projection expert_out = expert.down_proj(gated_out) # (num_samples, hidden_dim) # Store results out_per_sample[mask] = expert_out.to(out_per_sample.dtype) # Apply routing weights out_per_sample = out_per_sample * sample_weights.unsqueeze(-1) # (S, hidden_dim) # Accumulate results using deterministic reshape+sum instead of index_add_ # (index_add_ with duplicate indices is non-deterministic on CUDA due to atomicAdd) final_hidden_states = out_per_sample.view(num_tokens, num_top_k, hidden_dim).sum(dim=1) # Reshape back to original format if input was [batch_size, seq_len, hidden_dim] if batch_size is not None: final_hidden_states = final_hidden_states.view(batch_size, seq_len, hidden_dim) return final_hidden_states def batched_input_experts_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: """Run defused experts for models that feed experts as expert-major input batches. Llama4 is the current example: upstream code repeats and pre-weights tokens, then calls ``experts(hidden_states)`` where the leading dimension is laid out as ``[expert0_tokens, expert1_tokens, ...]``. This forward keeps that public contract while still executing per-expert ``nn.Linear`` modules internally. """ if DEBUG_ON: logger.debug(f"Using {BATCHED_INPUT_IMPL} experts forward for {self.__class__.__name__}") hidden_dim = hidden_states.size(-1) expert_inputs = hidden_states.view(self.num_experts, -1, hidden_dim) expert_outputs = [] for expert_idx in range(self.num_experts): expert = getattr(self, str(expert_idx)) expert_input = expert_inputs[expert_idx] gate_out = expert.gate_proj(expert_input) up_out = expert.up_proj(expert_input) gated_out = _apply_expert_gate(self, gate_out, up_out) expert_outputs.append(expert.down_proj(gated_out)) return torch.stack(expert_outputs, dim=0).reshape(-1, hidden_dim) def register_linear_loop_experts() -> bool: """Register the linear_loop experts implementation with transformers. Returns: True if registration was successful, False otherwise. """ if not HAS_EXPERTS_INTERFACE: logger.warn( "transformers.integrations.moe.ALL_EXPERTS_FUNCTIONS not available. " "linear_loop experts implementation not registered. " "Requires transformers >= 5.0.0" ) return False if LINEAR_LOOP_IMPL not in ALL_EXPERTS_FUNCTIONS._global_mapping: ALL_EXPERTS_FUNCTIONS._global_mapping[LINEAR_LOOP_IMPL] = linear_loop_experts_forward if DEBUG_ON: logger.debug(f"Registered '{LINEAR_LOOP_IMPL}' experts implementation") return True def _model_experts_defuse_specs(model: nn.Module) -> list[dict]: """Return declarative experts-defusion specs for the current model type.""" config = getattr(model, "config", None) model_type = getattr(config, "model_type", None) if model_type is None: return [] specs = MODEL_CONFIG.get(model_type, {}).get(PATCH.EXPERTS_DEFUSE, []) if isinstance(specs, dict): return [specs] return list(specs) def _module_class_path(module: nn.Module) -> str: """Return a stable import-style class path for matching model specs.""" return f"{module.__class__.__module__}.{module.__class__.__name__}" def _matching_experts_defuse_spec(module: nn.Module, specs: list[dict]) -> dict | None: """Find the first declarative experts-defusion spec that matches ``module``.""" module_path = _module_class_path(module) module_name = module.__class__.__name__ for spec in specs: target = spec.get("module_class") if target in {module_path, module_name}: return spec return None def _install_instance_forward(module: nn.Module, implementation: str) -> None: """Attach a generic forward implementation directly to one experts module.""" if implementation == LINEAR_LOOP_IMPL: module.forward = MethodType(linear_loop_experts_forward, module) return if implementation == BATCHED_INPUT_IMPL: module.forward = MethodType(batched_input_experts_forward, module) return raise ValueError(f"Unsupported experts forward implementation: {implementation}") def _detect_expert_projections(module: nn.Module) -> dict[str, dict]: """Detect which expert projections exist in the module. This function scans the module for any registered 3D Parameter / Tensor attributes. It first checks known projection names, then discovers any unknown 3D registered tensors. Returns: Dict mapping projection names to their config, only for projections that exist as 3D registered Parameter / Tensor in the module. """ detected = {} # Only inspect registered tensors here. Scanning arbitrary attributes can # trigger unrelated properties such as Transformers' `loss_function`. local_tensors = { name: tensor for registry in (module._parameters, module._buffers) for name, tensor in registry.items() if isinstance(tensor, torch.Tensor) } # First, check known projection patterns for proj_name, config in KNOWN_PROJECTION_PATTERNS.items(): param = local_tensors.get(proj_name) if param is not None and param.dim() == 3: detected[proj_name] = config # If no known patterns found, scan for any 3D registered tensor (future-proofing) if not detected: for attr_name, param in local_tensors.items(): if attr_name.startswith("_"): continue if param is not None and param.dim() == 3: # Use default config for unknown projections if DEBUG_ON: logger.debug(f"Discovered unknown 3D projection: {attr_name}") detected[attr_name] = {"is_input_proj": True, "output_multiplier": 1} return detected def _get_registered_tensor(module: nn.Module, name: str) -> torch.Tensor | None: """Return a direct registered parameter or buffer without touching properties.""" tensor = module._parameters.get(name) if isinstance(tensor, torch.Tensor): return tensor tensor = module._buffers.get(name) if isinstance(tensor, torch.Tensor): return tensor return None def _set_registered_tensor_like(module: nn.Module, name: str, tensor: torch.Tensor, source: torch.Tensor) -> None: """Register ``tensor`` using the same parameter-vs-buffer kind as ``source``.""" if isinstance(source, nn.Parameter): module.register_parameter(name, nn.Parameter(tensor, requires_grad=source.requires_grad)) else: module.register_buffer(name, tensor) def _experts_supports_decorator(module: nn.Module) -> bool: """Check if experts module supports @use_experts_implementation decorator. Only experts classes decorated with @use_experts_implementation will use our linear_loop forward. Others need full module replacement. """ forward_method = getattr(module.__class__, "forward", None) if forward_method is None: return False # @use_experts_implementation sets __wrapped__ on the decorated method return hasattr(forward_method, "__wrapped__") def _infer_dimensions(param: torch.Tensor, config: dict, is_transposed: bool) -> tuple[int, int]: """Infer input and output dimensions for a projection. Args: param: The 3D projection tensor (num_experts, dim1, dim2) config: Projection config with is_input_proj and output_multiplier is_transposed: Whether weights are stored transposed Returns: (in_features, out_features) for the Linear layer """ dim1, dim2 = param.shape[1], param.shape[2] multiplier = config.get("output_multiplier", 1) if is_transposed: # transposed: (num_experts, in_features, out_features) in_features, out_features = dim1, dim2 else: # not transposed: (num_experts, out_features, in_features) out_features, in_features = dim1, dim2 # Adjust for multiplier (e.g., gate_up has 2x intermediate) if multiplier > 1: out_features = out_features // multiplier * multiplier # ensure divisible return in_features, out_features def _unfuse_single_projection( module: nn.Module, proj_name: str, num_experts: int, is_transposed: bool, dtype: torch.dtype, target_device: torch.device, ) -> list | None: """Unfuse a single projection from a 3D registered tensor to Linear layers. Optimized to keep peak device memory low while preserving the module's original device placement: - Moves the full 3D tensor to CPU in a single transfer - Performs batch transpose on CPU if needed - Releases the original fused parameter before allocating defused linears - Re-materializes each expert linear back onto ``target_device`` Args: module: The experts module proj_name: Name of the projection attribute num_experts: Number of experts is_transposed: Whether weights are stored transposed dtype: Data type for the Linear layers target_device: Device for the Linear layers Returns: List of Linear layers, or None if projection doesn't exist """ param = _get_registered_tensor(module, proj_name) if param is None or param.dim() != 3: return None # Get projection config config = KNOWN_PROJECTION_PATTERNS.get(proj_name, {"is_input_proj": True, "output_multiplier": 1}) # Infer dimensions in_features, out_features = _infer_dimensions(param, config, is_transposed) # Check for bias bias_name = f"{proj_name}_bias" bias_param = _get_registered_tensor(module, bias_name) has_bias = isinstance(bias_param, torch.Tensor) source_device = param.device is_meta = source_device.type == "meta" weight_requires_grad = param.requires_grad if isinstance(param, nn.Parameter) else False # Prepare weight slices on CPU in batch (single D2H transfer + batch transpose) if not is_meta: # Single transfer: GPU -> CPU (or no-op if already on CPU) weights_cpu = param.detach().cpu() # (num_experts, dim1, dim2) if is_transposed: # Batch transpose: (num_experts, in, out) -> (num_experts, out, in) weights_cpu = weights_cpu.transpose(1, 2) # Ensure contiguous layout so unbind produces contiguous 2D slices. # This is needed when the param comes from a chunk split (Phase 1) # which produces non-contiguous views even on CPU. if not weights_cpu.is_contiguous(): weights_cpu = weights_cpu.contiguous() # Unbind into a tuple of 2D tensors (zero-copy views since contiguous) weight_slices = weights_cpu.unbind(0) if has_bias: bias_cpu = bias_param.detach().cpu() # Ensure contiguous — bias may come from a chunk split (Phase 1) if not bias_cpu.is_contiguous(): bias_cpu = bias_cpu.contiguous() bias_slices = bias_cpu.unbind(0) bias_requires_grad = bias_param.requires_grad if isinstance(bias_param, nn.Parameter) else False # Drop the original fused parameter before allocating the defused # per-expert linears back on the original device. try: setattr(module, proj_name, to_meta(param)) param = None if has_bias: setattr(module, bias_name, to_meta(bias_param)) bias_param = None if DEBUG_ON: logger.debug(f"Released memory for {proj_name} using to_meta()") except Exception: pass # Create Linear shells on meta device (no memory allocation) linears = [] for i in range(num_experts): # meta device: creates the module structure without allocating weight storage linear = nn.Linear(in_features, out_features, bias=has_bias, dtype=dtype, device="meta") if not is_meta: weight = weight_slices[i] if target_device.type != "cpu": weight = weight.to(device=target_device, dtype=dtype) linear.weight = nn.Parameter(weight, requires_grad=weight_requires_grad) if has_bias: bias = bias_slices[i] if target_device.type != "cpu": bias = bias.to(device=target_device, dtype=bias.dtype) linear.bias = nn.Parameter(bias, requires_grad=bias_requires_grad) linears.append(linear) return linears def _install_compact_expert_repr(module: nn.Module) -> None: """Install compact __repr__ on the module's class. Collapses identical _ExpertContainer children into a range display, similar to nn.ModuleList:: GptOssExperts( (0-63): 64 x _ExpertContainer( (gate_proj): Linear(...) (up_proj): Linear(...) (down_proj): Linear(...) ) ) """ cls = module.__class__ if "_compact_expert_repr_installed" in cls.__dict__: return _original_repr = cls.__repr__ def _compact_repr(self): from torch.nn.modules.module import _addindent # Separate expert containers from other children expert_items = [] other_items = [] for key, child in self._modules.items(): if isinstance(child, _ExpertContainer): expert_items.append((key, child)) else: other_items.append((key, child)) if not expert_items: return _original_repr(self) lines = [] extra_repr = self.extra_repr() if extra_repr: lines.extend(extra_repr.split("\n")) # Print non-expert children normally for key, child in other_items: mod_str = repr(child) mod_str = _addindent(mod_str, 2) lines.append(f"({key}): {mod_str}") # Collapse expert containers into a single range line first_key = expert_items[0][0] last_key = expert_items[-1][0] sample_repr = repr(expert_items[0][1]) n = len(expert_items) sample_repr = _addindent(sample_repr, 2) if n > 1: lines.append(f"({first_key}-{last_key}): {n} x {sample_repr}") else: lines.append(f"({first_key}): {sample_repr}") main_str = self._get_name() + "(" if lines: main_str += "\n " + "\n ".join(lines) + "\n" main_str += ")" return main_str cls.__repr__ = _compact_repr cls._compact_expert_repr_installed = True # Make the module subscriptable and iterable: module[0], module[1], len(module), for e in module if not hasattr(cls, "__getitem__") or "__getitem__" not in cls.__dict__: def _getitem(self, idx): return getattr(self, str(idx)) def _len(self): return self.num_experts def _iter(self): for i in range(self.num_experts): yield getattr(self, str(i)) cls.__getitem__ = _getitem cls.__len__ = _len cls.__iter__ = _iter def _unfuse_experts_weights_inplace( module: nn.Module, check_decorator: bool = True, projection_names: list[str] | None = None, ) -> bool: """Convert fused 3D expert weights to per-expert containers of nn.Linear layers. This function modifies the module in-place, replacing fused 3D Parameters with numbered _ExpertContainer children, each holding nn.Linear projections. Args: module: The experts module to unfuse check_decorator: If True, only unfuse if the module supports @use_experts_implementation decorator. Default is True. projection_names: Optional list of projection names to unfuse. If None, auto-detects from KNOWN_PROJECTION_PATTERNS. Returns: True if unfusing was successful, False if module doesn't match pattern """ # Detect available projections if projection_names: detected_projections = { name: config for name, config in KNOWN_PROJECTION_PATTERNS.items() if name in projection_names } else: detected_projections = _detect_expert_projections(module) if not detected_projections: return False # Only unfuse if the module supports the decorator (unless check_decorator is False) if check_decorator and not _experts_supports_decorator(module): logger.debug(f"Skipping unfuse for {module.__class__.__name__}: does not support @use_experts_implementation") return False # Get first projection to determine num_experts and layout first_proj_name = next(iter(detected_projections)) first_param = _get_registered_tensor(module, first_proj_name) num_experts = first_param.shape[0] # Detect if transposed is_transposed = getattr(module, "is_transposed", None) if is_transposed is None: # Infer from shape: typically hidden_dim < intermediate_dim if ( first_proj_name in {"up_proj", "down_proj"} and "gate_up_proj" not in detected_projections and "gate_proj" not in detected_projections ): is_transposed = False else: dim1, dim2 = first_param.shape[1], first_param.shape[2] is_transposed = dim1 < dim2 dtype = first_param.dtype target_device = first_param.device if first_param.device.type != "meta" else "cpu" # Phase 1: Split fused projections (e.g., gate_up_proj -> gate_proj + up_proj) into separate 3D tensors extra_projections = {} fused_to_remove = [] for proj_name, config in detected_projections.items(): split_into = config.get("split_into") if not split_into: continue param = _get_registered_tensor(module, proj_name) if param is None or param.dim() != 3: continue # Split along output dimension split_dim = 2 if is_transposed else 1 split_params = param.chunk(len(split_into), dim=split_dim) # Also split bias if present (e.g., gate_up_proj_bias -> gate_proj_bias + up_proj_bias) bias_name = f"{proj_name}_bias" bias_param = _get_registered_tensor(module, bias_name) bias_splits = None if isinstance(bias_param, torch.Tensor) and bias_param.dim() == 2: bias_splits = bias_param.chunk(len(split_into), dim=1) for i, (split_name, split_param) in enumerate(zip(split_into, split_params)): # Avoid .contiguous() here — _unfuse_single_projection will handle it # during batch transpose/unbind, saving a full-tensor copy _set_registered_tensor_like(module, split_name, split_param, param) if bias_splits is not None: _set_registered_tensor_like(module, f"{split_name}_bias", bias_splits[i], bias_param) extra_projections[split_name] = KNOWN_PROJECTION_PATTERNS.get( split_name, {"is_input_proj": True, "output_multiplier": 1} ) delattr(module, proj_name) if bias_param is not None: delattr(module, bias_name) fused_to_remove.append(proj_name) if DEBUG_ON: logger.debug(f"Split {proj_name} -> {split_into}: {num_experts} experts") # Remove fused entries and add split entries for name in fused_to_remove: del detected_projections[name] detected_projections.update(extra_projections) # Phase 2: Unfuse all 3D tensors into per-expert Linear layers proj_linears = {} # proj_name -> [Linear_expert0, Linear_expert1, ...] for proj_name in detected_projections: linears = _unfuse_single_projection(module, proj_name, num_experts, is_transposed, dtype, target_device) if linears is not None: delattr(module, proj_name) # Also remove the bias parameter if it exists (already absorbed into Linear.bias) bias_name = f"{proj_name}_bias" if hasattr(module, bias_name): delattr(module, bias_name) proj_linears[proj_name] = linears if not proj_linears: return False # Phase 3: Assemble per-expert containers as direct numbered children. # This produces state_dict keys like: {prefix}0.gate_proj.weight # which naturally matches the standard checkpoint format — no hooks needed. for expert_idx in range(num_experts): container = _ExpertContainer() for proj_name, linears in proj_linears.items(): setattr(container, proj_name, linears[expert_idx]) module.add_module(str(expert_idx), container) # Ensure num_experts is set if not hasattr(module, "num_experts"): module.num_experts = num_experts # Mark as unfused for detection during save module._unfused_experts = True # Mark all unfused modules as already initialized so transformers' # _initialize_weights skips them. Without this, the model's _init_weights # tries to access the original fused attributes (e.g., gate_up_proj) which # no longer exist after unfusing. module._is_hf_initialized = True for child in module.modules(): child._is_hf_initialized = True # Install compact repr to collapse identical expert containers in print output _install_compact_expert_repr(module) return True def prepare_model_for_moe_quantization( model: nn.Module, implementation: str = LINEAR_LOOP_IMPL, filter_rules=None, ) -> list[str]: """Prepare a model for MOE quantization using transformers' experts interface. This function: 1. Registers the linear_loop experts implementation with transformers 2. Sets model.config._experts_implementation = implementation 3. Unfuses all fused MOE expert weights into per-expert containers of nn.Linear After calling this function, the model's forward pass will use individual nn.Linear layers per expert, which can be quantized normally. Args: model: The model to prepare implementation: The experts implementation to use (default: "linear_loop") Returns: List of module names that were unfused """ if not register_linear_loop_experts(): raise RuntimeError( "Failed to register linear_loop experts implementation. " "This requires transformers >= 5.0.0 with MOE integration support." ) # Unfuse all fused experts modules, including models that need a generic # instance-level forward override instead of transformers' decorator path. unfused_modules = [] decorated_unfused_modules = [] experts_defuse_specs = _model_experts_defuse_specs(model) module_name_filter = compile_module_name_filter(filter_rules) for name, module in model.named_modules(): if not matches_module_name_filter(name, module_name_filter): continue spec = _matching_experts_defuse_spec(module, experts_defuse_specs) if spec is not None and _unfuse_experts_weights_inplace( module, check_decorator=False, projection_names=spec.get("projection_names"), ): _install_instance_forward(module, spec["forward_impl"]) unfused_modules.append(name) if DEBUG_ON: logger.debug(f"[MoE Prep] Unfused '{name}' via declarative spec") continue if _unfuse_experts_weights_inplace(module): unfused_modules.append(name) decorated_unfused_modules.append(name) if DEBUG_ON: logger.debug(f"[MoE Prep] Unfused '{name}'") # Only set config if we actually unfused something # Models that don't support the decorator (like Llama4) won't have anything unfused # and should use full module replacement instead if unfused_modules: if DEBUG_ON: logger.info(f"[MoE Prep] Unfused {len(unfused_modules)} MOE experts modules") clear_memory() # Set config for linear_loop forward only when the upstream model uses # the decorator-based experts interface. if decorated_unfused_modules and hasattr(model, "config"): saved_impl = getattr(model.config, "experts_implementation", None) impl_to_set = saved_impl if saved_impl else implementation model.config._experts_implementation = impl_to_set if DEBUG_ON: logger.debug(f"Set model.config._experts_implementation = '{impl_to_set}'") return unfused_modules