Nonlinear equations work on CPU but not GPU #914
Labels
Comments
nicholaskl97
changed the title
|
I've done a bit more digging and it's more complicated than just "nonlinear" vs. "linear". In particular terms like |
Bug: GPU Computation Failures with Certain Nonlinear Equations in NeuralPDEIssue Description Some nonlinear equations and boundary conditions fail to compute correctly on GPU, while working as expected on CPU. This significantly limits GPU support for more complex physics-informed neural network (PINN) problems. Specific Observations
Detailed Breakdown
Minimal Reproducible Exampleusing NeuralPDE
using Lux, LuxCUDA, Optimization, OptimizationOptimisers, Random, ComponentArrays
Random.seed!(200)
@parameters x
@variables u(..)
# Problematic equation
eq = [u(x)^2 ~ 0.0]
bc = [u(0.0)^2 ~ 0.0]
domain = [x ∈ (-1.0, 1.0)]
# Neural network setup
chain = Chain(
Dense(1, 10, tanh),
Dense(10, 10, tanh),
Dense(10, 1)
)
ps = first(Lux.setup(Random.default_rng(), chain))
const gpud = gpu_device()
ps = ps |> ComponentArray |> gpud |> f64
strategy = QuasiRandomTraining(1000)
# GPU computation fails
discretization = PhysicsInformedNN(chain, strategy; init_params = ps)
@named pde_system = PDESystem(eq, bc, domain, [x], [u(x)])
prob = discretize(pde_system, discretization)
res = Optimization.solve(prob, OptimizationOptimisers.Adam(); maxiters = 10) |
5 tasks
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Describe the bug 🐞
Some nonlinear equations (or boundary conditions) lead to NaNs after optimization on GPU. As shown in the MRE below, the equation
u(x) ~ 0works on GPU (and CPU, not shown), but the equationu(x)^2 ~ 0only works on CPU, not GPU. The same is true for the boundary conditionsu(0) ~ 0andu(0)^2 ~ 0.Expected behavior
If a PDESystem works on CPU, it should also work on GPU. Equations that are nonlinear in the dependent variable shouldn't result in NaN. Users should be able to specify the equivalent boundary conditions
u(0) ~ 0andu(0)^2 ~ 0and get similar results.Minimal Reproducible Example 👇
Environment (please complete the following information):
using Pkg; Pkg.status()using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)versioninfo()Additional context
@ChrisRackauckas, we spoke about this on a call on Thursday. It prevents GPU support for NeuralLyapunov in all but the simplest cases.
The text was updated successfully, but these errors were encountered: