CollectivePermute support

[rpsilva-aws]

🐛 Bug

We currently discourage the use of CollectivePermute in #2384. There is no context behind why this is the case, including whether the motivation was hardware specific. It seems that we have enabled it for All-to-All (#2472) - do we have the same guidance/information from the XLA team?

I can not find any relevant reference to 'all_to_all_emitter'.

cc: @miladm @ManfeiBai @JackCaoG

[ManfeiBai]

Hi, @rpsilva-aws, thanks,

IIUC, all_to_all_emitter is defined internally for all_to_all, which helped fixed all_to_all unstable issue;

for collective_permute, we would need to sync with XLA team, cc @ddunl

btw, what's the current failure did you met with collective_permute now? any context or reproduce material?

[rpsilva-aws]

Thanks @ManfeiBai. I have not encountered any failure yet - at least with Neuron's TRN1, but this is a generally concerning docstring/call-out when trying to productionize with this collective on XLA. We need to use this instead of P2P send/recv for other HW specific reasons.

IIUC, all_to_all_emitter is defined internally for all_to_all, which helped fixed all_to_all unstable issue;

This is what I understood as well from Jack's PR above, but it would be nice if we had a reference point that we can use to cross check with other collectives (particularly CollectivePermute). I'll wait on the XLA team's comment.

[miladm]

@yaochengji - can you please share the latest technical updates on the support of this op outside of SPMD path?

[rpsilva-aws]

Any updates on whether there are still known XLA limitations with this op? cc: @yaochengji @ddunl

[yaochengji]

Hi @rpsilva-aws , thanks for asking.

Currently I have a PoC script https://github.com/pytorch/xla/blob/chengji/cm/test/torch_distributed/cm_perf.py to demostrate the CM support.

And the current main blocker is that sometimes in real workload, pin_layout should be set to True. The CM optimization in XLA compiler cannot handle this at the moment.

[rpsilva-aws]

Thanks for the context @yaochengji . Is this (ENABLE_COLLECTIVE_MATMUL_IN_MP) TPU specific? The #2384 PR was not clear if this was the case, and it hinted at a more fundamental issue with XLA that is hardware agnostic. If it's TPU specific, I'll reword the title but wanted to confirm. We do have similar einsum optimizations for Neuron, but iirc, that is separate from functionalizing CollectivePermute - perhaps we could add more context behind the docstr here.

[yaochengji]

Yes, ENABLE_COLLECTIVE_MATMUL_IN_MP is TPU specific, it will turn on some optimization flag in the tpu compiler.