DIDx aware ATen evaluation for matmul and linear.

[Priya2698]

Motivation: Using MatmulOp/LinearOp in multi-GPU work.

Issue: We currently do not check for the presence of DID axis in the inputs to evaluate for any of the nodes.

Example: Consider: A (M x K), B (K x N) are sharded on 2-devices, with rank: 0 computing the first M/2 rows of C (MxN):
In the current approach, the inputs to evaluate will be: A = [1 x M/2 x K], B = [2 x K x N/2] (assuming there is an all-gather step to get the complete B matrix so we compute the first half of the output, instead of only the first quadrant). The output will be: C = [2 x M/2 x N/2].

Correct computation: reorder and reshape B = [K, N] before matmul to get [1 x M/2 x N].

CC: @wujingyue @cowanmeg. Correct me, if my understanding is incorrect.

[cowanmeg]

Yes, that is the right understanding. We can assume the resharding passes will properly add appropriate collectives. One correction is that B's layout will be [K, N] after it is all-gathered. If the tensor is not sharded it won't have an extra dimension.
A: [1, M/2, K] and B: [K, N] -> C: [1, M/2, N]
So only A, C will have to be reshaped (unsqueezing the DID axes)

1. Reshape input(s) to remove the extra DID dimensions, reorder into ATen format
2. ATen computation
3. Reshape to add DID dimensions back and reorder back to the nvFuser format

[cowanmeg]

This unit test is a good one to look at: https://github.com/NVIDIA/Fuser/blob/main/tests/cpp/test_multidevice_matmul.cpp#L69
It doesn't have any communication just the matmul computation.

[Priya2698]

Yes, that is the right understanding. We can assume the resharding passes will properly add appropriate collectives. One correction is that B's layout will be [K, N] after it is all-gathered. If the tensor is not sharded it won't have an extra dimension. A: [1, M/2, K] and B: [K, N] -> C: [1, M/2, N] So only A, C will have to be reshaped (unsqueezing the DID axes)

1. Reshape input(s) to remove the extra DID dimensions, reorder into ATen format
2. ATen computation

What is the expected output from the ATen computation -- [1, M/2, N] or [M/2, N]?
If A and B are as described (i.e. there is no reshape and reorder needed on B), then the current implementation will output C as [1, M/2, N] for matmul.

3. Reshape to add DID dimensions back and reorder back to the nvFuser format

What do you mean by reordering to nvFuser format?

[cowanmeg]

What is the expected output from the ATen computation -- [1, M/2, N] or [M/2, N]?

[1, M/2, N]. The logical domain of C will be [DIDx(2), M, N]. We have DID axes as size 1 during binding which is why we have the extra 1.

nvFuser format

Sorry disregard that. The output of the matmul should always be [M, N]

[Priya2698]

What is the expected output from the ATen computation -- [1, M/2, N] or [M/2, N]?

[1, M/2, N]. The logical domain of C will be [DIDx(2), M, N]. We have DID axes as size 1 during binding which is why we have the extra 1.

This is what the current implementation should output. I will verify the tests.

[Priya2698]

@cowanmeg: We likely don't need any changes to MatmulOp::evaluate with the layouts you described.
I have the updated tests in my branch here: https://github.com/NVIDIA/Fuser/tree/pm/test_dist_matmul.

One important different between mul-sum/matmul is that the logical shape required for matmul is [M,K]x[K,N].
If you want, I can create a PR for review.

[Priya2698]

Closed with #2386