fix(moe): Handle dtype mismatch in torch._grouped_mm with autocast

[Mr-Neutr0n]

What does this PR do?

Fixes a RuntimeError: expected mat1 and mat2 to have the same dtype error when using torch.autocast with MoE models like microsoft/Phi-tiny-MoE-instruct.

Problem

torch._grouped_mm is not autocast-enabled, meaning it doesn't automatically cast input tensors to match when torch.autocast is active. This causes a dtype mismatch when:

1. Model weights are in bfloat16 (e.g., loaded with dtype="bfloat16")
2. Input hidden states may be float32 at certain computation stages
3. The grouped GEMM operation fails with dtype mismatch error

Error trace:

RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::BFloat16

Solution

Add explicit dtype casting in _grouped_linear() to ensure input and weight tensors have the same dtype before calling torch._grouped_mm:

if input.dtype != weight.dtype:
    input = input.to(weight.dtype)

This is the same pattern used in other low-level ops that aren't autocast-enabled.

Note: _batched_linear doesn't need this fix because torch.bmm IS autocast-enabled.

Fixes #43828

[IlyasMoutawwakil]

hi ! wouldn't it make more sense to disable autocast around grouped_mm ?

[hxrikp1729]

@IlyasMoutawwakil yeah, that's a fair point. Disabling autocast around grouped_mm would be cleaner since it addresses the root cause directly rather than adding a manual cast. The current approach works but it does mean we're doing an explicit dtype cast that could technically diverge from what autocast would have chosen.

Would something like wrapping the torch._grouped_mm calls with torch.amp.autocast(enabled=False) be what you had in mind? Happy to rework the PR in that direction if so.

[IlyasMoutawwakil]

yes let's do that, and open an issue in https://github.com/pytorch/pytorch/issues
@ArthurZucker wdyt ? i saw a couple PRs trying to fix this by casting inputs/outputs

[Mr-Neutr0n]

@IlyasMoutawwakil Done — reworked to wrap the torch._grouped_mm calls with torch.amp.autocast(device_type=..., enabled=False) instead of the manual dtype cast. Much cleaner.

Also opened the upstream PyTorch issue: pytorch/pytorch#174763

[IlyasMoutawwakil]

LGTM ! did you verify that it solves the issue on the relevant model ?

[ArthurZucker]

#43833 came first, but this fix is more elegant so let's go with this one!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[Mr-Neutr0n]

Thank you for the kind words @ArthurZucker