GPU Benchmarks

.gitignore

@@ -32,5 +32,3 @@ coverage/lcov.info
 !MPI_BGC_wide_E_neg_cmyk.png
 !logo.png
 !logo_sq.png
-test/setup_local_testrun.jl
-docs/setup_local_docsrun.jl

docs/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
@@ -10,9 +11,13 @@ EasyHybrid = "61bb816a-e6af-4913-ab9e-91bff2e122e3"
 Hyperopt = "93e5fe13-2215-51db-baaf-2e9a34fb2712"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
+oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
 
 [sources]
 EasyHybrid = {path = ".."}

docs/literate/tutorials/synthetic_respiration_gpu.jl

@@ -6,10 +6,19 @@
 # on synthetic data for respiration modeling with Q10 temperature sensitivity.
 #
 
+## for development, use local setup
+# include("../../setup_local_docsrun.jl")
+
 using EasyHybrid
-using Metal
-## using CUDA
-using Lux
+using MLDataDevices
+
+# Load whichever GPU backend is available on this machine (CUDA / AMDGPU /
+# Metal / oneAPI), or fall back to CPU. Defines `GPU_BACKEND_PKG` and
+# `GPU_DEVICE_TYPE` in the caller's scope.
+include("../../setup_gpu_backend.jl")
+
+cpu_device() isa CPUDevice
+gpu_device() isa GPU_DEVICE_TYPE
 
 # ## Data Loading and Preprocessing
 #
@@ -18,9 +27,9 @@ using Lux
 df = load_timeseries_netcdf("https://github.com/bask0/q10hybrid/raw/master/data/Synthetic4BookChap.nc");
 
 # Select a subset of data for faster execution
-df = df[1:5000, :];
-nothing #hide
-first(df, 5)
+#df = df[1:20000, :];
+#nothing #hide
+#first(df, 5)
 
 # ## Define the Physical Model
 #
@@ -84,7 +93,7 @@ large_nn_hybrid_model = constructHybridModel(
     parameters,              # Parameter definitions
     neural_param_names,      # NN-predicted parameters
     global_param_names,      # Global parameters
-    hidden_layers = [1024, 512, 256, 128, 64], # Neural network architecture
+    hidden_layers = [512, 512, 512], # Neural network architecture
     activation = sigmoid,      # Activation function
     scale_nn_outputs = true, # Scale neural network outputs
     input_batchnorm = true   # Apply batch normalization to inputs
@@ -94,53 +103,191 @@ large_nn_hybrid_model = constructHybridModel(
 # Train the hybrid model
 
 cfg = EasyHybrid.TrainConfig(
-    nepochs = 20,
-    batchsize = 64,
-    opt = RMSProp(0.01),
+    nepochs = 10,
+    batchsize = 512,
+    opt = Adam(0.001),
     loss_types = [:mse, :nse],
     show_progress = false,
     keep_history = false, # set to true to keep per-epoch history, losses, predictions, etc.
     save_training = false, # Set to true to enable saving training history and checkpoints
+    plotting = false,
 )
 
-using BenchmarkTools
+using Suppressor
+using Printf
 
-gpu_small_nn() = tune(
-    small_nn_hybrid_model, df, cfg;
-    gdev = gpu_device(), model_name = "small_nn_gpu"
-)
+function _pure(s)
+    return s.time - s.gctime - s.compile_time - s.recompile_time
+end
 
-cpu_small_nn() = tune(
-    small_nn_hybrid_model, df, cfg;
-    gdev = cpu_device(), model_name = "small_nn_cpu"
-)
+function _warn_overhead(tag, label, s, warn_frac)
+    gc_frac = s.gctime / s.time
+    compile_frac = s.compile_time / s.time
+    recompile_frac = s.recompile_time / s.time
+    overhead_frac = gc_frac + compile_frac + recompile_frac
+    return if overhead_frac > warn_frac
+        pct(x) = round(100 * x; digits = 1)
+        @warn """[$label/$tag] $(pct(overhead_frac))% of elapsed time was overhead:
+        gc=$(pct(gc_frac))% ($(round(s.gctime; digits = 3))s), \
+        compile=$(pct(compile_frac))% ($(round(s.compile_time; digits = 3))s), \
+        recompile=$(pct(recompile_frac))% ($(round(s.recompile_time; digits = 3))s).
+        Elapsed ratio will be misleading; prefer the pure ratio."""
+    end
+end
+
+# Benchmark `tune` on CPU vs GPU. Uses `@timed` so we can separate out
+# compile/recompile/GC overhead from steady-state work, and reports BOTH:
+#   - elapsed ratio (wall clock CPU / wall clock GPU)
+#   - pure ratio    (elapsed minus GC/compile/recompile, CPU/GPU)
+#
+# `warn_frac` is the fraction of total time spent in GC / (re)compile above
+# which a warning is emitted for that run; defaults to 5%.
+function bench_cpu_vs_gpu(model, df, cfg; label::AbstractString, warmup::Bool = true, warn_frac::Real = 0.05, tune_kwargs...)
+    # A short 1-epoch run on each device triggers Julia/Zygote specialization
+    # and GPU kernel compilation, so the timed runs below measure steady-state
+    # training rather than first-call compile latency.
+    if warmup
+        warm_cfg = EasyHybrid.TrainConfig(;
+            nepochs = 1,
+            batchsize = cfg.batchsize,
+            opt = cfg.opt,
+            show_progress = false,
+            save_training = false,
+            keep_history = false,
+            plotting = false,
+        )
+        @suppress tune(model, df, warm_cfg; gdev = cpu_device(), tune_kwargs...)
+        @suppress tune(model, df, warm_cfg; gdev = gpu_device(), tune_kwargs...)
+    end
 
-gpu_large_nn() = tune(
-    large_nn_hybrid_model, df, cfg;
-    gdev = gpu_device(), model_name = "large_nn_gpu"
+    cpu_stats = @suppress @timed tune(model, df, cfg; gdev = cpu_device(), tune_kwargs...)
+    gpu_stats = @suppress @timed tune(model, df, cfg; gdev = gpu_device(), tune_kwargs...)
+
+    _warn_overhead("CPU", label, cpu_stats, warn_frac)
+    _warn_overhead("GPU", label, gpu_stats, warn_frac)
+
+    cpu_pure = _pure(cpu_stats)
+    gpu_pure = _pure(gpu_stats)
+    elapsed_ratio = cpu_stats.time / gpu_stats.time
+    pure_ratio = cpu_pure / gpu_pure
+
+    @printf(
+        "[%s] CPU: total=%.3fs compile=%.3fs recompile=%.3fs gc=%.3fs pure=%.3fs\n",
+        label, cpu_stats.time, cpu_stats.compile_time, cpu_stats.recompile_time, cpu_stats.gctime, cpu_pure
+    )
+    @printf(
+        "[%s] GPU: total=%.3fs compile=%.3fs recompile=%.3fs gc=%.3fs pure=%.3fs\n",
+        label, gpu_stats.time, gpu_stats.compile_time, gpu_stats.recompile_time, gpu_stats.gctime, gpu_pure
+    )
+    @printf(
+        "[%s] with GPU we get: elapsed=%.2fx  pure=%.2fx\n",
+        label, elapsed_ratio, pure_ratio
+    )
+
+    # Return a slim summary. We deliberately drop `.value` (the full TrainResults)
+    # and `.gcstats` so that REPL/Literate printing of the return value stays compact.
+    slim(s) = (; s.time, s.bytes, s.gctime, s.compile_time, s.recompile_time)
+    return (; cpu = slim(cpu_stats), gpu = slim(gpu_stats), elapsed_ratio, pure_ratio)
+end
+
+bench_cpu_vs_gpu(small_nn_hybrid_model, df, cfg; label = "small NN"); # on our gpu1-hpc22 0.57x
+bench_cpu_vs_gpu(large_nn_hybrid_model, df, cfg; label = "large NN"); # on our gpu1-hpc22 2.73x
+
+# ## Sweep: GPU speedup vs hidden-layer width
+#
+# Build a hybrid model for a range of hidden-layer widths (fixed depth of 3),
+# benchmark each on CPU and GPU, and plot the elapsed and pure ratios.
+# A ratio > 1 means GPU is faster than CPU; a ratio < 1 means GPU is slower.
+
+widths = [16, 64, 256, 512, 1024]
+depths = 2:5
+
+results = [
+    begin
+            r = bench_cpu_vs_gpu(
+                small_nn_hybrid_model, df, cfg;
+                label = "d=$d w=$w", hidden_layers = fill(w, d)
+            )
+            (;
+                depth = d, width = w,
+                elapsed_ratio = r.elapsed_ratio, pure_ratio = r.pure_ratio,
+                cpu_time = r.cpu.time, gpu_time = r.gpu.time,
+            )
+        end
+        for d in depths for w in widths
+]
+
+using CairoMakie
+
+fig = Figure(size = (820, 520))
+ax = Axis(
+    fig[1, 1];
+    xlabel = "Hidden-layer width",
+    ylabel = "GPU speedup (CPU / GPU)",
+    title = "GPU speedup vs hidden-layer width, per depth",
+    xscale = log2,
+    xticks = (widths, string.(widths)),
 )
 
-cpu_large_nn() = tune(
-    large_nn_hybrid_model, df, cfg;
-    gdev = cpu_device(), model_name = "large_nn_cpu"
+hlines!(ax, [1.0]; color = :gray, linestyle = :dash) # break-even
+
+# One colour per depth; `pure` is solid, `elapsed` is dashed.
+palette = Makie.wong_colors()
+for (i, d) in enumerate(depths)
+    rs = filter(r -> r.depth == d, results)
+    ws = [r.width for r in rs]
+    el = [r.elapsed_ratio for r in rs]
+    pr = [r.pure_ratio    for r in rs]
+    c = palette[mod1(i, length(palette))]
+    # Only the solid (pure) line gets a label → legend entries are per-depth.
+    scatterlines!(ax, ws, pr; color = c, linestyle = :solid, marker = :circle, label = "depth=$d")
+    scatterlines!(ax, ws, el; color = c, linestyle = :dash, marker = :xcross)
+end
+
+# Legend 1: depth → colour (auto-picked up from the labeled solid lines).
+axislegend(ax, "depth"; position = :lt)
+
+# Legend 2: line style → metric (built manually).
+style_elements = [
+    LineElement(color = :black, linestyle = :solid),
+    LineElement(color = :black, linestyle = :dash),
+]
+Legend(fig[1, 2], style_elements, ["pure", "elapsed"], "metric")
+
+fig
+
+
+batchsizes = [64, 128, 256, 512, 1024, 2048, 4096, 8192]
+bs_hidden = [64, 64, 64]
+
+bs_results = map(batchsizes) do b
+    r = bench_cpu_vs_gpu(
+        small_nn_hybrid_model, df, cfg;
+        label = "b=$b", batchsize = b, hidden_layers = bs_hidden
+    )
+    (;
+        batchsize = b, elapsed_ratio = r.elapsed_ratio, pure_ratio = r.pure_ratio,
+        cpu_time = r.cpu.time, gpu_time = r.gpu.time,
+    )
+end
+
+fig_bs = Figure(size = (720, 480))
+ax_bs = Axis(
+    fig_bs[1, 1];
+    xlabel = "Batch size (hidden = $bs_hidden)",
+    ylabel = "GPU speedup (CPU / GPU)",
+    title = "GPU speedup vs batch size",
+    xscale = log2,
+    xticks = (batchsizes, string.(batchsizes)),
 )
 
-# warm-up to pay compilation once
-gpu_small_nn();
-cpu_small_nn();
-gpu_large_nn();
-cpu_large_nn();
-nothing # hide
-
-# ## With Large NN CPU is slower than GPU
-# Large NN on GPU
-@benchmark gpu_large_nn() samples = 4 evals = 1
-# Large NN on CPU
-@benchmark cpu_large_nn() samples = 4 evals = 1
-
-# ## With Small NN CPU and GPU are on par
-using BenchmarkTools
-# Small NN on GPU
-@benchmark gpu_small_nn() samples = 4 evals = 1
-# Small NN on CPU
-@benchmark cpu_small_nn() samples = 4 evals = 1
+bs = [r.batchsize     for r in bs_results]
+bs_elapsed = [r.elapsed_ratio for r in bs_results]
+bs_pure = [r.pure_ratio    for r in bs_results]
+
+hlines!(ax_bs, [1.0]; color = :gray, linestyle = :dash) # break-even
+scatterlines!(ax_bs, bs, bs_pure; color = :steelblue, linestyle = :solid, marker = :circle, label = "pure")
+scatterlines!(ax_bs, bs, bs_elapsed; color = :steelblue, linestyle = :dash, marker = :xcross, label = "elapsed")
+
+axislegend(ax_bs; position = :lt)
+fig_bs

docs/setup_gpu_backend.jl

@@ -0,0 +1,57 @@
+# Load the appropriate GPU backend for `MLDataDevices`, independent of machine.
+#
+# Tries the common GPU trigger packages in order:
+#   - LuxCUDA  -> CUDADevice   (NVIDIA)
+#   - AMDGPU   -> AMDGPUDevice (AMD ROCm)
+#   - Metal    -> MetalDevice  (Apple Silicon)
+#   - oneAPI   -> oneAPIDevice (Intel)
+#
+# The first backend that is both installed AND functional on this machine
+# is kept. If none are available, we silently fall back to CPU and
+# `gpu_device()` will return a `CPUDevice`.
+#
+# After `include`ing this file, the following are defined in the caller's
+# namespace:
+#   - `GPU_BACKEND_PKG`  : `Symbol` of the loaded package (or `nothing`)
+#   - `GPU_DEVICE_TYPE`  : the `MLDataDevices` device type (e.g. `CUDADevice`,
+#                          or `CPUDevice` as fallback)
+
+using MLDataDevices
+
+const _GPU_TRIGGERS = [
+    (:LuxCUDA, :CUDADevice),
+    (:AMDGPU, :AMDGPUDevice),
+    (:Metal, :MetalDevice),
+    (:oneAPI, :oneAPIDevice),
+]
+
+function _load_gpu_backend()
+    for (pkg, dev) in _GPU_TRIGGERS
+        try
+            @eval Main using $pkg
+        catch err
+            @debug "Skipping $pkg (not installed / failed to load)" exception = err
+            continue
+        end
+
+        # `using $pkg` may load an `MLDataDevices` package extension that
+        # adds new methods (e.g. a real `functional(::Type{CUDADevice})`).
+        # Those methods live in a newer world age than this function, so
+        # we must dispatch via `invokelatest` or we'll get the stale
+        # (pre-extension) method that reports `false`.
+        DeviceT = Base.invokelatest(getfield, MLDataDevices, dev)
+        is_functional = Base.invokelatest(MLDataDevices.functional, DeviceT)
+
+        if is_functional
+            @info "GPU backend loaded: $pkg → $dev"
+            return pkg, DeviceT
+        else
+            @warn "$pkg loaded but $dev is not functional on this machine; trying next backend."
+        end
+    end
+
+    @info "No functional GPU backend found; falling back to CPU."
+    return nothing, CPUDevice
+end
+
+const GPU_BACKEND_PKG, GPU_DEVICE_TYPE = _load_gpu_backend()

docs/setup_local_docsrun.jl

@@ -0,0 +1,4 @@
+using Pkg
+Pkg.activate("docs")
+Pkg.develop(path = pwd())
+Pkg.instantiate()