Error when usecollective communication in torch_xla.core.xla_model.all_to_all in SPMD mopde

[DarkenStar]

🐛 Bug

F ./torch_xla/csrc/runtime/debug_macros.h:20] Non-OK-status: status.status() status: INTERNAL: during context [hlo verifier]: RET_CHECK failure (external/xla/xla/service/hlo_verifier.cc:594) hlo->operand_count() == split_count 

Failed after pipeline-start
*** Begin stack trace ***
        tsl::CurrentStackTrace()
        std::unique_ptr<xla::PjRtLoadedExecutable, std::default_delete<xla::PjRtLoadedExecutable> > ConsumeValue<std::unique_ptr<xla::PjRtLoadedExecutable, std::default_delete<xla::PjRtLoadedExecutable> > >(absl::lts_20230125::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable, std::default_delete<xla::PjRtLoadedExecutable> > >&&)
        torch_xla::runtime::PjRtComputationClient::Compile(std::vector<torch_xla::runtime::ComputationClient::CompileInstance, std::allocator<torch_xla::runtime::ComputationClient::CompileInstance> >)
        torch_xla::XLAGraphExecutor::Compile(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&, absl::lts_20230125::Span<std::string const>, torch::lazy::LazyGraphExecutor::SyncTensorCollection const&, torch::lazy::LazyGraphExecutor::PostOrderData*, std::vector<torch::lazy::Value, std::allocator<torch::lazy::Value> > const&)
        torch_xla::XLAGraphExecutor::SyncTensorsGraphInternal(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> > > >*, absl::lts_20230125::Span<std::string const>, torch::lazy::LazyGraphExecutor::SyncTensorsConfig const&, bool)
        torch_xla::XLAGraphExecutor::SyncTensorsGraph(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> > > >*, absl::lts_20230125::Span<std::string const>, bool, bool, bool)
        torch_xla::XLATensor::ApplyPendingGraph()
        torch_xla::XLATensor::GetXlaData()
        torch_xla::XLATensor::ToTensor(bool)
        torch_xla::XLANativeFunctions::_to_copy(at::Tensor const&, c10::optional<c10::ScalarType>, c10::optional<c10::Layout>, c10::optional<c10::Device>, c10::optional<bool>, bool, c10::optional<c10::MemoryFormat>)

        at::_ops::_to_copy::redispatch(c10::DispatchKeySet, at::Tensor const&, c10::optional<c10::ScalarType>, c10::optional<c10::Layout>, c10::optional<c10::Device>, c10::optional<bool>, bool, c10::optional<c10::MemoryFormat>)

        at::_ops::_to_copy::call(at::Tensor const&, c10::optional<c10::ScalarType>, c10::optional<c10::Layout>, c10::optional<c10::Device>, c10::optional<bool>, bool, c10::optional<c10::MemoryFormat>)


        at::_ops::_to_copy::redispatch(c10::DispatchKeySet, at::Tensor const&, c10::optional<c10::ScalarType>, c10::optional<c10::Layout>, c10::optional<c10::Device>, c10::optional<bool>, bool, c10::optional<c10::MemoryFormat>)


        at::_ops::_to_copy::call(at::Tensor const&, c10::optional<c10::ScalarType>, c10::optional<c10::Layout>, c10::optional<c10::Device>, c10::optional<bool>, bool, c10::optional<c10::MemoryFormat>)
        at::native::to(at::Tensor const&, c10::optional<c10::ScalarType>, c10::optional<c10::Layout>, c10::optional<c10::Device>, c10::optional<bool>, bool, bool, c10::optional<c10::MemoryFormat>)

        at::_ops::to_dtype_layout::call(at::Tensor const&, c10::optional<c10::ScalarType>, c10::optional<c10::Layout>, c10::optional<c10::Device>, c10::optional<bool>, bool, bool, c10::optional<c10::MemoryFormat>)



        _PyEval_EvalFrameDefault
        _PyFunction_Vectorcall
        _PyEval_EvalFrameDefault
        _PyFunction_Vectorcall
        _PyEval_EvalFrameDefault




        _PyEval_EvalFrameDefault
        _PyFunction_Vectorcall
        _PyEval_EvalFrameDefault

        _PyEval_EvalFrameDefault
        _PyFunction_Vectorcall
        _PyEval_EvalFrameDefault
        _PyFunction_Vectorcall
        _PyEval_EvalFrameDefault
        _PyFunction_Vectorcall
        _PyEval_EvalFrameDefault
        _PyFunction_Vectorcall
        _PyEval_EvalFrameDefault
        _PyFunction_Vectorcall
        _PyEval_EvalFrameDefault
        _PyFunction_Vectorcall
        _PyEval_EvalFrameDefault
        _PyFunction_Vectorcall
        _PyEval_EvalFrameDefault
        _PyFunction_Vectorcall
        _PyEval_EvalFrameDefault





        _PyEval_EvalFrameDefault
        _PyFunction_Vectorcall
        _PyEval_EvalFrameDefault

        PyEval_EvalCode



        _PyRun_SimpleFileObject
        _PyRun_AnyFileObject
        Py_RunMain
        Py_BytesMain
        __libc_start_main

*** End stack trace ***

To Reproduce

Codes to reproduce the behavior:

