[RFC] Improved coverage for native distributed collective operations

[bfolie]

🚀 Feature

Improve coverage of the PyTorch collective operations so that native distributed code is more likely to work without any user modification required. Collective ops would ideally work both in LazyTensor and in compiled mode, but many do not have an upstream Dynamo path.

The following are in scope.

• All Gather Into Tensor
• Broadcast
• Scatter
• Reduce Scatter
• Send and Recv
• Gather
• Reduce
• All to All

Collectives that operate on objects instead of tensors are out of scope.

Pitch

All Gather Into Tensor

Gathers tensors into a single output tensor. Like the already-implemented all_gather but outputs a single tensor instead of a list of tensors. It appears in FSDP (example), and other places.

This is implemented by wrapping all_gather, so it fails in some cases because it has not been set up to support both stacking and concatenation of the input tensors. This zip, which is meant for a list of tensors, happens to work in some cases by splitting the tensors along the 0th dimension.

This should be a simple refactor, isolating the common logic but letting stacking/concatenation/copying to individual tensors happens only in the correct cases.

Done in #9332

Broadcast

This has been solved, but PR 7956 and PT PR 135171 were not merged. They should be revived. Broadcast is used in DDP.

Scatter

Scatters a list of tensors across processes. It is used by sharding code. Scatter does not have a corresponding XLA op, but we could implement it similarly to how broadcast is implemented, and perform a reduce_scatter after multiplying all but the non-source device tensors times 0.

Done in #9365

Reduce Scatter

Reduces a list of tensors, then scatters across processes. This is like the already-implemented reduce_scatter_tensor, but it acts on a list of tensors instead of a single tensor. It is not used in the torch.distributed code but may be necessary to implement scatter.

Reduce Scatter works in LazyTensor mode, but fails when compiled because there's no Dynamo mapping. There isn't even a reduce_scatter functional collective but reduce_scatter_tensor_coalesced might be usable. Once implemented we would bind it in pt/xla (the analogous binding for reduce_scatter_tensor is here). After implementing the binding the logic can be shared with the existing reduce_scatter_tensor, and the underlying XLA Op can accept a tuple of arrays, so the rest should be straight-forward.

Send and Recv

Send and receive are implemented but do not work. As seen in #8074, the XLA ops are missing "_xla_send_recv_source_target_pairs". Those fields get set here. XLA's Send and Recv aren't meant to be called directly. Instead the user is expected to use CollectivePermute and specify every source-target pair. If there are no cycles then the HLO is decomposed into Send and Recv ops.

We can replace send/recv with calls to xm.collective_permute..

Done in #9373

Gather

This is not used in the torch.distributed code but was requested in #9069. It could be implemented using all_gather and only keeping the result on the dst rank.

Reduce

This is not used in the torch.distributed code. It could be implemented using all_reduce and only keeping the result on the dst rank.

All to All

This is not used in the torch.distributed code. Since all_to_all_single is implemented we could probably stack the tensors, run all_to_all, then chunk them.

[rpsilva-aws]

Thanks @bfolie, this is great. For "Send and Recv", aren't the referenced passes for GPU XLA? I am not sure the motivation there was that Send/Recv was not suitable for XLA, but rather that they needed a way to convert CollectivePermute to Send/Recv, since afaik, they don't support the former (also just found a reference). I was under the the same impression that the push for frontend attribute to include src/tgt pairs stemmed from the same reason.

Is there any other reasons for "XLA's Send and Recv aren't meant to be called directly"? They should still be valid XLA operation semantics.

[bfolie]

For "Send and Recv", aren't the referenced passes for GPU XLA? I am not sure the motivation there was that Send/Recv was not suitable for XLA, but rather that they needed a way to convert CollectivePermute to Send/Recv, since afaik, they don't support the former (also just found a reference). I was under the the same impression that the push for frontend attribute to include src/tgt pairs stemmed from the same reason.

Is there any other reasons for "XLA's Send and Recv aren't meant to be called directly"? They should still be valid XLA operation semantics.

It's possible that collective_permute_decomposer is only for GPU. Admittedly I didn't look to see where in the XLA codebase it is called. Do you know where that happens? Appreciate the linked NCCL issue -- that provides more context.

Regardless, the important fact is that CollectivePermute should work for all accelerators regardless of whether and how it gets decomposed. This provides us with a safe (though not necessarily the most performant) path to supporting native PT code that uses torch.dist.send/recv.

As for whether XLA's Send/Recv are meant to be called directly, you're right that they are valid. But when making something work requires setting a hidden attribute that isn't mentioned in the docs, that feels to me like it's not meant to be used. And since we can get to what I think is the same place via the XLA CollectivePermute op, that strikes me as the best path.

[rpsilva-aws]

We can follow-up on these, but I am not sure that PyTorch/XLA remains the place to enforce this decision on.

Regardless, the important fact is that CollectivePermute should work for all accelerators regardless of whether and how it gets decomposed.

But when making something work requires setting a hidden attribute that isn't mentioned in the docs, that feels to me like it's not meant to be used.

These accelerators are usually in control of how they want to decompose any ops. For instance, TRN supports send/recv in some context, and we'd be voiding users to implicitly lower to send/recv. I'd understand some documentation, warning, etc - but I think we need to be particularly careful with doing these transformations here, as it'll be hard to move back moving forward. On the second reference above, we need to make sure that this is for any accelerator, and not GPU only. I believe "supporting native PT code that uses torch.dist.send/recv" is definitely a valid point, but slightly orthogonal. We can always intercept and change it to a CollectivePermute in the context of native PT code, but we can decouple it from the discussion of doing the transformation on all cases.

But when making something work requires setting a hidden attribute that isn't mentioned in the docs, that feels to me like it's not meant to be used.

This is for XLA GPU, and iirc, it's could be partially to mimic the CollectivePermute behavior with Send/Recv. We can check if they have any other valid use case, as with any other accelerators.

A few other things to check:
1/ Does JAX do this transformation implicitly?
2/ There's a chance that Send/Recv will no longer have this "hidden attribute", as GPU XLA looks into standardizing that spec. See https://github.com/openxla/stablehlo/blob/335a3fb58b3dd7cfa877835584e002ebad01ba56/rfcs/20250421-source-target-pairs.md (cc: @GleasonK).

3/

Admittedly I didn't look to see where in the XLA codebase it is called. Do you know where that happens?

I haven't had the time to, but I think this could help with clarifying the motivation from XLA further.

[bfolie]

I believe "supporting native PT code that uses torch.dist.send/recv" is definitely a valid point, but slightly orthogonal.

There may be some miscommunication, because that is the point of this RFC. It's all about making the collective ops in torch.distributed work without any changes to user code. Transformations at the XLA level would be an implementation detail and are not being proposed in this case.

The implementation is here. Do you see any issues with it?

[rpsilva-aws]

Yes, I agree with the point to tie torch.distribute with the XLA backend, making that an explicit association when using distributed process groups. We would still need to document this discrepancy if so. Users can still use xla_model to write send/recv outside of torch.distributed, which is a fine separation.

The concern still surrounds using XLA GPU references (with the known limitations above) as a motivation with explicitly lowering to a different CC op. I think it could be valuable to go over the check points above, just to double-check that this is the right path forward. I generally find CollectivePermute to capture the same P2P semantics, supporting both channel id (seen some cases with it on XLA) and not having ring constraints, so I was interested in capturing those details (e.g., if 2/ above were to land, do we switch back?). At least for the Pipeline Parallelism workstream, I have no concerns and was planning on using it for communication, but was targeting doing so on XLA/PJRT - I see JAX usually explicitly lowering to CollectivePermute for SPMD-only cases, and XLA GPU decomposing it to Send/Recv :) At the very least, if we go down this path, should we start with a warning for 1-2 versions, before making the change? I am not sure we have data points if any user is not using it as intended.

IIRC, for #8074, TRN actually worked fine, whereas GPU needed the "hidden attributes" - do you know if TPU has the same requirements?

I'll look at #9373 independently, thanks!

[bfolie]

IIRC, for #8074, TRN actually worked fine, whereas GPU needed the "hidden attributes" - do you know if TPU has the same requirements?

Yes, TPU requires the frontend attribute "_xla_send_recv_target_pairs". Is the current approach working for Neuron? And if so, where do you implement make_send_channel_id and make_recv_channel_id?

My sense is that this proposal is an internal, non-breaking change. And it's changing something that, at least on TPU, wasn't working. The language you're using, of multi-version user warnings and "hard to move back moving forward", is what I associate with breaking API changes. It implies that there is a contract about the mapping between operations on the torch frontend and XLA ops. Or that users might rely on the specific text of the IR, and that modifying it, even if the resulting calculation is the same, is a breaking change. Is that your impression?

Even if it's not a breaking change from a user's point of view it could still be a breaking change if Neuron relies on the specific IR. See this comment on 9373 for why I don't think it breaks Neuron functionally. But if it causes a performance degradation that's a reason to take another approach, so that is worth checking.

1/ Does JAX do this transformation implicitly?

As far as I can find, JAX doesn't expose send/recv operations.

2/ There's a chance that Send/Recv will no longer have this "hidden attribute", as GPU XLA looks into standardizing that spec... if 2 above were to land, do we switch back?

We could, especially if it's more performant (although if it's more performant then that would be an argument for using Send/Recv XLA ops now and inserting the frontend attribute, in spite of the additional complexity). Taking the perspective that this as an internal implementation detail, we're free to switch. If this is treated as a breaking change then I agree that it's a conversation is different.

[rpsilva-aws]

Sounds good. We don't require _xla_send_recv_target_pairs, with the channel id being sufficient. However, I've only ever tested for MPMD only. We don't use this to build anything else (or even for MPMD anymore) at the moment. I'll let you follow through with Kevin/Han over whether there is an implicit contract between torch frontend and XLA ops - but I agree that for the distributed XLA backend, there is more flexibility to argue for.

I have confirmed that it does not break Neuron functionality, so no concerns on our side. Thanks for checking all the other variants.