Skip to content

next_probes_implicit_step.md

Latest commit

 

History

History
287 lines (227 loc) · 13.9 KB

next_probes_implicit_step.md

File metadata and controls

287 lines (227 loc) · 13.9 KB
Raw

Next probes: GPU seam tracer bug localised to implicit_step!

Plan file for the next debugging session. The bug we are chasing is documented in serial_vs_distributed_validation.md; this file picks up where that one left off.

RESOLVED 2026-05-16 — fixed upstream in Oceananigans 0.107.7

The bug isolated by Probes A/B/C (a CPU↔GPU divergence in solve_batched_tridiagonal_system_kernel! on the rank-1 fold-owning path) is fixed upstream and pulled in via the bp/offline_ACCESS-OM2_v3 fork bump from 0.107.6 to 0.107.7 (commit e5c4e36 here, Oceananigans tree-sha 8844e43).

End-to-end verification on the original failing case (OM2-1 1968-1977, centered2_AB2, gpuvolta):

run PBS job wallclock notes
serial 1×1 168447353 7:14 1-year, gpuvolta
dist 1×2 168447357 8:42 1-year, gpuvolta
compare 168447438 4:45 DURATION_TAG=1year

Visual end-of-year age plot from scripts/plotting/compare_runs_across_architectures.sh confirms the seam-row divergence is gone. Stop condition reached.

The trail below is preserved for reference; sections are unchanged.

TL;DR — what we know as of 2026-05-15

The bug is in implicit_step! (implicit vertical-diffusion tridiagonal solve), on rank 1 only.

Evidence (from test/probe_tracer_tendency.jl manual decomposition of time_step!, CPU 1×2 vs GPU 1×2, OM2-1 1968-1977, centered2_AB2 default config):

stage in the iter 0 → iter 1 step rank 0 age max|diff| rank 1 age max|diff| rank 1 seam row (parent y=14)
iter 0 inputs (age, u, v, w, η, σ, κV, metrics, bottom, …) 0 (except w at 1 Float64 ULP) 0 (except w at 1 Float64 ULP) 0
compute_tracer_tendencies! → Gⁿ.age 0 0 0 ← bit-identical
_ab2_step_tracer_field! → age (post_explicit) 1.8e-12 s (Float64 ULP, 120 cells) 1.8e-12 s (Float64 ULP, 130 cells) 0 ← bit-identical
implicit_step! → age (post_implicit) 3.3e-11 s (Float64 ULP, 388K cells — clean) 4534.677 s (250K cells) 4534.677 s in 3769 cells ← BUG
update_state! → iter 1 age 4534.677 s (494K cells, ULP-spread elsewhere) 4534.677 s (340K cells) 4534.677 s in 4266 cells

So:

  1. Everything upstream of implicit_step! is CPU↔GPU bit-identical to within Float64 ULP — including Gⁿ.age, the freshly computed tendency.
  2. implicit_step! itself produces the 4535 s seam-row divergence on rank 1.
  3. Rank 0 is clean (only Float64-ULP scatter from a slightly different tridiagonal-solver rounding path on GPU — harmless).
  4. The asymmetry between rank 0 and rank 1 is the key. Rank 1 has y-topology LeftConnectedRightFaceFolded (south rank-rank seam + north tripolar fold). Rank 0 has RightConnected (south wall + north rank-rank seam).

Probe artefacts on disk

Per-rank JLD2 in outputs/ACCESS-OM2-1/1deg_jra55_iaf_omip2_cycle6/1968-1977/standardrun/cgridtransports_wdiagnosed_centered2_AB2/1x2/probe/:

  • probe_tendency_{cpu|gpu}_iter{0,1}_rank{0,1}.jld2 — full state, ~1 GB each, includes age, u_fts/v_fts/w, eta_fts, Gn_age/Gm_age, sigma_cc/eta_n/dt_sigma, 9 grid-metric variants, bottom_height, κV, closure_fields (currently (nothing, nothing) for our setup).
  • probe_age_{cpu|gpu}_iter0_{post_explicit|post_implicit}_rank{0,1}.jld2 — lightweight age + Gⁿ_age snapshots straddling implicit_step!.

Total ~9 GB. Safe to delete after each round — they're regenerable in ~6 min wallclock via JOB_CHAIN=probetend-probetendcpu.

Where the suspect code lives

The implicit-vertical-diffusion solver path used by our model:

