How to define Custom Communication Operations for Custom Operators in Distributed Settings

[Doraemonzzz]

Thank you for your awesome project. I would like to ask how to solve the following issue:

I have implemented the logcumsumexp operator, where the input placement is Shard(-1) and the output placement is Replicate(). To obtain the final result, I need to create a custom all-reduce operator (instead of using the conventional sum). How should I go about implementing this?

More generally, for an operator function f, given an input placement1 and an output placement2, where should I implement various custom communication operations? I would greatly appreciate it if you could provide some examples for this.

[tianyu-l]

It sounds like you would like to register a sharding strategy for your customized op.

DTensor register_sharding API may be something you are looking for.
https://pytorch.org/docs/stable/distributed.tensor.html#torch.distributed.tensor.experimental.register_sharding

[Doraemonzzz]

Thank you for the information you provided. I checked the code, but I’m even more confused now. Suppose we apply softmax on the dim, and the input placement is Shard(dim) while the output placement is Replicate(). There should be a customized all-reduce for LSE (logsumexp) here, but I don’t seem to see it in the code.

[tianyu-l]

Hmm let's take a step back. Can you share more context so that I can better help?

When you say

where the input placement is Shard(-1) and the output placement is Replicate().

What do you exactly mean? Can you write some (pseudo)code to explain what's going on, e.g.

• what are the input and output of this function and what are their sharding placements
• what does the function do
• if communication needs to happen, why it has to happen and where do you hope it to happen