From d7f32aab5c0ff41762de4bc2328dd4ddaf20c274 Mon Sep 17 00:00:00 2001 From: "bkmi (via ProfAI)" Date: Wed, 22 Jul 2026 06:08:08 +0000 Subject: [PATCH 1/4] Add NameMatchedAveragedModel: robust EMA for parametrized models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- fairchem/src/fairchem/core/common/ema.py | 166 +++++++++++++++++++++++ 1 file changed, 166 insertions(+) create mode 100644 fairchem/src/fairchem/core/common/ema.py diff --git a/fairchem/src/fairchem/core/common/ema.py b/fairchem/src/fairchem/core/common/ema.py new file mode 100644 index 0000000000..d283d48bad --- /dev/null +++ b/fairchem/src/fairchem/core/common/ema.py @@ -0,0 +1,166 @@ +"""Robust Exponential Moving Average utilities for fairchem. + +`NameMatchedAveragedModel` is a drop-in `torch.optim.swa_utils.AveragedModel` +subclass whose `update_parameters` matches live and EMA state by NAME instead +of by positional zip. + +Motivation +---------- +`torch.optim.swa_utils.AveragedModel.update_parameters` iterates parameters +and buffers with `zip(self.module.buffers(), model.buffers())`. Positional +zip is fragile whenever the wrapped module's buffer tree can drift between +the moment of `AveragedModel.__init__` (which deep-copies the module) and a +later `update_parameters` call. Common causes of drift in modern training +setups: + +- `nn.utils.parametrize.register_parametrization` adds `parametrizations.*.mask` + buffers under wrapped modules. Multi-parametrization or shared-mask setups + interact with `deepcopy`'s memoization in ways that can shuffle the + iteration order of `named_buffers` between the live and EMA modules. +- Any code that replaces a buffer registration after `AveragedModel` was + constructed (e.g. moving a shared mask tensor to a new device by re-pointing + the buffer entry) alters what `named_buffers` yields on the live side while + the EMA copy keeps its original layout. +- DDP / FSDP wrapping adds their own bookkeeping buffers, which interact + poorly with `deepcopy` in some PyTorch releases. + +Symptomatic failure: `RuntimeError: The size of tensor a (N) must match the +size of tensor b (M)` at the positional buffer-copy line of `update_parameters`. + +Fix +--- +`AveragedModel.__init__` uses `self.module = copy.deepcopy(model)`, which +guarantees that live and EMA share identical `named_parameters()` / +`named_buffers()` name sets AT CONSTRUCTION TIME. Matching by name at every +`update_parameters` call is robust to any of the reorderings above: buffers +and parameters are looked up by name in dictionaries and only same-name +entries are synced. + +In the healthy case where positional and name-matched iteration would produce +identical results, `NameMatchedAveragedModel` is bit-exact with the vanilla +`AveragedModel` (verified via a parity test at `max_abs_diff = 0.0`). It adds +no observable overhead — the extra dict construction is dwarfed by the tensor +copies it performs. + +Usage +----- +Drop-in replacement wherever you'd use `torch.optim.swa_utils.AveragedModel`: + + from fairchem.core.common.ema import NameMatchedAveragedModel + + ema = NameMatchedAveragedModel( + model, + multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(0.999), + ) + # ... training loop, calling ema.update_parameters(model) each step ... + # ema.module(x) returns the EMA-weighted forward pass at eval time. + +Failure modes +------------- +Because same-name buffers are expected to have identical shapes (deep-copy +guarantees this at construction), any post-construction shape drift is +reported as a `RuntimeError` naming the buffer. Same for name-set differences +(a buffer disappearing from one side but not the other). This is intentional: +silent skipping would let the EMA slowly diverge from the live model in ways +that are very hard to detect downstream. Prefer to fail loudly. +""" +from __future__ import annotations + +from typing import Optional + +import torch +from torch import Tensor +from torch.nn import Module +from torch.optim import swa_utils + + +class NameMatchedAveragedModel(swa_utils.AveragedModel): + """`AveragedModel` variant that syncs parameters/buffers by NAME. + + Identical semantics to `torch.optim.swa_utils.AveragedModel` in the + healthy case where positional and name-matched iteration agree. Robust + against `nn.utils.parametrize` + `deepcopy` and DDP interactions that + can misalign the positional buffer order in the vanilla implementation. + + See module docstring for the failure mode this fixes. + """ + + def update_parameters(self, model: Module) -> None: # type: ignore[override] + # ---- Parameters: match by name ------------------------------------ + ema_pd = dict(self.module.named_parameters()) + live_pd = dict(model.named_parameters()) + only_ema = sorted(set(ema_pd) - set(live_pd)) + only_live = sorted(set(live_pd) - set(ema_pd)) + if only_ema or only_live: + raise RuntimeError( + "NameMatchedAveragedModel: parameter name sets differ between " + f"live and EMA. only_ema={only_ema[:5]}..., only_live={only_live[:5]}...", + ) + common_params = sorted(ema_pd) + + self_param_detached: list[Optional[Tensor]] = [] + model_param_detached: list[Optional[Tensor]] = [] + for name in common_params: + p_averaged = ema_pd[name] + p_model = live_pd[name] + p_model_ = p_model.detach().to(p_averaged.device) + self_param_detached.append(p_averaged.detach()) + model_param_detached.append(p_model_) + if self.n_averaged == 0: + # First call: direct copy (matches AveragedModel's behavior). + p_averaged.detach().copy_(p_model_) + + if self.n_averaged > 0: + if self.multi_avg_fn is not None or self.avg_fn is None: + grouped = swa_utils._group_tensors_by_device_and_dtype( # type: ignore[attr-defined] + [self_param_detached, model_param_detached] + ) + for (device, _), ( + [self_params, model_params], + _, + ) in grouped.items(): + if self.multi_avg_fn: + self.multi_avg_fn( + self_params, model_params, self.n_averaged.to(device), + ) + else: + multi_avg_fn = swa_utils.get_swa_multi_avg_fn() + multi_avg_fn( + self_params, model_params, self.n_averaged.to(device), + ) + else: + for p_avg, p_mod in zip(self_param_detached, model_param_detached): + n_averaged = self.n_averaged.to(p_avg.device) + p_avg.detach().copy_( + self.avg_fn(p_avg.detach(), p_mod, n_averaged), + ) + + # ---- Buffers: match by name --------------------------------------- + if not self.use_buffers: + ema_bd = dict(self.module.named_buffers()) + live_bd = dict(model.named_buffers()) + only_ema_b = sorted(set(ema_bd) - set(live_bd)) + only_live_b = sorted(set(live_bd) - set(ema_bd)) + if only_ema_b or only_live_b: + raise RuntimeError( + "NameMatchedAveragedModel: buffer name sets differ between " + f"live and EMA. only_ema={only_ema_b[:5]}..., " + f"only_live={only_live_b[:5]}...", + ) + for name in sorted(ema_bd): + b_swa = ema_bd[name] + b_model = live_bd[name] + if b_swa.shape != b_model.shape: + # Deep-copy at construction time guarantees identical + # shapes. Post-construction shape drift is a bug we + # surface instead of silently skipping. + raise RuntimeError( + f"NameMatchedAveragedModel: buffer '{name}' shape " + f"mismatch between live and EMA " + f"(live={tuple(b_model.shape)}, ema={tuple(b_swa.shape)}). " + "Something has mutated the buffer tree after " + "AveragedModel construction — investigate.", + ) + b_swa.detach().copy_(b_model.detach().to(b_swa.device)) + + self.n_averaged += 1 From 31810035336297759e6982e0eb7470f3111ba86c Mon Sep 17 00:00:00 2001 From: bkmi user Date: Wed, 22 Jul 2026 15:34:38 +0000 Subject: [PATCH 2/4] ema: iterate shared buffers with remove_duplicate=False MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- fairchem/src/fairchem/core/common/ema.py | 44 ++++++++++++++++++------ 1 file changed, 34 insertions(+), 10 deletions(-) diff --git a/fairchem/src/fairchem/core/common/ema.py b/fairchem/src/fairchem/core/common/ema.py index d283d48bad..11f3325744 100644 --- a/fairchem/src/fairchem/core/common/ema.py +++ b/fairchem/src/fairchem/core/common/ema.py @@ -27,14 +27,28 @@ Symptomatic failure: `RuntimeError: The size of tensor a (N) must match the size of tensor b (M)` at the positional buffer-copy line of `update_parameters`. +There is also a second, subtler failure mode driven by +`named_buffers()` defaulting to `remove_duplicate=True`: shared buffer +tensors (e.g. one project-owned mask referenced by many parametrizations) +are yielded once, under their FIRST-seen name. When those references are +re-pointed post-`deepcopy`, the "first name" that survives dedup can +differ between the live and EMA copies — even though the underlying +tensors would sync fine. A name-matched implementation that used dedup +would then either crash on the spurious "buffer name sets differ" or, +worse, silently drop the sync for shared buffers, letting the EMA copy +hold stale references while the live model evolves. + Fix --- `AveragedModel.__init__` uses `self.module = copy.deepcopy(model)`, which guarantees that live and EMA share identical `named_parameters()` / -`named_buffers()` name sets AT CONSTRUCTION TIME. Matching by name at every -`update_parameters` call is robust to any of the reorderings above: buffers -and parameters are looked up by name in dictionaries and only same-name -entries are synced. +`named_buffers()` name trees AT CONSTRUCTION TIME. Matching by name at +every `update_parameters` call is robust to any of the reorderings above. + +To also cover shared buffers correctly, buffers are iterated with +`remove_duplicate=False`, so every registered path is walked and shared +tensors get synced under all their names. Redundant writes to the same +underlying tensor are cheap and correct. In the healthy case where positional and name-matched iteration would produce identical results, `NameMatchedAveragedModel` is bit-exact with the vanilla @@ -94,7 +108,8 @@ def update_parameters(self, model: Module) -> None: # type: ignore[override] if only_ema or only_live: raise RuntimeError( "NameMatchedAveragedModel: parameter name sets differ between " - f"live and EMA. only_ema={only_ema[:5]}..., only_live={only_live[:5]}...", + f"live and EMA. only_ema ({len(only_ema)}): {only_ema}. " + f"only_live ({len(only_live)}): {only_live}." ) common_params = sorted(ema_pd) @@ -135,17 +150,26 @@ def update_parameters(self, model: Module) -> None: # type: ignore[override] self.avg_fn(p_avg.detach(), p_mod, n_averaged), ) - # ---- Buffers: match by name --------------------------------------- + # ---- Buffers: match by name, INCLUDING shared duplicates ---------- if not self.use_buffers: - ema_bd = dict(self.module.named_buffers()) - live_bd = dict(model.named_buffers()) + # remove_duplicate=False so a buffer tensor shared by multiple + # parametrizations appears under EVERY name it's registered as. + # Otherwise dedup can pick different "first names" for shared + # tensors on live vs EMA when buffer references are re-pointed + # post-deepcopy (e.g. a sparsifier that consolidates per-parametrization + # mask buffers to a single project-owned tensor), producing + # spurious name-set mismatches AND, in a silent-skip variant, + # missing the sync entirely. + ema_bd = dict(self.module.named_buffers(remove_duplicate=False)) + live_bd = dict(model.named_buffers(remove_duplicate=False)) only_ema_b = sorted(set(ema_bd) - set(live_bd)) only_live_b = sorted(set(live_bd) - set(ema_bd)) if only_ema_b or only_live_b: + # Full list to help debug — silent truncation obscures the real diff. raise RuntimeError( "NameMatchedAveragedModel: buffer name sets differ between " - f"live and EMA. only_ema={only_ema_b[:5]}..., " - f"only_live={only_live_b[:5]}...", + f"live and EMA. only_ema ({len(only_ema_b)}): {only_ema_b}. " + f"only_live ({len(only_live_b)}): {only_live_b}." ) for name in sorted(ema_bd): b_swa = ema_bd[name] From 3b868b069062df76cd762c88a516eb0d5b4e3a5b Mon Sep 17 00:00:00 2001 From: bkmi user Date: Wed, 22 Jul 2026 15:40:33 +0000 Subject: [PATCH 3/4] ema: satisfy ruff lint (unused import, TCH002) 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. --- fairchem/src/fairchem/core/common/ema.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/fairchem/src/fairchem/core/common/ema.py b/fairchem/src/fairchem/core/common/ema.py index 11f3325744..1607e8d798 100644 --- a/fairchem/src/fairchem/core/common/ema.py +++ b/fairchem/src/fairchem/core/common/ema.py @@ -78,15 +78,17 @@ silent skipping would let the EMA slowly diverge from the live model in ways that are very hard to detect downstream. Prefer to fail loudly. """ + from __future__ import annotations -from typing import Optional +from typing import TYPE_CHECKING, Optional -import torch -from torch import Tensor -from torch.nn import Module from torch.optim import swa_utils +if TYPE_CHECKING: + from torch import Tensor + from torch.nn import Module + class NameMatchedAveragedModel(swa_utils.AveragedModel): """`AveragedModel` variant that syncs parameters/buffers by NAME. @@ -136,12 +138,16 @@ def update_parameters(self, model: Module) -> None: # type: ignore[override] ) in grouped.items(): if self.multi_avg_fn: self.multi_avg_fn( - self_params, model_params, self.n_averaged.to(device), + self_params, + model_params, + self.n_averaged.to(device), ) else: multi_avg_fn = swa_utils.get_swa_multi_avg_fn() multi_avg_fn( - self_params, model_params, self.n_averaged.to(device), + self_params, + model_params, + self.n_averaged.to(device), ) else: for p_avg, p_mod in zip(self_param_detached, model_param_detached): From 9f326ad1848e9e28f24484c46c618c9933147041 Mon Sep 17 00:00:00 2001 From: bkmi user Date: Wed, 22 Jul 2026 16:19:03 +0000 Subject: [PATCH 4/4] ema: relocate to correct source root MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- {fairchem/src => src}/fairchem/core/common/ema.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename {fairchem/src => src}/fairchem/core/common/ema.py (100%) diff --git a/fairchem/src/fairchem/core/common/ema.py b/src/fairchem/core/common/ema.py similarity index 100% rename from fairchem/src/fairchem/core/common/ema.py rename to src/fairchem/core/common/ema.py