Skip to content

Make the adjoint design-region gradient independent of the MPI chunk division - #3264

Open
jin-castle wants to merge 3 commits into
NanoComp:masterfrom
jin-castle:pr/np-invariant-adjoint-gradient
Open

Make the adjoint design-region gradient independent of the MPI chunk division#3264
jin-castle wants to merge 3 commits into
NanoComp:masterfrom
jin-castle:pr/np-invariant-adjoint-gradient

Conversation

@jin-castle

Copy link
Copy Markdown

Fixes #2578 (adjoint gradients inconsistent with finite differences in 3D).

Symptom. The adjoint gradient of a MaterialGrid design region depends on
the number of MPI processes while the objective value does not. Measured on a
20×20×6 design grid (3D waveguide, both with and without subpixel smoothing):

  • f0 identical to 8+ digits at every process count;
  • gradient max-norm difference between -np 1 and -np 8: 57% of
    |g|_max (smoothing off) / 18% (smoothing on), spread over ~80% of the
    design nodes in chunk-sized patches;
  • directional finite-difference checks that pass at -np 1 (ratio 1.0002)
    fail under MPI (ratios 0.86–1.22 at -np 16).

Root causes — all in material_grids_addgradient(), which computed the
result from per-chunk DFT views:

  1. chunks of different components were paired by list index
    (forward_dft_chunks[ci_forward][cur_chunk]); the per-component chunk sets
    can legitimately differ near the monitor edges, after which every
    subsequent pair associates the wrong regions, and cross terms were dropped
    entirely when the counts differed (cur_chunk >= num_f_chunks → continue);
  2. neighboring-point lookups bounds-checked only the chunk-local linear
    index, so an out-of-range point in one dimension could alias to a valid
    index of a different point instead of returning zero;
  3. the ±1-pixel padding of persist dft chunks is clamped to the owning
    fields chunk and cannot cross process boundaries, so the data those
    lookups need does not exist at rank boundaries;
  4. the same-component term indexed the forward chunk with the adjoint loop
    index, relying on the chunk-alignment assumption of (1).

Fix. Gather each monitor component onto its full (padded) design-region
grid first — monitor-owned points authoritatively (the is_old..ie_old
decomposition is disjoint), the padding ring as a multiplicity-average of the
ghost copies so its coverage does not depend on the chunk division — then run
the same per-point math on this chunk-independent view, with the work strided
across processes and per-dimension bounds checks on neighbor lookups. Memory
cost is the gathered design-region DFTs (6 complex arrays over the monitor
grid per frequency), negligible for realistic design regions.

Validation.

check before after
-np 8 vs -np 1, 20×20×6, smoothing off 57% 7×10⁻¹²
-np 8/16/32/128 vs -np 1 (both smoothing cases) up to 57% ≤7×10⁻¹²
FD-vs-adjoint at -np 16 and -np 128 (six grid configurations incl. incommensurate/offset/coarse/fine grids) ratios 0.86–1.22 identical to serial (1.000 ± FD noise)
serial behavior preserved (test_subpixel_3d, adjoint E2E unchanged)

🤖 Generated with Claude Code

material_grids_addgradient computed its result from per-chunk DFT views:
chunks of different components were paired by list index, cross terms were
dropped when per-component chunk counts differed, neighboring-point lookups
used chunk-local linear indices that could alias across dimensions or
silently return zero at chunk boundaries, and the +-1-pixel padding of
persist dft chunks cannot cross rank boundaries. Together these made the
adjoint gradient depend on the number of MPI processes (measured: up to
57% of |g|max on a 20x20x6 3D design region between -np 1 and -np 8, with
the objective value identical; see NanoComp#2578).

The computation now first gathers each monitor component onto its full
(padded) design-region grid -- owned points authoritatively, the padding
ring as a multiplicity-average of the ghost copies -- and then evaluates
the same per-point math on this chunk-independent view, with the work
strided across processes and per-dimension bounds checks on neighbor
lookups. Results are now identical (up to summation roundoff) for any
number of processes, and serial semantics are preserved.
@jin-castle
jin-castle force-pushed the pr/np-invariant-adjoint-gradient branch from 299afb7 to 6d39640 Compare August 9, 2026 02:38
The rewrite in the previous commit has no coverage: nothing in the suite
fails if material_grids_addgradient goes back to reading the per-chunk DFT
views. Add a test that computes the design-region gradient twice for the
same problem under two different chunk divisions and requires the two
vectors to agree.

Two things make this cheap to run. The defects are triggered by a chunk
boundary crossing the design region, not by MPI as such, so forcing the
split with Simulation(num_chunks=...) reproduces them in a serial run --
the test therefore belongs in TESTS and guards the serial jobs too.  And
comparing the gradient vectors is far more sensitive than the usual
directional finite-difference check, which probes along g/|g| and so
largely cancels the error: measured against the pre-fix code, the gradient
was off by 44% in L2 at four chunks while fd/adjoint still read 1.0008.

Measured with this test: 4.7e-1 relative difference before the rewrite,
3e-16 after (serial), 45 s runtime.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Inconsistent finite-difference and adjoint gradients for a simple 3D example

2 participants