Repository: huggingface/transformers
File: src/transformers/integrations/moe.py (Python)

Code before:

def _grouped_linear(
    input: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None = None,
    offs: torch.Tensor | None = None,
    is_transposed: bool = False,
) -> torch.Tensor:
    """Grouped linear layer supporting optional bias and transposed weights.

    Args:
        input (`torch.Tensor`):
            Input tensor of shape (S, input_dim).
        weight (`torch.Tensor`):
            Weight tensor of shape (num_experts, output_dim, input_dim) if transposed is `False`,
            else of shape (num_experts, input_dim, output_dim).
        bias (`torch.Tensor`, *optional*):
            Bias tensor of shape (num_experts, output_dim). Default is `None`.
        offs (`torch.Tensor`, *optional*):
            Offsets tensor indicating the boundaries of each group in the input tensor.
        is_transposed (`bool`, *optional*, defaults to `False`):
            Whether the weight tensor is transposed.
    Returns:
        `torch.Tensor`: Output tensor of shape (S, output_dim).
    """
    # torch.nn.functional.grouped_mm is not autocast-enabled, so we disable autocast to avoid dtype mismatch.
    # See: https://github.com/pytorch/pytorch/issues/174763
    with torch.amp.autocast(device_type=input.device.type, enabled=False):
        if is_transposed:
            # (S, input_dim) @ grouped (num_experts, input_dim, output_dim) -> (S, output_dim)
            out = torch.nn.functional.grouped_mm(input, weight, offs=offs)
        else:
            # (S, input_dim) @ grouped (num_experts, output_dim, input_dim).T -> (S, output_dim)
            out = torch.nn.functional.grouped_mm(input, weight.transpose(-2, -1), offs=offs)

    if bias is not None:
        # We should be able to pass bias to the grouped_mm call, but it's not yet supported.
        out = out + bias

    return out


def grouped_mm_experts_forward(
    self: torch.nn.Module,
    hidden_states: torch.Tensor,
    top_k_index: torch.Tensor,
    top_k_weights: torch.Tensor,
) -> torch.Tensor:
    if not hasattr(torch.nn.functional, "grouped_mm"):
        raise ImportError(
            "torch.nn.functional.grouped_mm is not available. "

Review comment (suggestion):
This PR would now force users to torch 2.10 to access the grouped_mm apis, leaving those still using 2.9 unable to use it whereas they were previously able to. Perhaps something with a fallback like
```python
grouped_mm_fn = torch.nn.functional.grouped_mm if hasattr(torch.nn.functional, "grouped_mm") else torch._grouped_mm
```

Code after:

def _grouped_mm(
    input: torch.Tensor,
    weight: torch.Tensor,
    offs: torch.Tensor | None = None,
) -> torch.Tensor:
    """Grouped matrix multiplication dispatcher that uses torch.nn.functional.grouped_mm if available, else falls back to torch._grouped_mm.
    Args:
        input (`torch.Tensor`):
            Input tensor of shape (S, input_dim).
        weight (`torch.Tensor`):
            Weight tensor of shape (num_experts, output_dim, input_dim).
        offs (`torch.Tensor`, *optional*):
            Offsets tensor indicating the boundaries of each group in the input tensor.
    Returns:
        `torch.Tensor`: Output tensor of shape (S, output_dim).
    Raises:
        ImportError: If neither `torch.nn.functional.grouped_mm` nor `torch._grouped_mm` is available, indicating that the PyTorch version is incompatible.
    """

    # torch.nn.functional.grouped_mm is not autocast-enabled, so if autocast is enabled we might end up with input tensors in fp32 (e.g. LayerNorm output) and weight tensors in fp16/bf16
    # In that case we need to cast the input to the weight dtype to avoid dtype mismatch errors.
    # See: https://github.com/pytorch/pytorch/issues/174763
    if hasattr(torch.nn.functional, "grouped_mm"):
        return torch.nn.functional.grouped_mm(input.to(weight.dtype), weight, offs=offs)
    elif hasattr(torch, "_grouped_mm"):
        return torch._grouped_mm(input.to(weight.dtype), weight, offs=offs)
    else:
        raise ImportError(
            "Neither torch.nn.functional.grouped_mm nor torch._grouped_mm is available. "
            "Please make sure you are using a PyTorch version that includes grouped_mm (2.9+)."
        )


def _grouped_linear(
    input: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None = None,
    offs: torch.Tensor | None = None,
    is_transposed: bool = False,
) -> torch.Tensor:
    """Grouped linear layer supporting optional bias and transposed weights.

    Args:
        input (`torch.Tensor`):
            Input tensor of shape (S, input_dim).
        weight (`torch.Tensor`):
            Weight tensor of shape (num_experts, output_dim, input_dim) if transposed is `False`,
            else of shape (num_experts, input_dim, output_dim).
        bias (`torch.Tensor`, *optional*):
            Bias tensor of shape (num_experts, output_dim). Default is `None`.