import torch
import torch.nn as nn
import numpy as np
import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.runtime as xr]
import os

    os.environ['PJRT_DEVICE'] = 'CPU'
    os.environ['CPU_NUM_DEVICES'] = '5'
    mesh_shape = (5, 1)  # 5 个设备
    num_devices = xr.global_runtime_device_count()
    device_ids = np.array(range(num_devices))
    mesh = xs.Mesh(device_ids, mesh_shape, ('x', 'y'))
    B, N, C = 1, 5, 30  
    H = 10
    input_tensor = torch.randn(B, N, C).to(xm.xla_device())

    partition_spec = (None, 0, None)
    xs.mark_sharding(input_tensor, mesh, partition_spec)
    temp = xm.all_to_all(input_tensor, split_dimension=2, concat_dimension=1, split_count=num_devices,)
    hlo_text = torch_xla._XLAC._get_xla_tensors_hlo([temp])
    print(hlo_text)

Expected behavior

an all_to_all operation in mlir

Environment

• Reproducible on XLA backend [CPU/TPU/CUDA]: CPU
• torch_xla version: 2.1.0
• pytorch version: 2.1.0

Additional context

reduce_scatter .etc also produce bugs

[ysiraichi]

Thank you for filing this issue.
I'm not able to reproduce this error on nightly. Could you try a more recent version of PyTorch and PyTorch/XLA, say 2.6.0?

HloModule IrToHlo.22, entry_computation_layout={(f32[], f32[1,5,30]{2,1,0})->(f32[1,25,6]{2,1,0})}

ENTRY %IrToHlo.22 (p0.1: f32[], p1.2: f32[1,5,30]) -> (f32[1,25,6]) {
  %p0.1 = f32[] parameter(0), sharding={unknown}
  %p1.2 = f32[1,5,30]{2,1,0} parameter(1), sharding={devices=[1,5,1]0,1,2,3,4}
  %broadcast.3 = f32[1,5,30]{2,1,0} broadcast(f32[] %p0.1), dimensions={}
  %add.4 = f32[1,5,30]{2,1,0} add(f32[1,5,30]{2,1,0} %p1.2, f32[1,5,30]{2,1,0} %broadcast.3)
  %slice.5 = f32[1,5,6]{2,1,0} slice(f32[1,5,30]{2,1,0} %add.4), slice={[0:1], [0:5], [0:6]}
  %slice.6 = f32[1,5,6]{2,1,0} slice(f32[1,5,30]{2,1,0} %add.4), slice={[0:1], [0:5], [6:12]}
  %slice.7 = f32[1,5,6]{2,1,0} slice(f32[1,5,30]{2,1,0} %add.4), slice={[0:1], [0:5], [12:18]}
  %slice.8 = f32[1,5,6]{2,1,0} slice(f32[1,5,30]{2,1,0} %add.4), slice={[0:1], [0:5], [18:24]}
  %slice.9 = f32[1,5,6]{2,1,0} slice(f32[1,5,30]{2,1,0} %add.4), slice={[0:1], [0:5], [24:30]}
  %all-to-all.10 = (f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}) all-to-all(f32[1,5,6]{2,1,0} %slice.5, f32[1,5,6]{2,1,0} %slice.6, f32[1,5,6]{2,1,0} %slice.7, f32[1,5,6]{2,1,0} %slice.8, f32[1,5,6]{2,1,0} %slice.9), replica_groups={}, constrain_layout=true
  %get-tuple-element.11 = f32[1,5,6]{2,1,0} get-tuple-element((f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}) %all-to-all.10), index=0
  %get-tuple-element.12 = f32[1,5,6]{2,1,0} get-tuple-element((f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}) %all-to-all.10), index=1
  %get-tuple-element.13 = f32[1,5,6]{2,1,0} get-tuple-element((f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}) %all-to-all.10), index=2
  %get-tuple-element.14 = f32[1,5,6]{2,1,0} get-tuple-element((f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}) %all-to-all.10), index=3
  %get-tuple-element.15 = f32[1,5,6]{2,1,0} get-tuple-element((f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}, f32[1,5,6]{2,1,0}) %all-to-all.10), index=4
  %concatenate.16 = f32[1,25,6]{2,1,0} concatenate(f32[1,5,6]{2,1,0} %get-tuple-element.11, f32[1,5,6]{2,1,0} %get-tuple-element.12, f32[1,5,6]{2,1,0} %get-tuple-element.13, f32[1,5,6]{2,1,0} %get-tuple-element.14, f32[1,5,6]{2,1,0} %get-tuple-element.15), dimensions={1}
  %constant.17 = s32[] constant(0)
  %broadcast.18 = s32[3]{0} broadcast(s32[] %constant.17), dimensions={}
  %gather.19 = f32[] gather(f32[1,25,6]{2,1,0} %concatenate.16, s32[3]{0} %broadcast.18), offset_dims={}, collapsed_slice_dims={0,1,2}, start_index_map={0,1,2}, index_vector_dim=0, slice_sizes={1,1,1}
  %multiply.20 = f32[] multiply(f32[] %p0.1, f32[] %gather.19)
  ROOT %tuple.21 = (f32[1,25,6]{2,1,0}) tuple(f32[1,25,6]{2,1,0} %concatenate.16)
}