[core][compiled graphs] Supporting allreduce on list of input nodes

[anyadontfly]

Why are these changes needed?

Currently, only one dag node per actor can be bound to allreduce. When the users want to schedule output tensors from several dag nodes into one allreduce call, they would have to collect tensors into a tuple and return the tuple of tensors from a single dag node. For example,

@ray.remote
class Actor:
    def comp1(self, _):
        return torch.ones(1,)

    def comp2(self, _):
        return torch.ones(1,) * 2

    def get_results(self, *args):
        return tuple(args)

To launch allreduce on the result of comp1 and comp2, users need an additional function get_results to gather the result of comp1 and comp2.

with InputNode() as inp:
    res_comp1 = [actor.comp1.bind(inp) for actor in actors]
    res_comp2 = [actor.comp2.bind(inp) for actor in actors]
    res_tuple = [actor.get_results.bind(res_comp1[i], res_comp2[i]) for i, actor in enumerate(actors)]
    res_ar = allreduce.bind(res_tuple)

In this PR, users can simply put the result of comp1 and comp2 in a list, and launch allreduce on the list of outputs, which no longer requires the intermediate function get_results.

with InputNode() as inp:
    res_comp1 = [actor.comp1.bind(inp) for actor in actors]
    res_comp2 = [actor.comp2.bind(inp) for actor in actors]
    res_ar = allreduce.bind([res_comp1, res_comp2])

Related issue number

Meta-issue: #47983

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[jjyao]

@anyadontfly could you rebase, there is merge conflict

[anyadontfly]

Hi Stephanie, the PR is ready to run more tests. Thanks! @stephanie-wang

[stephanie-wang]

Actually I don't quite understand this logic. I thought we should always return the same type of node, just sometimes it will be a list and sometimes it will be a nested list. Is this method now returning either a list of ClassMethodNodes or a list of CollectiveOutputNodes? If so, please update it so that we only return one type of node.

Also, please update the type signature accordingly.

[anyadontfly]

I used the same logic as ClassMethodNode with multiple return nodes so that we can execute the nccl operation once with multiple return nodes.

If no nested list passed in (normal allreduce), _bind will return CollectiveNode which is equivalent to a ClassMethodNode with _is_class_method_output=False.

If a nested list passed in (bucketed allreduce), _bind will return ClassMethodNode with _is_class_method_output=True and set _class_method_output to the original CollectiveNode so that the nccl operation will be execute only once at runtime instead of once for each object.

[stephanie-wang]

Can we modify this so that we always use the same logic for len(input_node_list) == 1 or N? Conceptually, it would be better if we always think of a CollectiveNode as taking in a nested list, and the only difference here should be whether we return a flat or nested list.

[anyadontfly]

I changed the returned node for N input_node_lists to CollectiveNode.

I'm trying to inherit ClassMethodNode implementation of num_returns == 1 or N.

For allreducing 1 object, we can perform the allreduce and return the result from the CollectiveNode. But for allreducing N objects, we still want to perform allreduce only once on the CollectiveNode. So we need intermediate nodes to distribute results to different nodes.

If I use the same logic for 1 object or N objects, I can return an intermediate node from .bind as the case for N objects, but I think that would result in a redundant node.

[stephanie-wang]

This comment is confusing. There is no input node list in this function?

[anyadontfly]

Yes I think so. The inputs and outputs are tensors in execute.

[stephanie-wang]

Why do we allocate a separate torch tensor for each input instead of one flat tensor?

[anyadontfly]

I changed the implementation here to recv_buf pointing to flat_buf to avoid unnecessary memory allocation.

[stephanie-wang]

Can we modify this so that we always use the same logic for len(input_node_list) == 1 or N? Conceptually, it would be better if we always think of a CollectiveNode as taking in a nested list, and the only difference here should be whether we return a flat or nested list.

[stephanie-wang]

Thanks, this looks good! Please just address the comment about dtypes.

[stephanie-wang]

You can support this case by using torch.view. But it's fine to do it in a follow-up PR.

[anyadontfly]

got it, thanks!