perf(mode): pin sliced mode-solver plane to single-device sharding - #401
perf(mode): pin sliced mode-solver plane to single-device sharding#401pic-asso wants to merge 2 commits into
Conversation
On multi-GPU runs the field/material arrays are x-sharded. Slicing a thin transverse plane out of one of these arrays to feed the Tidy3D mode-solver pure_callback (which runs on a single device) could otherwise make XLA all-gather the full ~100M-cell parent array to satisfy the callback's single-device requirement. pin_to_single_device replicates just the already -sliced plane across the mesh instead, so only the thin band is gathered. No-op on single-device runs.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (5)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughAdds replicated sharding support for mode-related arrays. Sliced material tensors and mode ChangesMode callback sharding
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #401 +/- ##
==========================================
- Coverage 90.77% 90.73% -0.05%
==========================================
Files 91 92 +1
Lines 10588 11459 +871
Branches 1582 1730 +148
==========================================
+ Hits 9611 10397 +786
- Misses 686 730 +44
- Partials 291 332 +41 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| inv_permittivity_slice = inv_permittivities[:, *self.grid_slice] | ||
| # Pin the sliced PLANE to a single-device replicated layout (no-op on 1 GPU): the slice | ||
| # feeds the host mode-solver pure_callback (device 0), so this gathers only the thin band | ||
| # rather than letting XLA all-gather the full x-sharded parent array. |
There was a problem hiding this comment.
This comment is duplicated, a bit overdocumented
|
Hi @pic-asso, thanks for the PR! I currently do not have easy access to a multi-gpu setup to verify this behavior. Could you provide some numbers / screenshot /logs or something similar to show that this PR does achieve what it claims? Don't get me wrong, the reasoning and code all makes sense, its just that I currently have no way of verifying. |
pin_to_single_device only covered the mode solver's sliced INPUT plane;
the pure_callback's own OUTPUTS (mode_E_raw, mode_H_raw, eff_idx) carry
a maximal/single-device GSPMDSharding that is incompatible with the
surrounding x-sharded computation and was left unpinned, so a 2-GPU run
still crashed inside compute_mode.
jax.lax.with_sharding_constraint cannot convert that GSPMDSharding
either (raises "Cannot convert GSPMDSharding {maximal device=0} into
SdyArray" via Shardy), so pin_to_single_device now uses jax.device_put,
which handles both the already-NamedSharding sliced-input case and the
maximal-sharded callback-output case.
Verified bit-exact (1-GPU vs 2-GPU overlap identical to full float32
precision) on a non-PDK benchmark, with no speedup from 2 GPUs (fixed
per-step sync + serial host callback dominate) but a real per-device
memory reduction, matching the point of ymahlau#401.
Also: shorten the duplicated pinning comment in sources/mode.py and
detectors/mode.py per review feedback, and extract
_replicated_mesh_sharding() as a separately-testable pure function to
close the patch-coverage gap on sharding.py (mocked-device-list tests
cover the mesh/sharding construction directly; the final device_put
placement still requires genuinely distinct physical devices).
|
Pushed a follow-up commit ( What's new in this commit
Benchmark (real 2×H100 hardware, non-PDK)To validate this against real multi-GPU hardware without touching any confidential PDK geometry, I wrote a generic standalone benchmark: a straight strip waveguide (n_core=2.0 / n_clad=1.44, made-up numbers, no tie to any real platform), ~121M cells, with a Ran 3 phases in one SLURM allocation (1×H100 reference, 2×H100 without the callback-output pinning, 2×H100 with it):
Bit-exact correctness (identical overlap to full float32 precision) between 1-GPU and the fully-patched 2-GPU run — that, plus no longer crashing, is the actual point of this fix. I'd flag the memory/timing columns above as inconclusive on their own: GPU0 only drops 18.3→17.3GB (~5%), the 2-GPU total (30.5GB) is actually higher than the 1-GPU total (18.3GB), and the split is uneven (17.3 vs 13.2) — at this size the fixed overhead (mesh setup, program binary, small non-sharded state) is a large enough fraction of the total that it can obscure a real halving one way or the other. Wall-clock is flat-to-slightly-worse too (~168s vs ~171s). To get a cleaner read, I re-ran the same scene ~1.6× larger (NX 2500→4000, ~194M cells instead of ~121M):
Here the 2-GPU split is genuinely even (17.51 vs 17.51 GB, not 17.3 vs 13.2) and both devices sit below the 1-GPU figure — a much more convincing sharded-not-replicated signature, even though the wall-clock still doesn't improve (same story: fixed per-step sync + the serial callback dominate). I don't want to overclaim a big memory win from either benchmark, but the larger one is a materially cleaner demonstration than the original table, so pairing them here for completeness. Edit: removed a closing claim here about Final wrap-up: while chasing the above I also found that gradient-mode multi-GPU (both |
|
|
||
|
|
||
| def _replicated_mesh_sharding(compute_devices: list[jax.Device]) -> jax.sharding.NamedSharding: | ||
| """Build a REPLICATED NamedSharding over ``compute_devices`` -- NOT SingleDeviceSharding, whose |
There was a problem hiding this comment.
Please remove all of the AI chain-of-thoughts comments. We want to have actual docstrings, not a development story
|
|
||
| # inv_permittivities shape: (3, Nx, Ny, Nz) - slice with component dimension | ||
| inv_permittivity_slice = inv_permittivities[:, *self.grid_slice] | ||
| # Pin the sliced plane to a single device (see pin_to_single_device docstring for why). |
There was a problem hiding this comment.
references to other comments are generally not good style
| assert isinstance(result.sharding, jax.sharding.NamedSharding) | ||
|
|
||
|
|
||
| # ---- _replicated_mesh_sharding ---- |
There was a problem hiding this comment.
remove excessive comments and unnecessary spacing
|
|
||
|
|
||
| class TestReplicatedMeshSharding: | ||
| """Tests for the _replicated_mesh_sharding helper, split out of pin_to_single_device so this |
| _run_in_subprocess_with_devices( | ||
| 2, | ||
| """ | ||
| import jax, jax.numpy as jnp |
There was a problem hiding this comment.
Is there any way to run these unit tests without having python code as strings? This kind of code is hard to maintain, maybe a better variant would be defining a helper function directly as python code and executing that
|
Hi @pic-asso, the |
2 participants
What
On multi-GPU runs, the field/material arrays (
inv_permittivities/inv_permeabilities)are x-sharded across the device mesh.
ModePlaneSourceandModeOverlapDetectorslice athin transverse plane out of these arrays and feed it to the Tidy3D mode-solver, which runs
as a
pure_callbackpinned to a single device.Without an explicit sharding constraint on the sliced plane, XLA can conservatively
all-gather the full x-sharded parent array (~100M cells) to satisfy the callback's
single-device requirement, rather than gathering just the thin sliced plane (a few hundred
KB) it actually needs.
How
Adds
pin_to_single_deviceinfdtdx/core/jax/sharding.py: applies awith_sharding_constraintthat replicates the already-sliced plane across the full device mesh (not a plain
SingleDeviceSharding, which would raise "incompatible devices for jitted computation" againstthe surrounding x-sharded computation). This tells XLA to gather only the small plane.
No-op on single-device runs (
len(jax.devices()) == 1), so the single-GPU path is byte-identicalto before.
Applied at the two mode-solver call sites:
ModePlaneSource.place_on_gridandModeOverlapDetector.update.Tests
tests/unit/core/jax/test_sharding.py(TestPinToSingleDevice, 4 new tests):CPU devices via
XLA_FLAGS=--xla_force_host_platform_device_count=2, since the replicatedpath needs genuinely distinct physical devices — mocking
jax.devices()with a duplicateddevice raises
ValueError: ... must be from distinct deviceson the eagerwith_sharding_constraintcall);NamedSharding(emptyPartitionSpec);inside
jax.jitand pinning it must not raise the "incompatible devices" error.All new tests verified to fail if the multi-device branch is stubbed out (sanity-checked
locally by temporarily forcing the no-op path unconditionally).
Summary by CodeRabbit
New Features
Bug Fixes
Tests