[Distributed] Switch all_reduce to use the new functional collective op

[yifuwang]

PyTorch has implemented a new set of functional collective ops and is planning to remove the old ops. Migrating all_reduce to use the new op.

See context in pytorch/pytorch#93173 (comment)

[alanwaketan]

Let me know if you need any helps to setup XLA environment?

[yifuwang]

Let me know if you need any helps to setup XLA environment?

Hey @alanwaketan, I'd appreciate your help on this. Is it possible to setup the dev env without having access to docker?

[alanwaketan]

Let me know if you need any helps to setup XLA environment?

Hey @alanwaketan, I'd appreciate your help on this. Is it possible to setup the dev env without having access to docker?

Right...

[alanwaketan]

Let me know if you need any helps to setup XLA environment?

Hey @alanwaketan, I'd appreciate your help on this. Is it possible to setup the dev env without having access to docker?

Have you followed https://github.com/pytorch/xla/blob/master/CONTRIBUTING.md? What error do you get? I wonder may be you can just develop and test this on the cpu machine and no need to access TPU.

[yifuwang]

Have you followed https://github.com/pytorch/xla/blob/master/CONTRIBUTING.md? What error do you get? I wonder may be you can just develop and test this on the cpu machine and no need to access TPU.

Yea I tried to follow it. But building from source requires docker which is not available in my dev env :(

[alanwaketan]

Have you followed master/CONTRIBUTING.md? What error do you get? I wonder may be you can just develop and test this on the cpu machine and no need to access TPU.

Yea I tried to follow it. But building from source requires docker which is not available in my dev env :(

Oh, so you cannot use docker or you cannot access our docker?

[yifuwang]

Oh, so you cannot use docker or you cannot access our docker?

Cannot use docker. It's a restriction of my corporate dev env :(

[alanwaketan]

Maybe you can leave the development to our team then. I don't know how difficult it is to use non-docker env... @JackCaoG Do you know?

[yifuwang]

Maybe you can leave the development to our team then.

Sounds good. Thanks for the help @alanwaketan! I just wanted to make sure there wasn't a gap for torch-xla to adopt the new API.

[yifuwang]

@alanwaketan - following your advices, the CI is now green.

Note this PR doesn't address the TODO re. generating groups. It merely switches to the new API while being consistent with the old behavior. Happy to continue discussing how to plumb through group information.

Let me know what you think.

[alanwaketan]

Thanks, Yifu. We don't need group information in XLA in general. So we can follow up on that later. Let me trigger the TPU CI. Once everything is green. I will just merge the PR.

[alanwaketan]

LGTM.

[alanwaketan]

Oh, we cannot run TPU CI on fork I guess... Never mind. I will just merge it. The head CI will take care of the rest.