Use matmul in distributed tests

[Priya2698]

Issue #2372.
Modifying the tests to use matmul in place of mul-sum.
Note: matmul API requires the logical shapes [M,K] x [K,N] and the output has the same dtype as the input.

[wujingyue]

Is it worth adding a check to make sure ATen kicked in? E.g.

Fuser/tests/cpp/test_matmul_aten_evaluation.cpp

Line 318 in 7b780e5

EXPECT_FALSE(executors.front().hasCompiledKernel());

[cowanmeg]

Thanks! I'm glad the ATen pathway worked out of the box!
I assume the extra size 1 axes in the aten tensors are being treated as batch dimensions which is why it works out ok right?

[cowanmeg]

This is actually an error from me....but K should be Ko.
I'm surprised it compiled and was correct before.

[wujingyue]

Thanks! I'm glad the ATen pathway worked out of the box! I assume the extra size 1 axes in the aten tensors are being treated as batch dimensions which is why it works out ok right?

Yes, that is treated as a BatchMatmul, case 5 in https://pytorch.org/docs/stable/generated/torch.matmul.html

[Priya2698]

EXPECT_FALSE(executors.front().hasCompiledKernel());

Is it worth adding a check to make sure ATen kicked in? E.g.

Fuser/tests/cpp/test_matmul_aten_evaluation.cpp

Line 318 in 7b780e5

EXPECT_FALSE(executors.front().hasCompiledKernel());

How do I access the executors from MultiDeviceCommunicator?
For these tests, since there will be multiple segments, we need to test for the second segment/executor (first is no-op for transpose).

[Priya2698]

!build

[cowanmeg]

Sorry for adding this late - can we add back the original TEST_F(DistributedMatmulTest, LayoutTN_NoComms) with mul-sum in addition to the updated one with the matmul node. We should keep one of these tests since they call the MatmulScheduler with sharded TensorViews.

How do I access the executors from MultiDeviceCommunicator?

MultiDeviceExecutor doesn't have accessor functions for executors so we would have to add that first. @wujingyue

Some background:
The MultiDeviceExecutor takes in a Fusion and segments it into compute and communication segments. The compute segments then go through FusionExecutorCache then they go through the standard segmentation.
So we need to do something like get the first "coarse" segment (which will correspond to the local compute at the beginning) then index into the proper segment/executor.

See here for the internals:
https://github.com/NVIDIA/Fuser/blob/main/csrc/multidevice/executor.h#L130
The MultiDeviceExecutor can use either FusionExecutorCache or directly use FusionExecutor for compute segments depending on flags (temporary for development purposes). For the DistributedMatmulTest I set the params to use FusionExecutorCache.

[wujingyue]

How do I access the executors from MultiDeviceExecutor?

Good question. We could plumb it through but I don't think it's worth the benefits at this moment.

[wujingyue]

Can we add back the original TEST_F(DistributedMatmulTest, LayoutTN_NoComms) with mul-sum in addition to the updated one with the matmul node. We should keep one of these tests since they call the MatmulScheduler with sharded TensorViews.

Good question!

AFAIK, nvFuser offers three ways (arguably too many, 🤷) to generate a matmul.

1. fusedMultiplySum, which takes broadcasted operands and generates an MmaOp (for MatmulScheduler)
2. matmul, which takes non-broadcasted operands and sometimes generates a MatmulOp (for ATen).
3. The same matmul API, which sometimes generates a decomposed broadcast+mul+sum pattern (e.g. when the inputs are vectors).

The current tests in this file exercise only (2).

I agree it's useful to exercise (1). MatmulScheduler is already DID-capable and eventually we do want nvFuser to generate matmul using the scheduler, so it's good to have tests to maintain that feature. However, I'm not aware that nvFuser can reliably turn a decomposed broadcast+mul+sum into an MmaOp/MatmulOp. Correct me if I'm wrong. So we may need to add new tests that use fusedMultiplySum instead.

I was told by @Priya2698 that (3) is a corner case that will go away. So I'll leave the decision to her whether to cover that as well.

[Priya2698]

