Shard the baseline axis so visibility memory divides across devices - #243
Draft
chrisfinlay wants to merge 16 commits into
Draft
chrisfinlay wants to merge 16 commits into
chrisfinlay wants to merge 16 commits into
Conversation
Documentation build overview
21 files changed ·
|
Sharding sources divides almost nothing. The psum in psum_over_rfi is declared out_specs=P(), so the result is replicated and every component after it runs outside the mesh on replicated arrays. Measured on four GH200s at 512 stations: the sharded work fell from 486.84 s to 133.05 s, but peak memory on the limiting device went from 36.02 GB to 37.80 GB. The only arrays that divide are rfi_A and rfi_phase, 0.24 GB of 36. Everything large carries a baseline axis and no source axis: the four visibility arrays at 1.256 GB each and the GP's mode arrays. Splitting baselines divides all of them, stops the sky model being computed once per device, and needs no cross-device sum at all -- each device keeps the baselines it computed, and only the scalar likelihood is reduced. This is the mesh and the map. The operator needs no change to use them: it already takes its index arrays as traced arguments and sizes its output from a1.shape[0], so inside the map each device sees its own baselines. The single-device path stays the function itself.
Each device takes a share of the group's baselines and none of anyone else's, so nothing needs summing across devices afterwards. That is the structural difference from sharding sources, where every device computes every baseline for a few sources and the partials must be added back -- a psum of the whole 1.256 GB visibility array on every iteration. The antenna set is deliberately not recompacted per device, which is what separates this from make_poly_time_group. Every device keeps the whole group's antennas so the per-antenna signal is identical and enters the map replicated; recompacting would give each device its own antenna count, and shard_map runs one program with one set of shapes. That signal is 0.24 GB at 512 stations against 1.256 GB for one visibility array, so replicating it is what buys dividing the visibilities. An uneven split raises and names the counts rather than dropping the remainder silently.
shard_map runs one program, so every device's operator call needs the same baseline count, and a group rarely divides evenly. There is no real baseline to pad with: the operator refuses a repeated (a1, a2), and at 512 stations all 130816 pairs are already in the list. A ghost antenna supplies fresh ones. Pairs (j, ghost) collide with nothing, distinct partners keep them distinct from each other, and the ghost carries no signal, so they contribute nothing and take no gradient -- the same device padded_rfi_count uses on the source axis. One ghost antenna covers any device count a real array will meet, since a device needs at most n_dev - 1 of them. The ghosts live only inside the operator call. n_real says how many of its output rows are real, and the rest are dropped before anything is written, so a ghost never reaches the visibility array and needs no place in it.
The first version chunked a group's baselines into equal sequential runs, which ignores where those baselines actually sit. A device owns a contiguous range of the visibility array in the order the data already has, so which of a group's baselines it holds is whatever the ordering puts there -- not an equal share, and not a share it can choose. Splitting by ownership is what lets the data stay as it is. The observed visibilities, the flags and the noise are sharded along the axis they already have, every array downstream keeps that split through to the likelihood, and nothing is reordered or gathered back: a device never needs anyone else's baselines, and only the scalar likelihood is reduced. Positions are local to the owning block, so a group writes into its device's own rows. Ownership makes the per-device counts uneven, which is what the ghost padding is for.
The scatter that places a group's results has to have a fixed shape, like everything else inside shard_map, but the number of real baselines a device owns varies. Padding the positions needs an index that is safe to write: any real row would have a ghost's zero overwrite a baseline another group had already computed. One past the block's last row is that index. The caller gives its local visibility array one spare row, lets every ghost land there and drops it, so all devices scatter the same number of rows and a ghost can never reach a real baseline. One row against a block of 32704 at 512 stations.
The route is now wired end to end. A device is handed the whole per-antenna signal and the index arrays for its own rows, computes the baselines in its block of the visibility array, and nothing is summed across devices: the map's output is the visibility array itself, one block per device. The source route is untouched and stays the default; TABASCAL_SHARD_AXIS picks between them. Index arrays go in as data. Those with a baseline axis are concatenated so the map hands each device its slice; the rest are stacked one row per device. Their values are built per device rather than sliced from the whole -- a sorter sorts the local baselines and pair_index names a local row, so a slice of the global arrays would point outside the shard and be wrong in silence. The observed visibilities, the flags and a resolved noise follow the same axis, so they divide too. A scalar noise has no baseline axis and stays replicated. The equivalence test is the point of it: on four devices, through the real component, the result matches what one device computes -- with ghost padding and without. It also asserts the sharded path was actually built, because the first version of that test patched the route away and compared the unsharded answer with itself.
The hybrid borrows build_forward from the quadrature route but derives from PolyInterpVisHybrid, so it inherited none of the sharded machinery and died reaching for it -- which is what the first 512-station run hit. It borrows those the same way now, and chooses per group which operator and which seam a device's baselines go through. That needed the closed form to expose the same seam, so analytic_eval_with_indices sits beside the quadrature one in ri_kernels. RFIAnalyticVisOp already shares its index arrays with RFIInterpVisOp, so the two take the same tuple. Without it the all-analytic route -- which is what the benchmarks run -- could not shard at all. The equivalence test now covers both routes. The hybrid case sets a cut so two groups actually form: left at its default everything goes to quadrature as one group, which divides evenly, and the padding the test exists to exercise would never be built.
The operator's index arrays have to match the signal it is handed. A group whose baselines divide evenly needs no padding and gets no ghost row, so building its operator one antenna wider leaves the indices longer than the inputs, which it rejects outright. That only shows on a configuration needing no padding -- 130816 baselines over four devices at 512 stations -- so the cases built to exercise the padding could never have caught it. The test now asserts the two agree on a case that divides.
The RFI visibility leaves the map already split, but the astronomical model and the gains are ordinary JAX built from parameters that arrive replicated, so XLA had no reason to divide what they produce: it computed a whole (n_bl, n_freq, n_time) array on every device and resharded to add the two together. Measured at 512 stations on four devices, that left the workers at 21.941 GB where the arrays alone should put them near a quarter of the 36.018 GB a single device holds, and made the run 9% slower than sharding sources rather than faster. Saying where the per-baseline entries belong, between components, is what keeps one layout across the whole forward. The counterpart of constrain_rfi_state for the other axis, and a no-op unless this run shards baselines.
The likelihood stacks real and imaginary parts before comparing them, which puts the baseline axis second. A constraint written for the leading axis cannot see that, so the log-probability, the mask, and the cotangents of all three -- the expensive half -- were free to be computed whole on every device even with every input to them divided.
Two limits in the first version were left as errors rather than handled. Padding assumed one dark antenna's pairs would cover any shortfall, and an indivisible baseline total was refused outright. Neither holds in general: ownership imbalance can exceed the antenna count, so the ghost antennas are now sized from it, and an indivisible total rounds the internal block up and trims the overhang. Placement is the larger change. shard_pytree now puts the astronomical parameters and constants on the baseline layout at init, so the optimiser buffers follow, rather than leaving them replicated for a constraint between components to chase. make_global reshards an existing global array collectively instead of round-tripping it through numpy, which a partly addressable array cannot survive. padded_rfi_count and the replicated sharding follow whichever axis the run splits, and the two kernels' capabilities are checked independently. The shipped quadrature_limit stays null, the measured crossover. All analytic is right at 512 stations and wrong below it -- 179 s against 57 s at 128 -- so the benchmark configurations set 0 explicitly rather than every run inheriting it.
512 stations now fit on four devices at 32 and 64 channels, which they did not before: 64 channels died on a single 18.71 GiB allocation. At 8 channels the limiting device holds 13.703 GB against 33.911 for source sharding, and the reduced chi-squared is unchanged at 83.42597198486328 on both axes. The memory was never in the optimiser step. That step needed 4.81 GB per device while each device peaked at 18, so four earlier attempts inside the model forward moved between 0.7 and 3.2 GB each. It was in setup and diagnostics: - GPVisAst encoded the padded FFT over every baseline before the blocked forward existed, allocating four whole visibility arrays. Bounding the batch at the existing sky budget removes that. - The sky scan padded and reshaped globally, which XLA satisfied with three latent all-gathers and a visibility-cotangent all-gather of 1.258 GB. Running the scan inside map_over_baselines keeps both the padding and the transpose local; the optimised step then contains no all-gather or collective-permute at all. - Reduced chi-squared and the truth moments compacted with a boolean mask, whose index vectors came to 1.884 GB. Masked reductions give the same numbers without them. - Components kept zero placeholders in state_outputs, so replacing the assembled model's state left whole cubes alive on device 0. Placing component outputs at setup puts them on their proper shards. What this does not fix: optimisation is still 6.1% slower than sharding sources, 140.97 s against 132.82. Removing every gather did not close it. Baseline workers materialise all RFI sources on the group antenna axis, and that redundant per-antenna work is a different problem. 64 channels finishes with 2.315 GB of a 96.905 GB pool to spare, so it is a result rather than headroom: a larger scratch budget, a longer observation or more stations will not fit as it stands.
A deterministic host-side plan that gives a complete four-worker cross-correlation graph a common antenna width of 308, against the 512 every worker carries today. That redundant per-antenna materialisation is the remaining explanation for baseline sharding being 6.1% slower than sharding sources. Nothing imports it and the runner does not use it. It is committed so the construction and its tests are not lost, not because it is ready: enabling it changes baseline ordering end to end, and it has not been runtime-benchmarked. Treat the numbers in its tests as properties of the construction, not of a fit.
Ragged masks, excluded NaNs and zero noise, all-flagged data, resolved noise, and integer-valued input -- where the sample count must widen to at least float32 rather than narrow to the data's dtype.
… divide TABASCAL_SHARD_AXIS was undocumented, and the only prose about baseline sharding was a paragraph at the end of config.md, under the gains section, about something that is not a config key. Move it to the usage guide beside the other environment switches, with the measurement and the divisibility caveat: a flagged or selected subset falls back to replicated placement and saves no memory.
chrisfinlay
force-pushed
the
poly-interp-hybrid
branch
from
September 13, 2026 14:40
9884ec5 to
adfca99
Compare
chrisfinlay
force-pushed
the
baseline-sharding
branch
from
September 13, 2026 14:40
30cb6c8 to
a7fd5e1
Compare
The analytic route's doc still said source-sharded GPU execution and peak memory needed validating. Job 4657092 did that on four GH200s, and this branch carries the measurements, so point at them instead.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #240. Depends on
ri-kernels#11, which
exposes the index arrays both operators need so a caller can shard baselines.
The result
512 stations at 32 and 64 channels now fit on four GH200s. They did not
before: 64 channels died on a single 18.71 GiB allocation.
Job 4657092 on four GH200s: all 102 iterations plus prediction, export and truth
diagnostics;
--skip-ms-write, analytic-only RFI,GPVisAst. The 8-channel χ²is identical on both axes and to a single device, so four devices each computing
a different quarter of the baselines agree to sixteen digits.
Why the source axis did not divide memory
psum_over_rfideclaresout_specs=P(), so its result is replicated and everycomponent after it runs outside the mesh. Every visibility-shaped array carries a
baseline axis and no source axis. Measured on four GH200s at 512 stations,
that left the optimisation scaling 3.66× while peak memory on the limiting
device went up, 36.02 → 37.80 GB.
Sharding baselines instead: each device owns a contiguous range of the visibility
array in the order the data already has, computes only its own rows, and nothing
is gathered back — only the scalar likelihood is reduced. Uneven per-device group
counts are padded with ghost baselines against dark antennas carrying no signal,
the same device
padded_rfi_countuses on the source axis; ghosts live onlyinside the operator call and are sliced off before the scatter.
TABASCAL_SHARD_AXIS=source|baselineselects the route.sourceremains thedefault and a single-device run is unaffected either way.
Where the memory actually was
Not in the optimiser step. That step needed 4.81 GB per device while each device
peaked at 18, so four successive attempts inside the model forward moved between
0.7 and 3.2 GB each. It was in setup and diagnostics: a padded FFT encoding every
baseline before the blocked forward existed; a global scan XLA satisfied with
three latent all-gathers and a 1.258 GB visibility-cotangent all-gather; boolean
compaction whose index vectors came to 1.884 GB; and component zero placeholders
left alive as whole cubes on device 0. The optimised step now contains no
all-gather or collective-permute at all.
What this does not fix
132.82). Removing every gather did not close it. Baseline workers materialise
all RFI sources on the group antenna axis, and that redundant per-antenna work
is a separate problem.
tabascal/baseline_partition.pycontains adeterministic antenna-reuse plan that would cut the common antenna width from
512 to 308 on a complete four-worker graph, but it is not wired to anything
and not runtime-benchmarked.
result, not headroom: a larger scratch budget, a longer observation or more
stations will not fit as it stands.
peak_bytes_in_useinside a preallocated pool. On asingle device that understated the true footprint by 29 GB (65.19 GB by
nvidia-smiagainst 36.02 reported). They are comparable to each other, notabsolute headroom.