transformers
007bb8c3 - fix(moe): Handle dtype mismatch in torch._grouped_mm with autocast (#43839)

Commit
191 days ago
fix(moe): Handle dtype mismatch in torch._grouped_mm with autocast (#43839) * fix(moe): Handle dtype mismatch in grouped_mm with autocast torch._grouped_mm is not autocast-enabled, so when using torch.autocast with MoE models like Phi-tiny-MoE, the input tensor may have a different dtype than the weights, causing RuntimeError. Add explicit dtype casting in _grouped_linear to ensure input and weight tensors have the same dtype before calling torch._grouped_mm. Fixes #43828 * refactor: use autocast(enabled=False) instead of manual dtype cast Replace manual input.to(weight.dtype) cast with torch.amp.autocast(enabled=False) wrapper around torch._grouped_mm calls. This is cleaner since it addresses the root cause (grouped_mm not being autocast-aware) rather than guessing the target dtype. Ref: https://github.com/pytorch/pytorch/issues/174763
Author
Parents
Loading