LLVM error during GRPO training with Apple M4 Max

[neslihaneti]

I have the below error while doing GRPO training. I am using HuggingFace example codes for GRPO. I couldn't run the model on MPS because of this issue.
How can I run GRPO on MPS?

loc("mps_matmul"("(mpsFileLoc): /AppleInternal/Library/BuildRoots/4B_wkugAG-524HdEQLaK0kvU7Y_D8Jtm6UxMaIoY/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm":43:0)): error: incompatible dimensions
loc("mps_matmul"("(mpsFileLoc): /AppleInternal/Library/BuildRoots/4B_wkugAG-524HdEQLaK0kvU7Y_D8Jtm6UxMaIoY/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm":43:0)): error: invalid shape
LLVM ERROR: Failed to infer result type(s).

Details:
OS: Tahoe 26.2
pytorch 2.9.1
trl: 0.26.1
MLX:0.30.0