Two independent memory reductions for SWIFT on high-dimensional data.#91
Open
copybara-service[bot] wants to merge 1 commit into
Open
Two independent memory reductions for SWIFT on high-dimensional data.#91copybara-service[bot] wants to merge 1 commit into
copybara-service[bot] wants to merge 1 commit into
Conversation
1. Memory-efficient marginal oracle for final estimation. SWIFT's final MirrorDescent estimation hardcoded the HUGIN marginal oracle (message_passing_stable), which materializes every junction-tree clique belief simultaneously. On high-dimensional data with large maximal cliques this drives peak memory to the sum of all clique beliefs (tens of GiB) and is numerically unstable with -inf potentials (HUGIN uses belief subtraction). Switch to message_passing_implicit with the einsum_materialized contraction: - Peak memory is bounded by the largest single clique super-factor rather than the sum of all beliefs. - Numerically stable: log-space add + logsumexp, no exp/log round-trip and no belief subtraction. - Fastest contraction on GPU. Both oracles perform exact junction-tree inference, so estimated marginals are identical up to floating-point error. 2. Free precomputed candidate marginals after measurement. The candidate marginals produced by from_projectable (plus the workload intermediates) are only needed through query selection and measurement, but they stayed live for the rest of __call__ -- through estimation and the column-by-column generation phase. On host-memory-constrained slices this residual is enough to push generation over the host RAM limit. Free them right after measurement so the peak during estimation and generation is lower. PiperOrigin-RevId: 947706110
copybara-service
Bot
force-pushed
the
cl/947706110
branch
from
f4b9602 to
1a07bc7
Compare
copybara-service
Bot
changed the title
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
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.
Two independent memory reductions for SWIFT on high-dimensional data.
SWIFT's final MirrorDescent estimation hardcoded the HUGIN marginal oracle
(message_passing_stable), which materializes every junction-tree clique
belief simultaneously. On high-dimensional data with large maximal cliques
this drives peak memory to the sum of all clique beliefs (tens of GiB) and
is numerically unstable with -inf potentials (HUGIN uses belief subtraction).
Switch to message_passing_implicit with the einsum_materialized contraction:
than the sum of all beliefs.
no belief subtraction.
Both oracles perform exact junction-tree inference, so estimated marginals
are identical up to floating-point error.
The candidate marginals produced by from_projectable (plus the workload
intermediates) are only needed through query selection and measurement, but
they stayed live for the rest of call -- through estimation and the
column-by-column generation phase. On host-memory-constrained slices this
residual is enough to push generation over the host RAM limit. Free them
right after measurement so the peak during estimation and generation is
lower.