You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
When a PDESystem has parameters (that aren't being estimated), solve fails if trying to solve on the GPU.
The below MRE shows a system without parameters being solved on the GPU, an identical system with parameters being solved on the CPU, but when I try solving the system with parameters on the GPU it errors in solve.
Expected behavior
Your PDESystem should be able to have parameters whether you solve on the CPU or the GPU.
Minimal Reproducible Example 👇
using Lux, LuxCUDA, Optimization, OptimizationOptimisers, Random, ComponentArrays
using NeuralPDE
import ModelingToolkit: Interval, infimum, supremum
const gpud =gpu_device()
# Based on the 1D PDE Dirichlet BC - CUDA test case in NeuralPDE.jl/test/NNPDE_cuda_tests.jl############################### Without parameters, With GPU ###############################
Random.seed!(100)
@parameters t x
@variablesu(..)
Dt =Differential(t)
Dxx =Differential(x)^2
eq =Dt(u(t, x)) ~Dxx(u(t, x))
bcs = [
u(0, x) ~cos(x),
u(t, 0) ~exp(-t),
u(t, 1) ~exp(-t) *cos(1)
]
domains = [t ∈Interval(0.0, 1.0), x ∈Interval(0.0, 1.0)]
@named pdesys =PDESystem(eq, bcs, domains, [t, x], [u(t, x)])
inner =30
chain =Chain(Dense(2, inner, σ), Dense(inner, inner, σ),
Dense(inner, inner, σ), Dense(inner, inner, σ),
Dense(inner, inner, σ), Dense(inner, inner, σ), Dense(inner, 1))
strategy =StochasticTraining(500)
ps = Lux.initialparameters(Random.default_rng(), chain) |> ComponentArray |> gpud |> f64
discretization =PhysicsInformedNN(chain, strategy; init_params = ps)
prob =discretize(pdesys, discretization)
res =solve(prob, OptimizationOptimisers.Adam(0.01); maxiters =1000)
prob =remake(prob, u0 = res.u)
res =solve(prob, OptimizationOptimisers.Adam(0.001); maxiters =1000)
phi = discretization.phi
u_exact = (t, x) ->exp.(-t) *cos.(x)
ts, xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
u_predict = [first(Array(phi([t, x], res.u))) for t in ts for x in xs]
u_real = [u_exact(t, x) for t in ts for x in xs]
diff_u =abs.(u_predict .- u_real)
@assertmaximum(diff_u) <0.01############################### With Parameters, Without GPU ###############################
Random.seed!(100)
@parameters t x
@variablesu(..)
@parameters k
Dt =Differential(t)
Dxx =Differential(x)^2
eq =Dt(u(t, x)) ~Dxx(u(t, x))
bcs = [
u(0, x) ~cos(x),
u(t, 0) ~exp(-k * t),
u(t, 1) ~exp(-k * t) *cos(1)
]
domains = [t ∈Interval(0.0, 1.0), x ∈Interval(0.0, 1.0)]
defaults =Dict([k =>1.0])
@named pdesys =PDESystem(eq, bcs, domains, [t, x], [u(t, x)], [k]; defaults)
inner =30
chain =Chain(Dense(2, inner, σ), Dense(inner, inner, σ),
Dense(inner, inner, σ), Dense(inner, inner, σ),
Dense(inner, inner, σ), Dense(inner, inner, σ), Dense(inner, 1))
strategy =StochasticTraining(500)
discretization =PhysicsInformedNN(chain, strategy)
prob =discretize(pdesys, discretization)
res =solve(prob, OptimizationOptimisers.Adam(0.01); maxiters =1000)
prob =remake(prob, u0 = res.u)
res =solve(prob, OptimizationOptimisers.Adam(0.001); maxiters =1000)
phi = discretization.phi
u_exact = (t, x) ->exp.(-t) *cos.(x)
ts, xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
u_predict = [first(Array(phi([t, x], res.u))) for t in ts for x in xs]
u_real = [u_exact(t, x) for t in ts for x in xs]
diff_u =abs.(u_predict .- u_real)
@assertmaximum(diff_u) <0.01################################ With Parameters, With GPU #################################
Random.seed!(100)
@parameters t x
@variablesu(..)
@parameters k
Dt =Differential(t)
Dxx =Differential(x)^2
eq =Dt(u(t, x)) ~Dxx(u(t, x))
bcs = [
u(0, x) ~cos(x),
u(t, 0) ~exp(-k * t),
u(t, 1) ~exp(-k * t) *cos(1)
]
domains = [t ∈Interval(0.0, 1.0), x ∈Interval(0.0, 1.0)]
defaults =Dict([k =>1.0])
@named pdesys =PDESystem(eq, bcs, domains, [t, x], [u(t, x)], [k]; defaults)
inner =30
chain =Chain(Dense(2, inner, σ), Dense(inner, inner, σ),
Dense(inner, inner, σ), Dense(inner, inner, σ),
Dense(inner, inner, σ), Dense(inner, inner, σ), Dense(inner, 1))
strategy =StochasticTraining(500)
ps = Lux.initialparameters(Random.default_rng(), chain) |> ComponentArray |> gpud |> f64
discretization =PhysicsInformedNN(chain, strategy; init_params = ps)
prob =discretize(pdesys, discretization)
res =solve(prob, OptimizationOptimisers.Adam(0.01); maxiters =1000)
prob =remake(prob, u0 = res.u)
res =solve(prob, OptimizationOptimisers.Adam(0.001); maxiters =1000)
phi = discretization.phi
u_exact = (t, x) ->exp.(-t) *cos.(x)
ts, xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
u_predict = [first(Array(phi([t, x], res.u))) for t in ts for x in xs]
u_real = [u_exact(t, x) for t in ts for x in xs]
diff_u =abs.(u_predict .- u_real)
@assertmaximum(diff_u) <0.01
Error & Stacktrace ⚠️
ERROR: GPUCompiler.KernelError(GPUCompiler.CompilerJob{GPUCompiler.PTXCompilerTarget, CUDA.CUDACompilerParams}(MethodInstance for (::GPUArrays.var"#35#37")(::CUDA.CuKernelContext, ::CuDeviceMatrix{…}, ::Base.Broadcast.Broadcasted{…}, ::Int64), GPUCompiler.CompilerConfig{GPUCompiler.PTXCompilerTarget, CUDA.CUDACompilerParams}(GPUCompiler.PTXCompilerTarget(v"8.6.0", v"7.8.0", true, nothing, nothing, nothing, nothing, false, nothing, nothing), CUDA.CUDACompilerParams(v"8.6.0", v"8.5.0"), true, nothing, :specfunc, false, 2), 0x0000000000006a24), "passing and using non-bitstype argument", "Argument 4 to your kernel function is of type Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{2, CUDA.DeviceMemory}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, typeof(*), Tuple{Base.Broadcast.Extruded{Vector{Float64}, Tuple{Bool}, Tuple{Int64}}, Base.Broadcast.Extruded{CuDeviceMatrix{Float64, 1}, Tuple{Bool, Bool}, Tuple{Int64, Int64}}}}, which is not isbits:\n .args is of type Tuple{Base.Broadcast.Extruded{Vector{Float64}, Tuple{Bool}, Tuple{Int64}}, Base.Broadcast.Extruded{CuDeviceMatrix{Float64, 1}, Tuple{Bool, Bool}, Tuple{Int64, Int64}}} which is not isbits.\n .1 is of type Base.Broadcast.Extruded{Vector{Float64}, Tuple{Bool}, Tuple{Int64}} which is not isbits.\n .x is of type Vector{Float64} which is not isbits.\n .ref is of type MemoryRef{Float64} which is not isbits.\n .mem is of type Memory{Float64} which is not isbits.\n", Base.StackTraces.StackFrame[])
Stacktrace:
[1] check_invocation(job::GPUCompiler.CompilerJob)
@ GPUCompiler ~/.julia/packages/GPUCompiler/2CW9L/src/validation.jl:92
[2] macro expansion
@ ~/.julia/packages/GPUCompiler/2CW9L/src/driver.jl:92 [inlined]
[3] macro expansion
@ ~/.julia/packages/TimerOutputs/NRdsv/src/TimerOutput.jl:253 [inlined]
[4]
@ GPUCompiler ~/.julia/packages/GPUCompiler/2CW9L/src/driver.jl:90
[5] codegen
@ ~/.julia/packages/GPUCompiler/2CW9L/src/driver.jl:82 [inlined]
[6] compile(target::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
@ GPUCompiler ~/.julia/packages/GPUCompiler/2CW9L/src/driver.jl:79
[7] compile
@ ~/.julia/packages/GPUCompiler/2CW9L/src/driver.jl:74 [inlined]
[8] #1145
@ ~/.julia/packages/CUDA/2kjXI/src/compiler/compilation.jl:250 [inlined]
[9] JuliaContext(f::CUDA.var"#1145#1148"{GPUCompiler.CompilerJob{…}}; kwargs::@Kwargs{})
@ GPUCompiler ~/.julia/packages/GPUCompiler/2CW9L/src/driver.jl:34
[10] JuliaContext(f::Function)
@ GPUCompiler ~/.julia/packages/GPUCompiler/2CW9L/src/driver.jl:25
[11] compile(job::GPUCompiler.CompilerJob)
@ CUDA ~/.julia/packages/CUDA/2kjXI/src/compiler/compilation.jl:249
[12] actual_compilation(cache::Dict{…}, src::Core.MethodInstance, world::UInt64, cfg::GPUCompiler.CompilerConfig{…}, compiler::typeof(CUDA.compile), linker::typeof(CUDA.link))
@ GPUCompiler ~/.julia/packages/GPUCompiler/2CW9L/src/execution.jl:237
[13] cached_compilation(cache::Dict{…}, src::Core.MethodInstance, cfg::GPUCompiler.CompilerConfig{…}, compiler::Function, linker::Function)
@ GPUCompiler ~/.julia/packages/GPUCompiler/2CW9L/src/execution.jl:151
[14] macro expansion
@ ~/.julia/packages/CUDA/2kjXI/src/compiler/execution.jl:380 [inlined]
[15] macro expansion
@ ./lock.jl:273 [inlined]
[16] cufunction(f::GPUArrays.var"#35#37", tt::Type{Tuple{…}}; kwargs::@Kwargs{})
@ CUDA ~/.julia/packages/CUDA/2kjXI/src/compiler/execution.jl:375
[17] cufunction
@ ~/.julia/packages/CUDA/2kjXI/src/compiler/execution.jl:372 [inlined]
[18] macro expansion
@ ~/.julia/packages/CUDA/2kjXI/src/compiler/execution.jl:112 [inlined]
[19] #launch_heuristic#1200
@ ~/.julia/packages/CUDA/2kjXI/src/gpuarrays.jl:17 [inlined]
[20] launch_heuristic
@ ~/.julia/packages/CUDA/2kjXI/src/gpuarrays.jl:15 [inlined]
[21] _copyto!
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl:78 [inlined]
[22] copyto!
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl:44 [inlined]
[23] copy
@ ~/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl:29 [inlined]
[24] materialize
@ ./broadcast.jl:867 [inlined]
[25] adjoint
@ ~/.julia/packages/Zygote/nyzjS/src/lib/broadcast.jl:92 [inlined]
[26] _pullback
@ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
[27] generated_callfunc
@ ~/.julia/packages/NeuralPDE/nYBAW/src/discretize.jl:130 [inlined]
[28] _pullback(::Zygote.Context{…}, ::typeof(RuntimeGeneratedFunctions.generated_callfunc), ::RuntimeGeneratedFunctions.RuntimeGeneratedFunction{…}, ::CuArray{…}, ::ComponentVector{…}, ::NeuralPDE.Phi{…}, ::typeof(NeuralPDE.numeric_derivative), ::NeuralPDE.var"#239#246"{…}, ::NeuralPDE.var"#7#8", ::Vector{…})
@ Zygote ~/.julia/packages/Zygote/nyzjS/src/compiler/interface2.jl:0
[29] _apply(::Function, ::Vararg{Any})
@ Core ./boot.jl:946
[30] adjoint
@ ~/.julia/packages/Zygote/nyzjS/src/lib/lib.jl:203 [inlined]
[31] _pullback
@ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
[32] RuntimeGeneratedFunction
@ ~/.julia/packages/RuntimeGeneratedFunctions/M9ZX8/src/RuntimeGeneratedFunctions.jl:150 [inlined]
[33] _pullback(::Zygote.Context{…}, ::RuntimeGeneratedFunctions.RuntimeGeneratedFunction{…}, ::CuArray{…}, ::ComponentVector{…}, ::NeuralPDE.Phi{…}, ::typeof(NeuralPDE.numeric_derivative), ::NeuralPDE.var"#239#246"{…}, ::NeuralPDE.var"#7#8", ::Vector{…})
@ Zygote ~/.julia/packages/Zygote/nyzjS/src/compiler/interface2.jl:0
[34] #197
@ ~/.julia/packages/NeuralPDE/nYBAW/src/discretize.jl:150 [inlined]
[35] _pullback(::Zygote.Context{…}, ::NeuralPDE.var"#197#198"{…}, ::CuArray{…}, ::ComponentVector{…})
@ Zygote ~/.julia/packages/Zygote/nyzjS/src/compiler/interface2.jl:0
[36] #86
@ ~/.julia/packages/NeuralPDE/nYBAW/src/training_strategies.jl:122 [inlined]
[37] _pullback(ctx::Zygote.Context{…}, f::NeuralPDE.var"#86#87"{…}, args::ComponentVector{…})
@ Zygote ~/.julia/packages/Zygote/nyzjS/src/compiler/interface2.jl:0
[38] #264
@ ./none:0 [inlined]
[39] _pullback(ctx::Zygote.Context{…}, f::NeuralPDE.var"#264#285"{…}, args::NeuralPDE.var"#86#87"{…})
@ Zygote ~/.julia/packages/Zygote/nyzjS/src/compiler/interface2.jl:0
[40] (::Zygote.var"#666#670"{Zygote.Context{false}, NeuralPDE.var"#264#285"{ComponentVector{…}}})(args::Function)
@ Zygote ~/.julia/packages/Zygote/nyzjS/src/lib/array.jl:188
[41] iterate
@ ./generator.jl:48 [inlined]
[42] collect_to!(dest::Vector{Tuple{…}}, itr::Base.Generator{Vector{…}, Zygote.var"#666#670"{…}}, offs::Int64, st::Int64)
@ Base ./array.jl:838
[43] collect_to_with_first!(dest::Vector{…}, v1::Tuple{…}, itr::Base.Generator{…}, st::Int64)
@ Base ./array.jl:816
[44] _collect(c::Vector{…}, itr::Base.Generator{…}, ::Base.EltypeUnknown, isz::Base.HasShape{…})
@ Base ./array.jl:810
[45] collect_similar
@ ./array.jl:709 [inlined]
[46] map
@ ./abstractarray.jl:3371 [inlined]
[47] ∇map(cx::Zygote.Context{…}, f::NeuralPDE.var"#264#285"{…}, args::Vector{…})
@ Zygote ~/.julia/packages/Zygote/nyzjS/src/lib/array.jl:188
[48] _pullback
@ ~/.julia/packages/Zygote/nyzjS/src/lib/array.jl:231 [inlined]
[49] full_loss_function
@ ~/.julia/packages/NeuralPDE/nYBAW/src/discretize.jl:463 [inlined]
[50] _pullback(::Zygote.Context{…}, ::NeuralPDE.var"#full_loss_function#283"{…}, ::ComponentVector{…}, ::SciMLBase.NullParameters)
@ Zygote ~/.julia/packages/Zygote/nyzjS/src/compiler/interface2.jl:0
[51] pullback(::Function, ::Zygote.Context{false}, ::ComponentVector{Float64, CuArray{…}, Tuple{…}}, ::Vararg{Any})
@ Zygote ~/.julia/packages/Zygote/nyzjS/src/compiler/interface.jl:90
[52] pullback(::Function, ::ComponentVector{Float64, CuArray{…}, Tuple{…}}, ::SciMLBase.NullParameters)
@ Zygote ~/.julia/packages/Zygote/nyzjS/src/compiler/interface.jl:88
[53] withgradient(::Function, ::ComponentVector{Float64, CuArray{…}, Tuple{…}}, ::Vararg{Any})
@ Zygote ~/.julia/packages/Zygote/nyzjS/src/compiler/interface.jl:205
[54] value_and_gradient
@ ~/.julia/packages/DifferentiationInterface/bulUW/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl:92 [inlined]
[55] value_and_gradient!(f::Function, grad::ComponentVector{…}, prep::DifferentiationInterface.NoGradientPrep, backend::AutoZygote, x::ComponentVector{…}, contexts::DifferentiationInterface.Constant{…})
@ DifferentiationInterfaceZygoteExt ~/.julia/packages/DifferentiationInterface/bulUW/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl:105
[56] (::OptimizationZygoteExt.var"#fg!#16"{…})(res::ComponentVector{…}, θ::ComponentVector{…})
@ OptimizationZygoteExt ~/.julia/packages/OptimizationBase/gvXsf/ext/OptimizationZygoteExt.jl:53
[57] macro expansion
@ ~/.julia/packages/OptimizationOptimisers/JgTMl/src/OptimizationOptimisers.jl:101 [inlined]
[58] macro expansion
@ ~/.julia/packages/Optimization/cfp9i/src/utils.jl:32 [inlined]
[59] __solve(cache::OptimizationCache{…})
@ OptimizationOptimisers ~/.julia/packages/OptimizationOptimisers/JgTMl/src/OptimizationOptimisers.jl:83
[60] solve!(cache::OptimizationCache{…})
@ SciMLBase ~/.julia/packages/SciMLBase/gZRU3/src/solve.jl:186
[61] solve(::OptimizationProblem{…}, ::Adam; kwargs::@Kwargs{…})
@ SciMLBase ~/.julia/packages/SciMLBase/gZRU3/src/solve.jl:94
Some type information was truncated. Use `show(err)` to see complete types.
Environment (please complete the following information):
Output of using Pkg; Pkg.status()
Status `~/julia_projects/NeuralPDEtests/Project.toml`
[b0b7db55] ComponentArrays v0.15.19
⌃ [b2108857] Lux v1.2.3
[d0bbae9a] LuxCUDA v0.3.3
[961ee093] ModelingToolkit v9.51.0
[315f7962] NeuralPDE v5.17.0
[7f7a1694] Optimization v4.0.5
[42dfb2eb] OptimizationOptimisers v0.3.4
Info Packages marked with ⌃ have new versions available and may be upgradable.
Output of using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)
Julia Version 1.11.1
Commit 8f5b7ca12ad (2024-10-1610:53 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: Linux (x86_64-linux-gnu)
CPU:20×12th Gen Intel(R) Core(TM) i9-12900H
WORD_SIZE:64
LLVM: libLLVM-16.0.6 (ORCJIT, alderlake)
Threads:20 default, 0 interactive, 10 GC (on 20 virtual cores)
Environment:
JULIA_NUM_THREADS = auto
JULIA_EDITOR = code
Additional context
@ChrisRackauckas, we discussed this error on a call last Thursday. It came up when I was working on GPU support NeuralLyapunov.jl, which passes along any parameters from the ODESystem to the PDESystem.
The text was updated successfully, but these errors were encountered:
Describe the bug 🐞
When a
PDESystem
has parameters (that aren't being estimated),solve
fails if trying to solve on the GPU.The below MRE shows a system without parameters being solved on the GPU, an identical system with parameters being solved on the CPU, but when I try solving the system with parameters on the GPU it errors in
solve
.Expected behavior
Your
PDESystem
should be able to have parameters whether you solve on the CPU or the GPU.Minimal Reproducible Example 👇
Error & Stacktrace⚠️
Environment (please complete the following information):
using Pkg; Pkg.status()
using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)
versioninfo()
Additional context
@ChrisRackauckas, we discussed this error on a call last Thursday. It came up when I was working on GPU support NeuralLyapunov.jl, which passes along any parameters from the ODESystem to the PDESystem.
The text was updated successfully, but these errors were encountered: