torch.distributed.all_reduce not converted to stableHLO

[AleksKnezevic]

🐛 Bug

As the title suggests, torch.distributed.all_reduce is not being converted to stableHLO.

To Reproduce

I run the following test:

import os
import torch
from torch import nn
import torch
from torch.export import export
from torch_xla.stablehlo import exported_program_to_stablehlo


def test():
    class Basic(nn.Module):
        def __init__(self):
            super().__init__()

        def forward(self, x):
            return torch.distributed.all_reduce(x, op=torch.distributed.ReduceOp.SUM)

    model = Basic()
    prog = export(model, (torch.rand(20, 10), ))
    shlo = exported_program_to_stablehlo(prog)
    print(shlo.get_stablehlo_text())

os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "29501"
torch.distributed.init_process_group(world_size=1, rank=0)


if __name__ == "__main__":
    test()

Expected behavior

I would expect a stableHLO module with all_reduce, however I get the following error:

WARNING:root:Defaulting to PJRT_DEVICE=CPU
loc("all-reduce.10"): error: failed to legalize operation 'mhlo.all_reduce' that was explicitly marked illegal
[rank0]: Traceback (most recent call last):
[rank0]:   File "/localdev/aknezevic/xt/test_mp.py", line 28, in <module>
[rank0]:     test()
[rank0]:   File "/localdev/aknezevic/xt/test_mp.py", line 19, in test
[rank0]:     shlo = exported_program_to_stablehlo(prog)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/localdev/aknezevic/xt/venv/lib/python3.11/site-packages/torch_xla/stablehlo.py", line 626, in exported_program_to_stablehlo
[rank0]:     bundle = _exported_program_to_stablehlo_bundle(exported_model, options)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/localdev/aknezevic/xt/venv/lib/python3.11/site-packages/torch_xla/stablehlo.py", line 405, in _exported_program_to_stablehlo_bundle
[rank0]:     stablehlo_content = xm.get_stablehlo_bytecode(res)
[rank0]:                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/localdev/aknezevic/xt/venv/lib/python3.11/site-packages/torch_xla/core/xla_model.py", line 1103, in get_stablehlo_bytecode
[rank0]:     return torch_xla._XLAC._get_stablehlo(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: RuntimeError: torch_xla/csrc/runtime/stablehlo_helper.cc:109 : Check failed: status.ok() 
[rank0]: *** Begin stack trace ***
[rank0]:        tsl::CurrentStackTrace()
[rank0]:        torch_xla::ConvertHloToStableHlo(xla::HloModuleProto const*, mlir::ModuleOp*)
[rank0]:        torch_xla::hloToStablehlo(xla::HloModuleProto const*, bool)
[rank0]:        torch_xla::DumpUtil::ToHlo(c10::ArrayRef<torch::lazy::Value>, torch::lazy::BackendDevice const&, torch_xla::EmitMode)
[rank0]:        torch_xla::XLAGraphExecutor::DumpHloComputation(std::vector<c10::intrusive_ptr<torch_xla::XLATensor, c10::detail::intrusive_target_default_null_type<torch_xla::XLATensor> >, std::allocator<c10::intrusive_ptr<torch_xla::XLATensor, c10::detail::intrusive_target_default_null_type<torch_xla::XLATensor> > > > const&, torch_xla::EmitMode)
[rank0]: 
[rank0]: 
[rank0]: 
[rank0]: 
[rank0]:        _PyObject_MakeTpCall
[rank0]:        _PyEval_EvalFrameDefault
[rank0]: 
[rank0]:        PyEval_EvalCode
[rank0]: 
[rank0]: 
[rank0]: 
[rank0]:        _PyRun_SimpleFileObject
[rank0]:        _PyRun_AnyFileObject
[rank0]:        Py_RunMain
[rank0]:        Py_BytesMain
[rank0]: 
[rank0]:        __libc_start_main
[rank0]:        _start
[rank0]: *** End stack trace ***
[rank0]: MHLO -> StableHLO conversion failed.
[rank0]: StableHLO Module from MHLO -> StableHLO conversion is not leagal.Please open a github issue to PyTorch/XLA.
[rank0]: Original HLO dump:
[rank0]: HloModule IrToHlo.14, entry_computation_layout={(f32[], f32[20,10]{1,0})->(f32[20,10]{1,0}, f32[20,10]{1,0})}

[rank0]: %AddComputation.6 (x.7: f32[], y.8: f32[]) -> f32[] {
[rank0]:   %x.7 = f32[] parameter(0)
[rank0]:   %y.8 = f32[] parameter(1)
[rank0]:   ROOT %add.9 = f32[] add(f32[] %x.7, f32[] %y.8)
[rank0]: }

[rank0]: ENTRY %IrToHlo.14 (p0.1: f32[], p1.2: f32[20,10]) -> (f32[20,10], f32[20,10]) {
[rank0]:   %p1.2 = f32[20,10]{1,0} parameter(1)
[rank0]:   %p0.1 = f32[] parameter(0)
[rank0]:   %tuple.3 = (f32[20,10]{1,0}, f32[]) tuple(f32[20,10]{1,0} %p1.2, f32[] %p0.1)
[rank0]:   %get-tuple-element.4 = f32[20,10]{1,0} get-tuple-element((f32[20,10]{1,0}, f32[]) %tuple.3), index=0
[rank0]:   %get-tuple-element.5 = f32[] get-tuple-element((f32[20,10]{1,0}, f32[]) %tuple.3), index=1
[rank0]:   %all-reduce.10 = (f32[20,10]{1,0}, f32[]) all-reduce(f32[20,10]{1,0} %get-tuple-element.4, f32[] %get-tuple-element.5), replica_groups={}, constrain_layout=true, to_apply=%AddComputation.6
[rank0]:   %get-tuple-element.12 = f32[] get-tuple-element((f32[20,10]{1,0}, f32[]) %all-reduce.10), index=1
[rank0]:   %get-tuple-element.11 = f32[20,10]{1,0} get-tuple-element((f32[20,10]{1,0}, f32[]) %all-reduce.10), index=0
[rank0]:   ROOT %tuple.13 = (f32[20,10]{1,0}, f32[20,10]{1,0}) tuple(f32[20,10]{1,0} %get-tuple-element.11, f32[20,10]{1,0} %get-tuple-element.11)
[rank0]: }

Environment

• Reproducible on XLA backend CPU:
• torch_xla version 2.5.0 and 2.6.0 (I tried both):

[ysiraichi]

Thank you for filing this issue. I can confirm this still happens on 76b0ce5.
cc @tengyifei @bhavya01

[iwknow]

i feel this might not be a bug because you are asking torch_xla to compile a distributed operation for a non-distributed, single CPU target, and its compiler pathway doesn't support (or "legalize") that specific operation for that target. torch.distributed operations like all_reduce are inherently designed for multi-process or multi-device scenarios. While torch.export captures the operation in the graph, the torch_xla backend targeting a single CPU via PJRT might not have a meaningful or implemented way to execute this collective operation.

Try to set PJRT_DEVICE=CUDA or PJRT_DEVICE=TPU to see if that works.

[AleksKnezevic]

@iwknow I just tested on TPU and got the same error. Other ccl ops work and are able to generate a stableHLO graph.

[AleksKnezevic]

I've Looking at this a little more, it fails in conversion from MHLO to stableHLO. It seems related to this xla issue and this stable_hlo issue which @GleasonK marked as complete.

Is it simply an issue of removing this check in mlir-hlo?

EDIT
I guess the actual code that needs to be updated lives in open_xla

[GleasonK]

Have a fix in progress. Removing that line resolves the issue for all_reduce. Still hitting an issue on tuple all_to_all, will get that sorted out shortly!

[AleksKnezevic]

Awesome, thanks @GleasonK! Do you have an ETA?

[ghpvnist]

To give additional context, there's a semantic difference in StableHLO's variadic all to all and HLO's tuple all to all. We should resolve the all_reduce issue very soon, and will investigate further on all to all. Thanks for flagging this!

[GleasonK]

Awesome, thanks @GleasonK! Do you have an ETA?

Change is linked (see copybara notes before this comment), should be in by tomorrow. Then PT/XLA will need to bump their XLA revision, not sure how frequently this happens but I believe only a few days at most.

[GleasonK]

And let me know if you hit the variadic-all-to-all issue. That'll be more effort to fix. Our specced version of a2a_v diverges from the XlaBuilder impl (probably incorrectly on our end), so we'll need to figure out proper path forward.