Introduce standardized BatchLinear module for MoE architectures to facilitate PEFT and Quantization
Author: ITcarrot
Pull request: no
Labels: name: Feature request

### Feature request

Starting from transformer v5, Linear module of Moe experts are fused into one single module. While this improve speed of Moe, it is making downstream library difficult to adapt to this change.

I propose introducing a standardized `BatchLinear` (or `MoELinear`) module within transformers. This module would serve as the standard abstraction for expert layers, encapsulating the raw weights while providing a clean `nn.Module` interface.

This affects `QwenMoeExperts` and other Moe models in transformers.

API definition of proposed module:
```python
class BatchLinear(nn.Module):
    def __init__(self, num_experts, in_features, out_features, bias=False):
        super().__init__()
        self.num_experts = num_experts
        self.in_features = in_features
        self.out_features = out_features
        
        # Standardized weight storage
        self.weight = nn.Parameter(torch.empty(num_experts, out_features, in_features))
        if bias:
            self.bias = nn.Parameter(torch.empty(num_experts, out_features))
        else:
            self.register_parameter('bias', None)

    def forward(self, expert_id: int, input: torch.Tensor) -> torch.Tensor:
        """
        Applies the linear transformation for the specified expert.
        
        Args:
            expert_id (int): The index of the expert whose parameters to use.
            input (torch.Tensor): The input tensor of shape (batch_size, in_features).

        Returns:
            torch.Tensor: The output tensor of shape (batch_size, out_features).
        """
        return nn.functional.linear(input, self.weight[expert_id], self.bias[expert_id] if self.bias is not None else None)

    def grouped_forward(self, grouped_input: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
        """
        Applies the linear transformation for a batch of inputs, each group associated with a different expert.
        
        Args:
            grouped_input (torch.Tensor): The input tensor of shape (total_batch_size, in_features).
            offsets (torch.Tensor): A tensor of shape (num_experts,) indicating the start and end indices for each expert's inputs.
                e.g., grouped_input[offsets[i]:offsets[i+1]] corresponds input to expert i.

        Returns:
            torch.Tensor: The output tensor of shape (total_batch_size, out_features).
        """
        pass
```

### Motivation

1.  **Seamless PEFT Support:**
    If `transformers` models use `BatchLinear`, PEFT can simply implement a `LoraBatchLinear` layer. When converting a model, PEFT can detect `BatchLinear` and swap it out, enabling LoRA finetuning on experts immediately without model-specific hacks. [Related issue in peft](/huggingface/peft/issues/3009)

2.  **Seamless quantization Support:**
     Quantization library can also replace this module easily to implement its work. Related to #43284 , per-expert quantization is also supportted under `BatchLinear.forward`.



### Your contribution

I'd like to help implement it if the community accept this idea of design.

Comments:
Hi @ITcarrot — thanks for the idea! One question: it looks like this approach still concatenates all experts at initialization.
Could you clarify whether this method helps with the RAM concern here? https://github.com/huggingface/transformers/issues/43284#issuecomment-3782649821

@yiliu30 What's happening to the RAM usage in your case is likely because quantization library is NOT quantizing experts AT ALL. I share some more details back in https://github.com/huggingface/transformers/issues/43284#issuecomment-3796437061

Hi @ITcarrot ,

I’ve spent some time digging into the MoE block implementations for Mixtral and Qwen2Moe. I think standardizing this into a BatchLinear module is exactly what's needed for better PEFT/Quantization integration.

I’ve started drafting the implementation using a vectorized einsum approach for the grouped_forward pass to ensure it’s framework-native and efficient. My plan is to:

Implement the standardized BatchLinear module with support for both individual and grouped expert inference.

Provide a numerical parity test suite (CPU-compatible) to verify the new module against current loop-based implementations.

Submit a PR shortly with the core module and a refactor for Qwen2Moe.

I'll be relying on the CI for the final GPU/CUDA validation, but I'll have the logic and CPU tests ready to go.