run-multinode-gpu

name

run-multinode-gpu

description

Launch Aether.jl across many GPUs and nodes under MPI on NERSC Perlmutter (Cray Shasta, 4x A100/node). Covers the salloc recipe, the srun launch environment, node-local GPU binding, GPU-aware MPI, and — critically — the Cray-MPICH "malformed environ" workaround without which multi-node CUDA kernel launches deadlock. Use for any distributed GPU run, weak/strong scaling study, or production multi-node job on Perlmutter.

Running Aether.jl on many GPUs / nodes (NERSC Perlmutter)

Perlmutter GPU nodes have 4x A100-40GB each. srun does not source your login profile, so several environment details that "just work" interactively must be set explicitly — and multi-node runs need a Cray-MPICH workaround or they deadlock at the first CUDA kernel launch. This skill is the recipe of record. Keep throwaway driver scripts outside the repo (AGENTS.md forbids scratch files in the tree); $SCRATCH/aether_bench/ is one option.

TL;DR — the five things that bite you

Symptom	Cause	Fix
execve(): julia: No such file	srun has no login PATH (juliaup shim absent)	call julia by absolute path $HOME/.juliaup/bin/julia
CUDA_ERROR_NOT_INITIALIZED	step got no GPUs	add -G <total_gpus> to srun
Package … required but does not seem to be installed / extension precompile fails	login profile prepends $HOME/julialib/ to JULIA_LOAD_PATH, which shadows dep resolution for MPIExt/CUDAExt under srun	set JULIA_LOAD_PATH="@:@stdlib"
Multi-node hang at first CUDA kernel (alloc + memcpy work, kernel launch never returns; high system-time spinning)	Cray MPICH inserts a malformed (no-=) env entry after MPI_Init; CUDA.jl deadlocks reading it	call sanitize_environ!() right after MPI.Init() (code below)
GPU-aware hang cuIpcGetMemHandle: invalid argument	Cray MPICH GPU transport uses CUDA IPC, incompatible with CUDA.jl's async memory pool	JULIA_CUDA_MEMORY_POOL=none

1. Allocate (from a LOGIN node only)

salloc must run on a login node, never from inside a compute-node shell. Interactive GPU allocation on Perlmutter:

# --no-shell returns immediately and keeps the allocation alive; drive it with
# `srun --jobid=<ID>` and it persists for the wall-time. Interactive allocations
# are hard to re-acquire — do NOT scancel one you still need.
salloc --nodes 4 --qos interactive --time 04:00:00 -C gpu -A <account> --no-shell
# -> "salloc: Granted job allocation <job_id>"   (capture this ID)

One 4-node allocation covers 1/2/4 (single node) and 8/16-GPU (2/4 node) runs.

2. The launch environment (every srun)

JULIA=$HOME/.juliaup/bin/julia              # absolute path: srun has no login PATH
AETHER=$HOME/julialib/Aether.jl             # main project env: has MPI + CUDA as hard deps
JID="<job_id>"                              # replace with the ID from salloc

JULIA_LOAD_PATH="@:@stdlib" \               # clean load path (avoid profile pollution)
JULIA_CUDA_MEMORY_POOL=none \               # disable CUDA async malloc (Cray-MPICH incompat)
MPICH_GPU_SUPPORT_ENABLED=1 \               # GPU-aware MPI (device buffers to MPI); 0 = host-staged
  srun --jobid=$JID -N <nodes> --ntasks-per-node=<gpus_per_node> \
       -n <total_ranks> -G <total_gpus> --gpu-bind=none \
       $JULIA --project=$AETHER --startup-file=no driver.jl <args>

• --ntasks-per-node = GPUs used per node (≤4); -n = total ranks = total GPUs (one rank per GPU); -N = nodes. -G <total_gpus> is mandatory or the step gets no GPUs.
• --gpu-bind=none exposes all 4 node GPUs to every rank; the script picks its own by node-local rank (below). This is fine and what the recipe assumes.
• Precompile once before the first multi-rank launch so ranks don't race: JULIA_LOAD_PATH="@:@stdlib" srun --jobid=$JID -N1 -n1 -G1 --gpu-bind=none $JULIA --project=$AETHER -e 'using Pkg; Pkg.precompile()'
• First multi-node load can be slow from cold FS caches. Warm all nodes first (1 rank/node, ~30 s): srun --jobid=$JID -N<nodes> --ntasks-per-node=1 -G<nodes> --gpu-bind=none $JULIA --project=$AETHER -e 'import CUDA; CUDA.functional()'

3. Script skeleton (the workaround is load-bearing)

using MPI
MPI.Initialized() || MPI.Init()

