Gradient-mode multi-GPU: both "reversible" and "checkpointed" replicate the backward-pass state instead of sharding it

[pic-asso]

Summary

Neither of fdtdx's gradient-computation methods (GradientConfig.method="reversible" or "checkpointed") shards its backward-pass state across multiple GPUs — both replicate the full per-device buffer on every device instead of splitting it along the x-sharded axis. This is a different issue from #401 (which is specifically about the mode-solver's pure_callback, already fixed): this one shows up even in scenes that never touch the mode solver at all.

Reproduction (generic, non-PDK)

Straight strip waveguide (n_core=2.0 / n_clad=1.44, made-up numbers), a GaussianPlaneSource, a PoyntingFluxDetector (scalar FOM), and a trainable fdtdx.Device region — jax.value_and_grad over the device's latent parameters, matching the standard adjoint-optimization pattern (place_objects → apply_params → run_fdtd(gradient_config=...) → scalar reduction → grad).

config = fdtdx.SimulationConfig(...)
config = config.aset("gradient_config", fdtdx.GradientConfig(
    method="reversible",
    recorder=fdtdx.Recorder(modules=[
        fdtdx.LinearReconstructEveryK(k=5),
        fdtdx.DtypeConversion(dtype=jnp.bfloat16),
    ])))

16.0M-cell domain (NX=1600, NY=100, NZ=100), 7343 time steps, 2×H100:

Devices	Peak GPU memory	value_and_grad wall time	fom	grad_norm
1×H100	18.3 GB	65.1s	3.574450e-14	3.554558e-16
2×H100	42.7 GB / 43.2 GB (per device — not halved)	477.0s (7.3× slower)	3.574450e-14	3.554558e-16

Correctness is bit-exact between 1 and 2 GPUs (identical fom/grad_norm), so this isn't a numerics bug — the gradient is right, it's just that adding a second GPU replicates the whole buffer instead of splitting it, and the extra device-to-device synchronization makes things slower, not faster.

I tried the other gradient method (method="checkpointed") on the same scene expecting it might avoid this (it reuses checkpointed_fdtd, the same function used for plain forward-only runs, which does shard cleanly on 2 GPUs). It also replicates: 1×H100 and 2×H100 both peak at the same ~34.7GB for a given num_checkpoints. Separately, sweeping num_checkpoints on a larger real-world domain gave results that don't match the classical binomial-checkpointing memory/recompute tradeoff (num_checkpoints scaled up from 5→90 made memory much worse — 230GB→3.48TB — rather than better), so there may be a second, separate issue in how num_checkpoints interacts with memory in this code path; I haven't dug into that further and don't want to conflate an unconfirmed observation with the reproducible one above.

What I ruled out

• Not a PhasorDetector/mode-solver issue — this scene doesn't use either. Confirmed separately (see perf(mode): pin sliced mode-solver plane to single-device sharding #401's PR thread) that a scene with a PhasorDetector and mode-solving sources/detectors shards fine on the forward-only path once perf(mode): pin sliced mode-solver plane to single-device sharding #401's fix is applied.
• Not a correctness issue — bit-exact fom/grad match between 1 and 2 GPUs.
• Forward-only sharding is fine — the same domain, run without gradient_config (checkpointed_fdtd's plain forward path), shards evenly and shows a real per-device memory reduction as domain size grows (see perf(mode): pin sliced mode-solver plane to single-device sharding #401's PR thread for a larger, cleanly-split example). The gap is specific to the differentiable/backward path.

What I haven't done

• Root-caused why — my best guess is that reversible_fdtd's backward pass (fdtd.py's fdtd_bwd/body_fn, which nests a jax.vjp of a single step inside equinox.internal.while_loop(kind="lax")) loses track of the x-sharded NamedSharding somewhere in that chain, causing XLA's auto-partitioner to fall back to full replication. I tried adding explicit jax.lax.with_sharding_constraint calls at the loop-carry boundaries (re-asserting the original sharding on the way in and out of the backward loop body) — it didn't change the memory footprint at all (still fully replicated) and made things slower, so whatever's happening isn't fixed by a local hint at that point; the actual sharding decision seems to be made more globally by XLA across the whole custom_vjp-wrapped program. I don't have a fix, just the repro above.
• I haven't tried jax.experimental.shard_map (explicit SPMD, no auto-partitioner ambiguity) — that seems like the more likely real fix but is a bigger structural change than I wanted to make without checking with you first.

Happy to share the benchmark script if useful — it's fully generic/non-PDK (a handful of made-up numbers, no tie to any real platform), same spirit as the one in #401's PR thread.

[ymahlau]

I think there are a few things to untangle here: Checkpointing is an internal mechanism (we use the checkpointed while loop from a different library), so it will be very hard to add custom sharding to that.

For reversible mode it is technically possible to shard the recording state, but I think it would not make much sense to do so at the current state. This would require computing where exactly the PML corresponds to the part in the simulation that is recorded and then sharding somehow along that axis. That sharding would be inbalanced all the time (some shards always have more pml overlap than others).

I would put this as a feature that we can look at after the pytorch re-write (#349), but for now I don't think the implementation effort would be justified. In pytorch the recording state could live on cpu anyways, which would make the whole mechanism obsolete.