1. fusedMultiplySum, which takes broadcasted operands and generates an MmaOp (for MatmulScheduler)
2. matmul, which takes non-broadcasted operands and sometimes generates a MatmulOp (for ATen).
3. The same matmul API, which sometimes generates a decomposed broadcast+mul+sum pattern (e.g. when the inputs are vectors).

The current tests in this file exercise only (2).

I agree it's useful to exercise (1). MatmulScheduler is already DID-capable and eventually we do want nvFuser to generate matmul using the scheduler, so it's good to have tests to maintain that feature. However, I'm not aware that nvFuser can reliably turn a decomposed broadcast+mul+sum into an MmaOp/MatmulOp. Correct me if I'm wrong. So we may need to add new tests that use fusedMultiplySum instead.

I was told by @Priya2698 that (3) is a corner case that will go away. So I'll leave the decision to her whether to cover that as well.

Yes, (3) will be removed. matmul API will always generate the MatmulOp.

For now, from the python frontend API, the only way is to call fd.ops.matmul or similarly a fd.ops.mul + fd.ops.sum. However, thunder will plumb down matmuls using the former.

To exercise the matmul scheduler:

1. Create a mul-sum pattern or fusedMultiplySum which both convert to MmaOp: Generalize CombineMulSum as MatmulPatterns #2272

2. @jacobhinkle added translation for the MatmulOp to MmaOp for matmul scheduler: Translate MatmulOp and LinearOp #2236. Use NVFUSER_ENABLE=fuse_matmul which enables the matmul scheduler to take segments with matmul if expression evaluator cannot (for eg: with epilogues). We can also disable expression evaluator using NVFUSER_DISABLE=matmul_expr_eval to only use one or the other.

[wujingyue]

Thanks for clarifying. I missed

Fuser/csrc/scheduler/mma_utils.cpp

Line 1462 in 0f66dc2

void handle(ReductionOp* rop) override {

-- nvFuser does convert a decomposed mul+sum to MmaOp.

So we have three ways to exercise the matmul scheduler:

1. mul-sum
2. fusedMultiplySum
3. matmul with fuse_matmul enabled.

Wdyt? (3) seems to me the best option: we'll get fd.ops.matmul from Thunder and stick with that for at least O(months). Btw, when are you going to enable fd.ops.matmul for Thunder by default?

[Priya2698]

Wdyt? (3) seems to me the best option: we'll get fd.ops.matmul from Thunder and stick with that for at least O(months).

Yes, we can use the NVFUSER_ENABLE/DISABLE options to control which scheduler to use. It would good to have some way to verify the scheduler heuristic used though. There is no need to plumb down a separate option if I can use MultiDeviceCommunicator->fec_?

Btw, when are you going to enable fd.ops.matmul for Thunder by default?

We are still seeing some issues such as #2354, so it is not enabled by default yet. I plan on running all the thunder benchmarks and tests with this enabled to preemptively identify existing issues first.

[cowanmeg]

MultiDeviceCommunicator->fec_

We can write an accessor function to return the FusionExecutorCache(s) used by the MultiDeviceRuntime. It will be a vector because each segment has its own FusionExecutorCache. In this case since there's only one compute segment it will only have one entry.

Related translation for the MatmulOp to MmaOp for matmul scheduler we eventually need to update insertResharding pass to handle Matmul, Linear, and SDPA ops to insert set for non-reduction collectives. It's easy to work around by manually inserting set which the tests were already doing.

[Priya2698]

MultiDeviceCommunicator->fec_

We can write an accessor function to return the FusionExecutorCache(s) used by the MultiDeviceRuntime. It will be a vector because each segment has its own FusionExecutorCache. In this case since there's only one compute segment it will only have one entry.

Related translation for the MatmulOp to MmaOp for matmul scheduler we eventually need to update insertResharding pass to handle Matmul, Linear, and SDPA ops to insert set for non-reduction collectives. It's easy to work around by manually inserting set which the tests were already doing.

So currently, what you suggest is retaining one of the original test cases?
Or do we want to switch to only using MatmulOp and use the nvfuser options to pick one of the schedulers?

[cowanmeg]

Let's keep one of the original test cases.

[Priya2698]

!build

[cowanmeg]

LGTM! Thank you!

[Priya2698]

!build