[Bug]: Errors in edge cases in parallel matrix multiplication

[brownbaerchen]

What happened?

In trying to debug something, I encountered this weird edge case where the result of parallel matrix multiplication is wrong. The error is detailed below and appears only when run with two tasks. I am not sure what makes this edge case special.
I have been running into this issue first with QR decomposition. You can try:

import heat as ht

comm = ht.comm
A = ht.random.randn(comm.size, comm.size, dtype=ht.float32, split=0)
QR = ht.linalg.qr(A)

matmul_success_loc = ht.allclose(QR.Q @ QR.R, A)
matmul_success_glob = ht.allclose((QR.Q.resplit(None) @ QR.R.resplit(None)).resplit(QR.Q.split), A)
print(matmul_success_loc, matmul_success_glob)

which will print False, True.

I am a bit lost here. In particular, the matrix multiplication function is very long and hard to understand. So I thought I just report this bug because it seems pretty bad and maybe somebody can give me some pointers.

Code snippet triggering the error

import heat as ht

split = 0
shape = (4, 3)

A = ht.ones(shape, split=split)
B = ht.ones(shape[::-1], split=split)

C = A @ B
C_glob = (A.resplit(None) @ B.resplit(None)).resplit(C.split)

print(C)
print(C_glob)

Error message or erroneous outcome

Output when run with two tasks: 

DNDarray([[3., 3., 3., 3.],
          [6., 6., 6., 6.],
          [3., 3., 3., 3.],
          [6., 6., 6., 6.]], dtype=ht.float32, device=cpu:0, split=0)
DNDarray([[3., 3., 3., 3.],
          [3., 3., 3., 3.],
          [3., 3., 3., 3.],
          [3., 3., 3., 3.]], dtype=ht.float32, device=cpu:0, split=0)

Version

main (development branch)

Python version

3.13

PyTorch version

2.9

MPI version

OpenMPI 5.0.8

[mrfh92]

@brownbaerchen some edgecases that definitely lead to problems are the ones where matrices are split but so small that there are processes without data.

[brownbaerchen]

Another edge case that raises an error rather than producing the wrong result when the matrix size is not divisible by the number of tasks popped up in #2195. I haven't fully understood what happens there, but yet another bug that needs fixing!

[mrfh92]

@brownbaerchen fixing matmul is, I guess, one of the most challenging things in Heat 👍

[JiwaniZakir]

The doubling pattern (alternating rows of 3 and 6) is a classic double-accumulation symptom: with 2 tasks, each process owns 2 rows of A and 2 rows of B (split=0), and the contraction axis of B is also the split axis, so the algorithm must scatter/reduce partial products across ranks. The likely root cause is in the mm/matmul implementation (probably heat/core/linalg/matmul.py) in the branch that handles split_0 == 0, split_1 == 0 — specifically, the local partial product appears to be accumulated both into the running sum before the Reduce_scatter (or equivalent) and again as part of the received buffer, causing each rank's own contribution to be counted twice. The QR reproduction confirms this isn't specific to the trivial ones-matrix: any split=0 × split=0 matmul where the inner dimension is distributed will silently produce wrong results whenever the number of processes divides the row count evenly (making the off-by-one in index arithmetic invisible in other shapes).

[github-actions]

Branch bugs/2093-_Bug_Errors_in_edge_cases_in_parallel_matrix_multiplication created!