Two reasons the GPU+rank-1-only divergence might arise:

  1. Coefficient construction reads diffusivity κᶜᶜᶠ(i, j, k, grid, ..., closure, K, id) and possibly grid metrics at face locations. If any of these inputs has a fold-row / partition-aware specialisation that misfires on a rank whose y-topology is LeftConnectedRightFaceFolded, the coefficients themselves end up different from the serial-CPU baseline.
  2. The solver kernel launch parametrises over the active-cells map on GPU. If the active-cells map for rank 1 includes / excludes a different set of columns than CPU does, some columns get solved-with-junk on GPU.

Probes for the next session

In rough priority order. All probes should keep the existing PARTITION=1x2 CPU + GPU baseline from probe_tracer_tendency.jl as the control.

Probe A — κV halo at the seam (cheap; 0 new code) — DONE 2026-05-15: clean

The existing dumps already include κV. Read the rank-1 GPU vs CPU κV parent-array at parent y = 13 (south halo), 14 (first interior), and the matching north-halo rows; print min/max/n_differ. If CPU and GPU disagree on κV at the seam-adjacent halo, the bug is upstream of the solver, in the κV halo-fill path.

Implementation: scripts/debugging/probe_kV_at_seam.jl (commit e69499f). Run on the login node — pure JLD2 reading, no GPU needed:

julia --project scripts/debugging/probe_kV_at_seam.jl

Result: κV_age is CPU↔GPU bit-identical on both ranks. max|diff| = 0 on every row of interest (parent y = 13, 14, 15, 163, 164, 165) and globally (4.3M cells). Rules out "wrong κV halo on rank 1" as the upstream cause.

Probe B — disable implicit vertical diffusion (direct confirmation) — DONE 2026-05-15: bug vanishes → implicit_step! is the sole cause