# ── Cray-MPICH workaround ────────────────────────────────────────────────────
# Multi-node srun inserts a malformed (no-'=') entry into `environ` after
# MPI_Init; CUDA.jl deadlocks on it at the first kernel launch (alloc & memcpy
# still work, which makes it baffling). Strip such entries immediately.
function sanitize_environ!()
    envp = unsafe_load(cglobal(:environ, Ptr{Ptr{Cchar}}))
    valid = Tuple{String,String}[]; i = 1
    while (p = unsafe_load(envp, i)) != C_NULL
        s = unsafe_string(p); eq = findfirst('=', s)
        eq === nothing || push!(valid, (String(s[1:eq-1]), String(s[eq+1:end])))
        i += 1
    end
    ccall(:clearenv, Cint, ())
    for (k, v) in valid
        ccall(:setenv, Cint, (Cstring, Cstring, Cint), k, v, 1)
    end
    return nothing
end
sanitize_environ!()
# ─────────────────────────────────────────────────────────────────────────────

import CUDA
using Aether

comm = MPI.COMM_WORLD
rank = MPI.Comm_rank(comm); nranks = MPI.Comm_size(comm)

# Bind each rank to its node-local GPU (correct across nodes).
lcomm = MPI.Comm_split_type(comm, MPI.COMM_TYPE_SHARED, rank)
lrank = MPI.Comm_rank(lcomm)
CUDA.device!(lrank % length(collect(CUDA.devices())))
# (equivalently: parse(Int, ENV["SLURM_LOCALID"]))

FT = Float64
architecture = GPU(CUDA.CUDABackend())
mesh = Mesh(architecture, FT; size = (N1, N2, N3), extent = ext,
            cells_per_block = (B, B, B), nghost = 3,
            communicator = nranks == 1 ? nothing : comm)   # nothing when serial
simulation = Simulation(mesh; eos = IdealMHD(FT(5//3), FT(1e-12), FT(1e-10)),
                        reconstruction = PPM(), riemann_solver = HLLD(),
                        stepper = RK2(FT), cfl = FT(0.3),
                        gpu_aware_mpi = true)               # false = host-staged
set_initial_condition!(simulation; vector_potential = …) do x1, x2, x3
    …                                                       # see run-problem skill
end

# Timing: barrier, time, barrier; reduce the max wall across ranks.
run!(simulation; stop_cycle = 3); CUDA.synchronize(); MPI.Barrier(comm)  # warmup
c0 = simulation.clock.cycle
MPI.Barrier(comm); t0 = time_ns()
run!(simulation; stop_cycle = c0 + 500)
CUDA.synchronize(); MPI.Barrier(comm); t1 = time_ns()
wall = MPI.Allreduce((t1 - t0) / 1e9, max, comm)
if rank == 0
    zonecycles = prod((N1, N2, N3)) * (simulation.clock.cycle - c0)
    println("aggregate ", zonecycles / wall / 1e6, " Mzone-cycles/s over $nranks GPUs")
end
MPI.Barrier(comm); MPI.Finalize()

• run-problem skill (sibling) has the physics/IC details and the serial/CPU path. This skill only adds the distributed-GPU launch layer.
• Keep Printf out of MPI scripts unless the env has it — the main Aether project has no Printf dep; use println + round.

4. Weak vs strong scaling meshes

Keep a whole number of blocks per rank. For a weak-scaling ladder at 2 blocks per GPU (block B = 128), grow the mesh with the GPU count and scale extent so dx stays constant (extent = size .÷ B, integer so periodic ICs stay valid):

GPUs	nodes	mesh (128³ blocks)	blocks	srun
1	1	256×128×128	2	-N1 --ntasks-per-node=1 -n1 -G1
2	1	256×256×128	4	-N1 --ntasks-per-node=2 -n2 -G2
4	1	256×256×256	8	-N1 --ntasks-per-node=4 -n4 -G4
8	2	512×256×256	16	-N2 --ntasks-per-node=4 -n8 -G8
16	4	512×512×256	32	-N4 --ntasks-per-node=4 -n16 -G16

Rules (from the distributed section of AGENTS/run-problem): cells_per_block must divide size per dimension; total blocks ≥ nranks; blocks are dealt to ranks in contiguous Z-order slices; mesh.nblocks/block_geometry are rank-local.

5. Verify correctness, then trust the number

A run's diagnostics are invariant to rank/block decomposition (ghost exchange copies exact values; the dt reduction is an exact Allreduce min). So the final simulation.clock.t (and any global norm) must match the single-GPU value bit-for-bit across every GPU count — a free regression check that the speedups reflect real, correct work, not dropped exchanges.

6. Measured anchors (A100, MHD PPM + HLLD + RK2, cfl 0.3, 500 cyc, GPU-aware)

Weak scaling, 2×128³ blocks/GPU, Mzone-cycles/s per GPU (2026-07):

GPUs	1	2	4	8	16
Mzone-cyc/s/GPU	179	164	159	158	150

≈ 84 % weak-scaling efficiency at 16 GPUs; nearly flat across the single→multi node jump (159→158→150 at 4→8→16). GPU-aware MPI beat host-staged ≈ 1.4× at 4 GPUs on one node. Single-GPU single-256³-block hydro (PPM+HLLC+RK2) ≈ 337, MHD (PPM+HLLD) ≈ 188 Mzone-cycles/s. First kernel compile adds ~1 min to a cold Julia session — always warm up with a few stop_cycle steps before timing.