Skip to content

Error/Segmentation fault for DeviceIterator with parallel DataLoader on sharded ReactantDevice #1716

Description

@benjwweber

I am trying to load a data set via a DataLoader/DeviceIterator to a sharded ReactantDevice. I have replicated it here with a mock data set. However I run into a variety of errors or segfaults if I do this.

cpu = cpu_device()
gpus = Reactant.devices()
ngpus = length(gpus)

mesh = Reactant.Sharding.Mesh(
    gpus,
    (:batch,)
)
device_replicated = reactant_device(
    sharding=Reactant.Sharding.Replicated(mesh)
)
device_sharded = reactant_device(
    sharding=Reactant.Sharding.DimsSharding(mesh, (-1,), (:batch,))
)

batchsize = 16 * ngpus
data = [
    (rand(Float32, 256, 256, 3, 2), rand(Bool, 16, 16, 3, 2), rand(Float32, 128, 128, 3, 4))
    for i in 1:batchsize * 100
] # mock data
dataloader = DataLoader(
    data; # dinoset;
    batchsize,
    shuffle=true,
    parallel=true,
    partial=false,
    collate=true
);

function f(dataloader, n;)
    times = [time(),]
    i = 0
    for inputs in dataloader
        i += 1
        @show i
        push!(times, time())
        if i >= n
            break
        end
    end
    return times
end

times = f(dataloader |> device_sharded, 100;)
diff(times)

This results in

julia> times = f(dataloader |> device_sharded, 100;)
i = 1
⋮
i = 25
i =
[1612749] signal 11 (128): Segmentation fault
in expression starting at REPL[50]:1
[1]    1612749 segmentation fault (core dumped)  taskset -c 11-60 julia --project -t 8

or

