Add NameMatchedAveragedModel: robust EMA for parametrized models - #2098
Open
bkmi wants to merge 4 commits into
Open
Add NameMatchedAveragedModel: robust EMA for parametrized models#2098bkmi wants to merge 4 commits into
bkmi wants to merge 4 commits into
Conversation
torch.optim.swa_utils.AveragedModel.update_parameters iterates parameters
and buffers positionally via `zip(self.module.buffers(), model.buffers())`.
This is fragile when the wrapped module's buffer tree can drift between
AveragedModel construction (deep-copy) and later update_parameters calls —
which happens naturally with nn.utils.parametrize.register_parametrization,
shared-mask setups, mid-training buffer re-registration, or DDP/FSDP
interactions with deep-copy in some PyTorch releases.
The symptomatic failure is a cryptic shape error at swa_utils.py's buffer
copy line, e.g.:
RuntimeError: The size of tensor a (148) must match the size of
tensor b (64) at non-singleton dimension 0
NameMatchedAveragedModel is a drop-in subclass that matches by name via
`named_parameters()` / `named_buffers()` dictionaries. Deep-copy guarantees
identical name sets at construction time, so name-matched sync is robust to
any of the reorderings above.
In the healthy case where positional and name-matched iteration would agree,
NameMatchedAveragedModel is bit-exact with AveragedModel — verified via a
parity test with max_abs_diff = 0.0 on plain and parametrized toy models.
Any post-construction name-set or shape divergence raises a RuntimeError
naming the offending buffer/parameter. Silent skipping would let the EMA
slowly diverge from the live model in ways that are very hard to detect
downstream; failing loudly is the safer default.
Follow-up to the initial NameMatchedAveragedModel commit. Fixes a subtler failure mode: named_buffers()'s default remove_duplicate=True yields each unique buffer tensor once under its FIRST-seen name. When a shared buffer is re-pointed post-deepcopy — e.g. a sparsifier that consolidates per-parametrization mask buffers to a single project-owned tensor — the "first name" surviving dedup can differ between live and EMA even though the underlying values would sync fine. That produced spurious "buffer name sets differ" errors AND, before this fix, could silently skip the sync entirely, leaving the EMA copy with stale mask references while the live model evolved. Using remove_duplicate=False on both sides guarantees every registered path is walked. Shared tensors get synced under all their names — redundant writes are cheap and correct. Also switch the (very rare) mismatch error to print the FULL diverging name sets instead of truncating to the first 5, so diagnosing a real inconsistency doesn't require re-running with extra instrumentation.
added 2 commits
July 22, 2026 15:40
Move `Tensor` and `Module` imports into a TYPE_CHECKING block since `from __future__ import annotations` is already in effect. Drop the unused `import torch`. Apply `ruff format` on the file.
The file was accidentally placed at fairchem/src/fairchem/core/common/ema.py (a stray fairchem/ prefix from a checkout confusion), which is outside the paths ruff.toml globs — so its per-file-ignores never applied and, when CI enumerated the file explicitly, it was treated as if outside the project include, tripping several rules. Move to src/fairchem/core/common/ema.py where the rest of the module lives. Verified: `ruff check` / `ruff format --check` both pass under the project's ruff.toml at the new path.
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.
torch.optim.swa_utils.AveragedModel.update_parameters iterates parameters and buffers positionally via
zip(self.module.buffers(), model.buffers()). This is fragile when the wrapped module's buffer tree can drift between AveragedModel construction (deep-copy) and later update_parameters calls — which happens naturally with nn.utils.parametrize.register_parametrization, shared-mask setups, mid-training buffer re-registration, or DDP/FSDP interactions with deep-copy in some PyTorch releases.The symptomatic failure is a cryptic shape error at swa_utils.py's buffer copy line, e.g.:
NameMatchedAveragedModel is a drop-in subclass that matches by name via
named_parameters()/named_buffers()dictionaries. Deep-copy guarantees identical name sets at construction time, so name-matched sync is robust to any of the reorderings above.In the healthy case where positional and name-matched iteration would agree, NameMatchedAveragedModel is bit-exact with AveragedModel — verified via a parity test with max_abs_diff = 0.0 on plain and parametrized toy models.
Any post-construction name-set or shape divergence raises a RuntimeError naming the offending buffer/parameter. Silent skipping would let the EMA slowly diverge from the live model in ways that are very hard to detect downstream; failing loudly is the safer default.