ReduceScatter with DID loop split

[wujingyue]

For #2563

Tested: http://nv/eoZ

[wujingyue]

This allreduce test is merely for DID logical split. I don't think allreduce can support DID loop split because sum's reduction axes can only be logical. But I'd be happy to know otherwise.

[naoyam]

This allreduce test is merely for DID logical split.

Just to be clear, you meant DID parallelization of logical domains, right? I'm not sure what you meant by DID logical split otherwise.

[naoyam]

Assuming I understand what you meant correctly, I think this is where TensorView::rFactor could be used. That's what we use for intra-device hierarchical reductions. For example, I'd think that for multi-GPU reductions, we would have something like:

(I'm mixing the C++ and Python APIs)

self.out->split(0, num_devices, /*inner=*/false);
auto intermediate_result = self.out->rFactor({1});
intermediate_result->axis(0)->parallelize(DIDx);
self.out->axis(0)->parallelize(DIDx);

Here, intermediate_result would be the partial result of per-device reduction, which would be then reduced between all the devices and saved to self.out.

[wujingyue]

I agree it's something like rfactor, and did look into how TensorView::rfactor works in

Fuser/tests/cpp/test_tutorial.cpp

Line 344 in ecabd46

TensorView* tv2 = tv1_copy->rFactor({0});

. However, I failed to see how it applies here.

If we want to loop (but not logical) split an allreduce, the input would be a logical shape like [D*2,3] and the output would be of logical shape like [2,3]. Regardless of scheduling, what ops in fusion IR could do that? (Not a sum because that reduces an entire dimension to 1).

[naoyam]

Let's talk offline. It seems we are not using the same vocabulary (e.g., I don't understand what "loop split" and "logical split" mean).

[wujingyue]

#3543 is my failed attempt. It triggered an assertion at

Fuser/csrc/scheduler/vectorize_helper.cpp

Line 1063 in 9346c8f

NVF_THROW("Unexpected producer RF ID: ", producer_rf_id->toString())

. Is it because code there has been assuming that reductions are innermost?

[wujingyue]

Anyhow, this isn't a blocker. As we discussed yesterday, we'll probably stick with logical split for reductions in Allreduce and ReduceScatter due to MatmulOp's implementation.

[wujingyue]

!test

[naoyam]

LGTM

[wujingyue]

!test

[wujingyue]

!test

[wujingyue]

!test

[wujingyue]

!build