i = 1
⋮
i = 23
i =
[1614156] signal 11 (128): Segmentation fault
in expression starting at REPL[37]:1
_ZNK3xla4ifrt20SingleDeviceSharding4HashEN4absl12lts_202508149HashStateE at /home/bweber/.julia/artifacts/a4a178033ffe954af0f2e3ce566081781c945068/lib/libReactantExtra.so (unknown line)
ifrt_client_assemble_array_from_single_shards at /home/bweber/.julia/artifacts/a4a178033ffe954af0f2e3ce566081781c945068/lib/libReactantExtra.so (unknown line)
ifrt_client_assemble_array_from_single_shards at /home/bweber/.julia/packages/Reactant/puoGz/src/mlir/libMLIR_h.jl:14816 [inlined]
Array at /home/bweber/.julia/packages/Reactant/puoGz/src/xla/IFRT/Array.jl:105
Array at /home/bweber/.julia/packages/Reactant/puoGz/src/xla/IFRT/Array.jl:74
#AsyncArray#25 at /home/bweber/.julia/packages/Reactant/puoGz/src/xla/IFRT/AsyncArray.jl:8 [inlined]
AsyncArray at /home/bweber/.julia/packages/Reactant/puoGz/src/xla/IFRT/AsyncArray.jl:8 [inlined]
NamedSharding at /home/bweber/.julia/packages/Reactant/puoGz/src/Sharding.jl:481
unknown function (ip: 0x7f755dc49bfd) at (unknown file)
DimsSharding at /home/bweber/.julia/packages/Reactant/puoGz/src/Sharding.jl:727
#ConcreteIFRTArray#61 at /home/bweber/.julia/packages/Reactant/puoGz/src/Types.jl:388
ConcreteIFRTArray at /home/bweber/.julia/packages/Reactant/puoGz/src/Types.jl:378 [inlined]
macro expansion at /home/bweber/.julia/packages/MLDataDevices/Gr7bh/ext/ReactantExt.jl:109 [inlined]
#5 at /home/bweber/.julia/packages/Reactant/puoGz/src/Profiler.jl:295 [inlined]
#annotate#5 at /home/bweber/.julia/packages/Reactant/puoGz/src/Profiler.jl:270
annotate at /home/bweber/.julia/packages/Reactant/puoGz/src/Profiler.jl:262 [inlined]
annotate at /home/bweber/.julia/packages/Reactant/puoGz/src/Profiler.jl:262 [inlined]
adapt_storage at /home/bweber/.julia/packages/Reactant/puoGz/src/Profiler.jl:295 [inlined]
adapt_structure at /home/bweber/.julia/packages/Adapt/qhEjd/src/Adapt.jl:57 [inlined]
adapt at /home/bweber/.julia/packages/Adapt/qhEjd/src/Adapt.jl:40 [inlined]
#_#57 at ./operators.jl:1193 [inlined]
Fix at ./operators.jl:1193 [inlined]
ExcludeWalk at /home/bweber/.julia/packages/Functors/LbNAu/src/walks.jl:126 [inlined]
CachedWalk at /home/bweber/.julia/packages/Functors/LbNAu/src/walks.jl:177
recurse at /home/bweber/.julia/packages/Functors/LbNAu/src/walks.jl:54 [inlined]
map at ./tuple.jl:360 [inlined]
_map at /home/bweber/.julia/packages/Functors/LbNAu/src/walks.jl:3 [inlined]
DefaultWalk at /home/bweber/.julia/packages/Functors/LbNAu/src/walks.jl:73 [inlined]
ExcludeWalk at /home/bweber/.julia/packages/Functors/LbNAu/src/walks.jl:126 [inlined]
CachedWalk at /home/bweber/.julia/packages/Functors/LbNAu/src/walks.jl:177 [inlined]
execute at /home/bweber/.julia/packages/Functors/LbNAu/src/walks.jl:55 [inlined]
#fmap#32 at /home/bweber/.julia/packages/Functors/LbNAu/src/maps.jl:11 [inlined]
fmap at /home/bweber/.julia/packages/Functors/LbNAu/src/maps.jl:3 [inlined]
AbstractDevice at /home/bweber/.julia/packages/MLDataDevices/Gr7bh/src/public.jl:516
unknown function (ip: 0x7f755dc3b18d) at (unknown file)
#eachobsparallel##0 at /home/bweber/.julia/packages/MLDataDevices/Gr7bh/ext/MLUtilsExt.jl:35
unknown function (ip: 0x7f755dbc4a99) at (unknown file)
macro expansion at /home/bweber/.julia/packages/MLUtils/5jDrc/src/parallel.jl:127 [inlined]
#reducing_function#278#67 at /home/bweber/.julia/packages/FLoops/GN1Bm/src/reduce.jl:817 [inlined]
AdjoinIdentity at /home/bweber/.julia/packages/InitialValues/OWP8V/src/InitialValues.jl:306
next at /home/bweber/.julia/packages/Transducers/PqRuk/src/combinators.jl:290 [inlined]
next at /home/bweber/.julia/packages/Transducers/PqRuk/src/core.jl:287 [inlined]
macro expansion at /home/bweber/.julia/packages/Transducers/PqRuk/src/core.jl:181 [inlined]
macro expansion at /home/bweber/.julia/packages/Transducers/PqRuk/src/processes.jl:199 [inlined]
macro expansion at /home/bweber/.julia/packages/Transducers/PqRuk/src/simd.jl:41 [inlined]
_foldl_linear_bulk at /home/bweber/.julia/packages/Transducers/PqRuk/src/processes.jl:198 [inlined]
macro expansion at /home/bweber/.julia/packages/Transducers/PqRuk/src/processes.jl:192 [inlined]
macro expansion at /home/bweber/.julia/packages/Transducers/PqRuk/src/basics.jl:96 [inlined]
_foldl_array at /home/bweber/.julia/packages/Transducers/PqRuk/src/processes.jl:188 [inlined]
__foldl__ at /home/bweber/.julia/packages/Transducers/PqRuk/src/processes.jl:182 [inlined]
foldl_basecase at /home/bweber/.julia/packages/Transducers/PqRuk/src/processes.jl:361 [inlined]
_reduce_basecase at /home/bweber/.julia/packages/Transducers/PqRuk/src/threading_utils.jl:58
_reduce at /home/bweber/.julia/packages/Transducers/PqRuk/src/reduce.jl:139
_reduce at /home/bweber/.julia/packages/Transducers/PqRuk/src/reduce.jl:148
#_reduce##2 at /home/bweber/.julia/packages/Transducers/PqRuk/src/reduce.jl:147
unknown function (ip: 0x7f7bf946de3e) at (unknown file)
jl_apply at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/julia.h:2391 [inlined]
start_task at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/task.c:1252
Allocations: 120234999 (Pool: 120233519; Big: 1480); GC: 52
[1]    1614156 segmentation fault (core dumped)  taskset -c 11-60 julia --project -t 8

or

i = 1
⋮
i = 9
ERROR: ReadOnlyMemoryError()
Stacktrace:
 [1] check_channel_state
   @ ./channels.jl:187 [inlined]
 [2] take_buffered(c::Channel{Any})
   @ Base ./channels.jl:531
 [3] take!(c::Channel{Any})
   @ Base ./channels.jl:526
 [4] iterate(::MLUtils.Loader, state::MLUtils.LoaderState)
   @ MLUtils ~/.julia/packages/MLUtils/5jDrc/src/parallel.jl:143
 [5] iterate(c::DeviceIterator{typeof(identity), MLUtils.Loader}, ::Tuple{MLUtils.LoaderState, Tuple{…}})
   @ MLDataDevices ~/.julia/packages/MLDataDevices/33vyp/src/iterator.jl:61
 [6] f(dataloader::DeviceIterator{typeof(identity), MLUtils.Loader}, n::Int64)
   @ Main ./REPL[26]:11
 [7] top-level scope
   @ REPL[27]:1
Some type information was truncated. Use `show(err)` to see complete types.

I am unsure, what the underlaying problem is. Especially since the errors persist, even if the mock data fits on the device as a whole!

As it may be important the devices are:

julia> cpu = cpu_device()
(::CPUDevice{Missing}) (generic function with 1 method)

julia> gpus = Reactant.devices()
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1778503875.363267 1614630 pjrt_client.cc:569] PjRt-IFRT device count: total=1, addressable=1
I0000 00:00:1778503875.363350 1614630 pjrt_client.cc:573] Addressable PjRt-IFRT device: CpuDevice(id=0)
I0000 00:00:1778503876.431877 1614630 service.cc:154] XLA service 0x30ee8dd0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1778503876.432011 1614630 service.cc:170]   StreamExecutor [0]: NVIDIA A100-SXM4-80GB, Compute Capability 8.0 (Driver: 13.0.0[570.211.1]; Runtime: 13.1.0; Toolkit: 13.1.0; DNN: 9.14.0)
I0000 00:00:1778503876.432033 1614630 service.cc:170]   StreamExecutor [1]: NVIDIA A100-SXM4-80GB, Compute Capability 8.0 (Driver: 13.0.0[570.211.1]; Runtime: 13.1.0; Toolkit: 13.1.0; DNN: 9.14.0)
I0000 00:00:1778503876.432047 1614630 service.cc:170]   StreamExecutor [2]: NVIDIA A100-SXM4-80GB, Compute Capability 8.0 (Driver: 13.0.0[570.211.1]; Runtime: 13.1.0; Toolkit: 13.1.0; DNN: 9.14.0)
I0000 00:00:1778503876.432058 1614630 service.cc:170]   StreamExecutor [3]: NVIDIA A100-SXM4-80GB, Compute Capability 8.0 (Driver: 13.0.0[570.211.1]; Runtime: 13.1.0; Toolkit: 13.1.0; DNN: 9.14.0)
I0000 00:00:1778503876.439964 1614630 se_gpu_pjrt_client.cc:1540] Using BFC allocator.
I0000 00:00:1778503876.440052 1614630 gpu_helpers.cc:144] XLA backend will use up to 85095874560 bytes on device 0 for BFCAllocator.
I0000 00:00:1778503876.440125 1614630 gpu_helpers.cc:144] XLA backend will use up to 85095874560 bytes on device 1 for BFCAllocator.
I0000 00:00:1778503876.440171 1614630 gpu_helpers.cc:144] XLA backend will use up to 85095874560 bytes on device 2 for BFCAllocator.
I0000 00:00:1778503876.440214 1614630 gpu_helpers.cc:144] XLA backend will use up to 85095874560 bytes on device 3 for BFCAllocator.
I0000 00:00:1778503876.440258 1614630 gpu_helpers.cc:183] XLA backend will use up to 0 bytes on device 0 for CollectiveBFCAllocator.
I0000 00:00:1778503876.440301 1614630 gpu_helpers.cc:183] XLA backend will use up to 0 bytes on device 1 for CollectiveBFCAllocator.
I0000 00:00:1778503876.440346 1614630 gpu_helpers.cc:183] XLA backend will use up to 0 bytes on device 2 for CollectiveBFCAllocator.
I0000 00:00:1778503876.440382 1614630 gpu_helpers.cc:183] XLA backend will use up to 0 bytes on device 3 for CollectiveBFCAllocator.
I0000 00:00:1778503876.472828 1614630 cuda_dnn.cc:461] Loaded cuDNN version 91400
I0000 00:00:1778503876.473744 1614630 pjrt_client.cc:569] PjRt-IFRT device count: total=4, addressable=4
I0000 00:00:1778503876.473768 1614630 pjrt_client.cc:573] Addressable PjRt-IFRT device: CudaDevice(id=0)
I0000 00:00:1778503876.473777 1614630 pjrt_client.cc:573] Addressable PjRt-IFRT device: CudaDevice(id=1)
I0000 00:00:1778503876.473784 1614630 pjrt_client.cc:573] Addressable PjRt-IFRT device: CudaDevice(id=2)
I0000 00:00:1778503876.473791 1614630 pjrt_client.cc:573] Addressable PjRt-IFRT device: CudaDevice(id=3)
4-element Vector{Reactant.XLA.IFRT.Device}:
 Reactant.XLA.IFRT.Device(Ptr{Nothing}(0x00000000303f1650), "CUDA:0 NVIDIA A100-SXM4-80GB")
 Reactant.XLA.IFRT.Device(Ptr{Nothing}(0x0000000032602540), "CUDA:1 NVIDIA A100-SXM4-80GB")
 Reactant.XLA.IFRT.Device(Ptr{Nothing}(0x0000000032601ca0), "CUDA:2 NVIDIA A100-SXM4-80GB")
 Reactant.XLA.IFRT.Device(Ptr{Nothing}(0x0000000032601eb0), "CUDA:3 NVIDIA A100-SXM4-80GB")

julia> ngpus = length(gpus)
4

julia> mesh = Reactant.Sharding.Mesh(
           gpus,
           (:batch,)
       )
Reactant.Sharding.Mesh{1, Vector{Int64}}([0, 1, 2, 3], [0, 1, 2, 3], (:batch,), (4,))

julia> device_replicated = reactant_device(
           sharding=Reactant.Sharding.Replicated(mesh)
       )
(::ReactantDevice{Missing, Missing, Reactant.Sharding.Replicated{Reactant.Sharding.Mesh{1, Vector{Int64}}}, Missing, Union{}}) (generic function with 1 method)

julia> device_sharded = reactant_device(
           sharding=Reactant.Sharding.DimsSharding(mesh, (-1,), (:batch,))
       )
(::ReactantDevice{Missing, Missing, Reactant.Sharding.DimsSharding{1, Tuple{Symbol}, Reactant.Sharding.Mesh{1, Vector{Int64}}}, Missing, Union{}}) (generic function with 1 method)

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions