EnzymeAD / EnzymeAD/Reactant.jl
VRAM usage is limited to 75%
- Dominant language
- Julia
- Stars
- 370
- Forks
- 74
- Avg merge
- 18h 47m
- Merged PRs (30d)
- 30
Description
Follow the instructions in https://github.com/EnzymeAD/Reactant.jl/issues/877, I set `XLA_REACTANT_GPU_MEM_FRACTION=0.9` and indeed VRAM prelocates 90% of total VRAM (32GB), which can be verified by `nvtop`. `@compile` is successful but throws message
```bash
W external/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc:3021] Can't reduce memory use below 23.06GiB (24764946267 bytes) by rematerialization; only reduced to 24.56GiB (26374492544 bytes), down from 88.47GiB (94992687120 bytes) originally
856.173470 seconds (87.89 M allocations: 4.356 GiB, 0.17% gc time, 91.68% compilation time: <1% of which was recompilation)
```
executing the compiled function leads to out of memory error
```bash
W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 87.49GiB (rounded to 93936833536)requested by op
... lengthy message...
ERROR: LoadError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 93936833408 bytes. [tf-allocator-allocation-error='']
....
```
The julia code:
```julia
using Combinatorics, EinExprs, Tenet, Reactant, Adapt
Reactant.set_default_backend("gpu")
Nv,Ns,Nt,Nc = 200,48,96,3
Nx = Ns^3
clrs = permutations([1, 2, 3])
ϵ = (a, b, c) -> (a - b) * (b - c) * (c - a) / 2
ph = randn(ComplexF64, Nx)
vec = randn(ComplexF64,(Nv, Nx, Nc))
const batch = 12*12
const Nbatch = Nx ÷ batch
Tph = Tensor(ph[1:Nbatch], [:x])
vec1 = Tensor(vec[ :, 1:Nbatch, 1],[:i,:x])
vec2 = Tensor(vec[ :, 1:Nbatch, 2],[:j,:x])
vec3 = Tensor(vec[ :, 1:Nbatch, 3],[:k,:x])
TN = TensorNetwork([Tph, vec1, vec2, vec3])
path_opt = einexpr(TN; optimizer= Greedy())
Tph = nothing
vec1 = nothing
vec2 = nothing
vec3 = nothing
TN = nothing
function Rfunc(vec)
sum_tmp = zeros(ComplexF64,(Nv,Nv, Nv))
for (a,b,c) in clrs
ϵ_abc = ϵ(a,b,c)
for i in 0:batch-1
tph = Tensor(ph[i*Nbatch+1 : (i+1)*Nbatch], [:x])
vec_a = Tensor(vec[ :, i*Nbatch+1 : (i+1)*Nbatch, a],[:i,:x])
vec_b = Tensor(vec[ :, i*Nbatch+1 : (i+1)*Nbatch, b],[:j,:x])
vec_c = Tensor(vec[ :, i*Nbatch+1 : (i+1)*Nbatch, c],[:k,:x])
tn = TensorNetwork([tph, vec_a, vec_b, vec_c])
sum_tmp = sum_tmp + ϵ_abc * contract(tn; path=path_opt).data
end
end
return sum_tmp
end
Rvec = adapt(ConcreteRArray,vec)
@time cRfunc = @compile Rfunc(Rvec)
cRfunc(Rvec)
```
The full output:
```log
2025-04-27 21:29:20.940201: I external/xla/xla/service/service.cc:152] XLA service 0x1e45fe0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2025-04-27 21:29:20.940231: I external/xla/xla/service/service.cc:160] StreamExecutor device (0): Tesla V100-PCIE-32GB, Compute Capability 7.0
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1745760560.941089 37434 se_gpu_pjrt_client.cc:1038] Using BFC allocator.
I0000 00:00:1745760560.941158 37434 gpu_helpers.cc:136] XLA backend allocating 30665303654 bytes on device 0 for BFCAllocator.
I0000 00:00:1745760560.941189 37434 gpu_helpers.cc:177] XLA backend will use up to 3407255961 bytes on device 0 for CollectiveBFCAllocator.
I0000 00:00:1745760560.945638 37434 cuda_dnn.cc:527] Loaded cuDNN version 90400
2025-04-27 21:43:31.565916: W external/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc:3021] Can't reduce memory use below 23.06GiB (24764946267 bytes) by rematerialization; only reduced to 24.56GiB (26374492544 bytes), down from 88.47GiB (94992687120 bytes) originally
850.440218 seconds (87.89 M allocations: 4.356 GiB, 0.17% gc time, 91.59% compilation time: <1% of which was recompilation)
2025-04-27 21:43:49.785907: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 87.49GiB (rounded to 93936833536)requested by op
2025-04-27 21:43:49.785979: I external/xla/xla/tsl/framework/bfc_allocator.cc:1058] BFCAllocator dump for GPU_0_bfc
2025-04-27 21:43:49.786003: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (256): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786024: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (512): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786041: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (1024): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786058: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (2048): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786077: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (4096): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786094: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (8192): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786112: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (16384): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786129: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (32768): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786146: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (65536): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786163: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (131072): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786181: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (262144): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786198: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (524288): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786215: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (1048576): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786232: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (2097152): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786250: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (4194304): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786269: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (8388608): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786287: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (16777216): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786305: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (33554432): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786330: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (67108864): Total Chunks: 1, Chunks in use: 1. 122.07MiB allocated for chunks. 122.07MiB in use in bin. 122.07MiB client-requested in use in bin.
2025-04-27 21:43:49.786336: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (134217728): Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2025-04-27 21:43:49.786344: I external/xla/xla/tsl/framework/bfc_allocator.cc:1065] Bin (268435456): Total Chunks: 2, Chunks in use: 1. 28.44GiB allocated for chunks. 1012.50MiB in use in bin. 1012.50MiB client-requested in use in bin.
2025-04-27 21:43:49.786354: I external/xla/xla/tsl/framework/bfc_allocator.cc:1081] Bin for 87.49GiB was 256.00MiB, Chunk State:
2025-04-27 21:43:49.786387: I external/xla/xla/tsl/framework/bfc_allocator.cc:1087] Size: 27.45GiB | Requested Size: 18.34MiB | in_use: 0 | bin_num: 20, prev: Size: 122.07MiB | Requested Size: 122.07MiB | in_use: 1 | bin_num: -1
2025-04-27 21:43:49.786394: I external/xla/xla/tsl/framework/bfc_allocator.cc:1094] Next region of size 30665303552
2025-04-27 21:43:49.786402: I external/xla/xla/tsl/framework/bfc_allocator.cc:1114] InUse at 7fe69c000000 of size 1061683200 next 1
2025-04-27 21:43:49.786407: I external/xla/xla/tsl/framework/bfc_allocator.cc:1114] InUse at 7fe6db480000 of size 128000000 next 5
2025-04-27 21:43:49.786423: I external/xla/xla/tsl/framework/bfc_allocator.cc:1114] Free at 7fe6e2e92000 of size 29475620352 next 18446744073709551615
2025-04-27 21:43:49.786428: I external/xla/xla/tsl/framework/bfc_allocator.cc:1119] Summary of in-use Chunks by size:
2025-04-27 21:43:49.786434: I external/xla/xla/tsl/framework/bfc_allocator.cc:1122] 1 Chunks of size 128000000 totalling 122.07MiB
2025-04-27 21:43:49.786441: I external/xla/xla/tsl/framework/bfc_allocator.cc:1122] 1 Chunks of size 1061683200 totalling 1012.50MiB
2025-04-27 21:43:49.786447: I external/xla/xla/tsl/framework/bfc_allocator.cc:1126] Sum Total of in-use chunks: 1.11GiB
2025-04-27 21:43:49.786452: I external/xla/xla/tsl/framework/bfc_allocator.cc:1128] Total bytes in pool: 30665303552 memory_limit_: 30665303654 available bytes: 102 curr_region_allocation_bytes_: 61330607616
2025-04-27 21:43:49.786462: I external/xla/xla/tsl/framework/bfc_allocator.cc:1133] Stats:
Limit: 30665303654
InUse: 1189683200
MaxInUse: 2758770688
NumAllocs: 18
MaxAllocSize: 1061683200
Reserved: 0
PeakReserved: 0
LargestFreeBlock: 0
2025-04-27 21:43:49.786474: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] ****________________________________________________________________________________________________
E0000 00:00:1745761429.786500 37434 pjrt_stream_executor_client.cc:3024] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 93936833408 bytes. [tf-allocator-allocation-error='']
ERROR: LoadError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 93936833408 bytes. [tf-allocator-allocation-error='']
Stacktrace:
[1] reactant_err(msg::Cstring)
@ Reactant.XLA ~/.julia/packages/Reactant/WNy1L/src/xla/Utils.jl:12
[2] macro expansion
@ ~/.julia/packages/Reactant/WNy1L/src/xla/PJRT/LoadedExecutable.jl:195 [inlined]
[3] execute_sharded
@ ~/.julia/packages/Reactant/WNy1L/src/xla/PJRT/LoadedExecutable.jl:164 [inlined]
[4] macro expansion
@ ~/.julia/packages/Reactant/WNy1L/src/Compiler.jl:2452 [inlined]
[5] (::Reactant.Compiler.Thunk{typeof(Rfunc), Symbol("##Rfunc_reactant#5421"), false, Tuple{ConcretePJRTArray{ComplexF64, 3, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}}, Reactant.XLA.PJRT.LoadedExecutable, Reactant.XLA.PJRT.Device, Reactant.XLA.PJRT.Client, Tuple{}, Vector{Bool}})(args::ConcretePJRTArray{ComplexF64, 3, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}})
@ Reactant.Compiler ~/.julia/packages/Reactant/WNy1L/src/Compiler.jl:2880
[6] top-level scope
@ ~/issue.jl:56
in expression starting at /home/chimaoshu/issue.jl:56
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.