Toggle κVField (the closure's κ) to zero, or temporarily replace implicit_vertical_diffusion = VerticalScalarDiffusivity(...) with nothing for the age tracer. Rerun the 10-step diag pair (GPU 1×2 vs CPU 1×2). If the seam tracer drift vanishes, that confirms implicit_step! is the sole cause; if it persists, there is a second mechanism we haven't found.

Implementation: IMPLICIT_KAPPAV=yes|no env var (default yes; commit 4f57feb). When set to no, src/setup_model.jl drops VerticalScalarDiffusivity from the closure tuple; Oceananigans then returns implicit_solver = nothing, and implicit_step!(field, ::Nothing, …) = nothing makes the call site a clean no-op without touching the probe script's call site. Outputs are tagged with a _noKV suffix on MODEL_CONFIG.

Probe re-fix (commit ff77a7a): dump_kappa_v! in test/probe_tracer_tendency.jl previously hard-coded closure[2]. Now searches the tuple for VerticalScalarDiffusivity by type and returns early when none is present.

Submit:

PARENT_MODEL=ACCESS-OM2-1 GPU_QUEUE=gpuvolta PARTITION=1x2 \
    IMPLICIT_KAPPAV=no PROBE_NSTEPS=1 \
    JOB_CHAIN=probetend-probetendcpu bash scripts/test_driver.sh
PARENT_MODEL=ACCESS-OM2-1 PARTITION=1x2 IMPLICIT_KAPPAV=no \
    JOB_CHAIN=compareprobe bash scripts/test_driver.sh

Result: with IMPLICIT_KAPPAV=no, the seam-row divergence vanishes:

stage DEFAULT rank-1 max|diff| (seam) noKV rank-1 max|diff| (seam)
post_explicit 1.8e-12 s ULP (0) 1.8e-12 s ULP (0)
post_implicit 4534.677 s (3769) 1.8e-12 s ULP (0)
iter-1 age 4534.677 s (4266) 1.8e-12 s ULP (0)

post_implicit becomes equal to post_explicit exactly. Confirms implicit_step! (the BatchedTridiagonalSolver path) is the sole cause of the GPU rank-1 seam tracer bug. No second mechanism exists.

Probe C — solver coefficients dump — DONE 2026-05-15: coefficients clean → bug is in solve_batched_tridiagonal_system_kernel!

Extend the probe to dump the per-column tri-diagonal coefficients that implicit_step! constructs before the Thomas sweep, for every (i, j) in rank 1's parent-y = 14 row. This requires either:

  • calling implicit_step!'s coefficient-construction kernels directly with the same arguments and dumping their outputs; or
  • monkey-patching implicit_step! to also write the LHS / RHS arrays to JLD2 before calling the BatchedTridiagonalSolver.

Implementation: dump_implicit_coefficients(model, n, Δt) in test/probe_tracer_tendency.jl (commit 160620a) launches a kernel that mirrors BatchedTridiagonalSolver's per-cell get_coefficient calls — without the Thomas sweep — and writes a/b/c into Nx×Ny×Nz Float64 arrays. The RHS is already captured in the post_explicit age snapshot. Filters closure exactly the way implicit_step! does (only VerticallyImplicitTimeDiscretization components survive). Companion diff script: scripts/debugging/compare_implicit_coeffs.jl.

Submit:

PARENT_MODEL=ACCESS-OM2-1 GPU_QUEUE=gpuvolta PARTITION=1x2 \
    PROBE_NSTEPS=1 \
    JOB_CHAIN=probetend-probetendcpu bash scripts/test_driver.sh
PARENT_MODEL=ACCESS-OM2-1 PARTITION=1x2 PROBE_NSTEPS=1 \
    julia --project scripts/debugging/compare_implicit_coeffs.jl

Result: the LHS coefficients are CPU↔GPU bit-identical at the rank-1 seam row (interior j=1 = parent y=14):

coef seam row max|diff| seam row n_differ GLOBAL max|diff| (n_differ / total)
a 0 0 / 18000 9.9e-14 (135 / 2.7M, Float64 ULP)
b 0 0 / 18000 1.1e-13 (95 / 2.7M, Float64 ULP)
c 0 0 / 18000 7.1e-14 (123 / 2.7M, Float64 ULP)

So in the rank-1 seam row:

  • LHS (a, b, c): bit-identical CPU↔GPU (Probe C, above)
  • RHS (post_explicit age): bit-identical CPU↔GPU (Probe B re-confirmed)
  • post_implicit age: differs by 4534.677 s in 3769 / 18000 cells

Identical inputs → different outputs. The bug is in solve_batched_tridiagonal_system_kernel! (the Thomas-sweep kernel in Oceananigans src/Solvers/batched_tridiagonal_solver.jl), specifically on the GPU rank-1 path. The coefficient builder is exonerated. The remaining hypotheses (from this doc's TL;DR section) collapse to one: the GPU kernel launches / executes wrong on the rank-1 LeftConnectedRightFaceFolded y-topology, in a way that the CPU path doesn't reproduce. Probes D + E (and a structural look at the active-cells map for rank 1) are the next steps.

Halo-size note (housekeeping, 2026-05-15)

The serial grid was rebuilt mid-session with new defaults Hx=Hy=7, Hz=2 (down from Hx=Hy=13, Hz=7). Partitioned monthly FTS files had to be rebuilt to match (JOB_CHAIN=partition, ~3 min). Coefficient dumps in this section are therefore at the new halo; parent-array shapes in the TL;DR table (386 × 176 × 64) refer to the old halo and won't match a fresh probe run. Interior shapes (360 × 150 × 64 for age, 360 × 150 × 50 for the coefficient arrays) are halo-invariant.

Probe D — try a 2×1 partition (x-direction split)

If we partition along x instead of y, the rank-rank seam is purely in i and the y-topology on each rank stays Bounded (no LeftConnectedRightFaceFolded issue). If the seam bug also appears on 2×1, the trigger is "any MPI partition", not the fold-row topology. If 2×1 is clean, the trigger is specifically the fold-row topology that rank 1 of 1×2 has.

Submit:

GPU_QUEUE=gpuvolta PARENT_MODEL=ACCESS-OM2-1 PARTITION=2x1 \
    JOB_CHAIN=probetend-probetendcpu bash scripts/test_driver.sh
PARENT_MODEL=ACCESS-OM2-1 PARTITION=2x1 \
    JOB_CHAIN=compareprobe bash scripts/test_driver.sh

This relies on partitioned FTS data existing for 2×1. Check preprocessed_inputs/.../1968-1977/partitions/2x1/ first; if missing, run the partition step.

Probe E — substitute BatchedTridiagonalSolver with a CPU-only stub

For the GPU run, route the age tracer's implicit_step! through a CPU solve (copy field to host, solve, copy back). If this fixes the seam diff, the bug is in the GPU implementation of the solver itself (probably in the kernel that reads/writes column data). If the diff persists, the bug is upstream of the solver.

Lowest-cost expression of this idea is IMPLICIT_KAPPAV=no (Probe B), which removes the call site entirely. Probe E is the "remove the suspect, keep the rest of the physics" experiment; B is "remove the suspect and its contribution to physics."

Stop conditions

The probe loop is done when one of:

  • A specific Oceananigans upstream PR or local change makes the seam-row diff in step 0→1, post_implicit, rank1 drop to Float64-ULP level (currently 4534.677 s, expected < 1e-10 s).
  • We have a reproducible MWE we can hand off as an Oceananigans upstream issue.

Hygiene

Probe artefacts are ~9 GB per round. The existing probe always writes to the same path so reruns clobber old files. Do not push probe JLD2s into the repo; the existing .gitignore already excludes outputs/.