From 6a5ed9f3b3f66ed1d43eae57326c735340290c71 Mon Sep 17 00:00:00 2001 From: Sajal Kumar Jana Date: Mon, 2 Nov 2026 03:39:43 +0000 Subject: [PATCH 1/2] fix(aggregation): Register tensor attributes as buffers --- CHANGELOG.md | 5 +++++ src/torchjd/aggregation/_aligned_mtl.py | 1 + src/torchjd/aggregation/_config.py | 1 + src/torchjd/aggregation/_constant.py | 9 +++++---- src/torchjd/aggregation/_cr_mogm.py | 4 +++- src/torchjd/aggregation/_dualproj.py | 1 + src/torchjd/aggregation/_sdmgrad.py | 1 + src/torchjd/aggregation/_upgrad.py | 1 + tests/unit/aggregation/test_aligned_mtl.py | 14 ++++++++++++++ tests/unit/aggregation/test_config.py | 16 +++++++++++++++- tests/unit/aggregation/test_constant.py | 8 ++++++++ tests/unit/aggregation/test_cr_mogm.py | 10 ++++++++++ tests/unit/aggregation/test_dualproj.py | 16 +++++++++++++++- tests/unit/aggregation/test_sdmgrad.py | 9 +++++++++ tests/unit/aggregation/test_upgrad.py | 16 +++++++++++++++- 15 files changed, 104 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d04177a36..4a0c81315 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,11 @@ changelog does not include internal changes that do not affect the user. - Fixed `Krum` and `KrumWeighting` ranking rows by sums of Euclidean distances instead of squared Euclidean distances, as required by Krum and Multi-Krum. This could select a different gradient or a different set of gradients to average. +- Fixed `.to()` not moving the tensors given to `Constant`, `ConstantWeighting`, `UPGrad`, + `UPGradWeighting`, `DualProj`, `DualProjWeighting`, `AlignedMTL`, `AlignedMTLWeighting`, `ConFIG`, + `SDMGradWeighting` and `CRMOGMWeighting` (`weights`, `pref_vector` and `initial_weights`). These + tensors are now registered as non-persistent buffers, so they follow the dtype and device of the + aggregator or weighting without being added to its `state_dict`. ## [0.18.0] - 2026-10-08 diff --git a/src/torchjd/aggregation/_aligned_mtl.py b/src/torchjd/aggregation/_aligned_mtl.py index 7df1558cc..5ad48ee17 100644 --- a/src/torchjd/aggregation/_aligned_mtl.py +++ b/src/torchjd/aggregation/_aligned_mtl.py @@ -36,6 +36,7 @@ def __init__( scale_mode: SUPPORTED_SCALE_MODE = "min", ) -> None: super().__init__() + self.register_buffer("_pref_vector", None, persistent=False) self.pref_vector = pref_vector self.scale_mode: SUPPORTED_SCALE_MODE = scale_mode diff --git a/src/torchjd/aggregation/_config.py b/src/torchjd/aggregation/_config.py index 8485a83d4..b81689a2f 100644 --- a/src/torchjd/aggregation/_config.py +++ b/src/torchjd/aggregation/_config.py @@ -31,6 +31,7 @@ class ConFIG(Aggregator, _NonDifferentiable): def __init__(self, pref_vector: Tensor | None = None) -> None: super().__init__() + self.register_buffer("_pref_vector", None, persistent=False) self.pref_vector = pref_vector def forward(self, matrix: Matrix, /) -> Tensor: diff --git a/src/torchjd/aggregation/_constant.py b/src/torchjd/aggregation/_constant.py index c849155e9..4dc7ed8e0 100644 --- a/src/torchjd/aggregation/_constant.py +++ b/src/torchjd/aggregation/_constant.py @@ -14,6 +14,8 @@ class ConstantWeighting(_MatrixWeighting): :param weights: The weights to return at each call. """ + weights: Tensor + def __init__(self, weights: Tensor) -> None: if weights.dim() != 1: raise ValueError( @@ -22,7 +24,7 @@ def __init__(self, weights: Tensor) -> None: ) super().__init__() - self.weights = weights + self.register_buffer("weights", weights, persistent=False) def forward(self, matrix: Tensor, /) -> Tensor: self._check_matrix_shape(matrix) @@ -48,11 +50,10 @@ class Constant(WeightedAggregator): def __init__(self, weights: Tensor) -> None: super().__init__(weighting=ConstantWeighting(weights=weights)) - self._weights = weights def __repr__(self) -> str: - return f"{self.__class__.__name__}(weights={repr(self._weights)})" + return f"{self.__class__.__name__}(weights={repr(self.weighting.weights)})" def __str__(self) -> str: - weights_str = vector_to_str(self._weights) + weights_str = vector_to_str(self.weighting.weights) return f"{self.__class__.__name__}([{weights_str}])" diff --git a/src/torchjd/aggregation/_cr_mogm.py b/src/torchjd/aggregation/_cr_mogm.py index c1b7cadbb..54b92d69e 100644 --- a/src/torchjd/aggregation/_cr_mogm.py +++ b/src/torchjd/aggregation/_cr_mogm.py @@ -98,13 +98,15 @@ class CRMOGMWeighting(Weighting[_T], Stateful): cr_mogm.alpha = 1 - current_lr / initial_lr """ + _initial_weights: Tensor | None + def __init__( self, weighting: Weighting[_T], alpha: float = 0.9, initial_weights: Tensor | None = None ) -> None: super().__init__() self.weighting = weighting self.alpha = alpha - self._initial_weights = initial_weights + self.register_buffer("_initial_weights", initial_weights, persistent=False) self.register_buffer("_lambda", None) @property diff --git a/src/torchjd/aggregation/_dualproj.py b/src/torchjd/aggregation/_dualproj.py index a9a4172ac..1e4b4dbf8 100644 --- a/src/torchjd/aggregation/_dualproj.py +++ b/src/torchjd/aggregation/_dualproj.py @@ -28,6 +28,7 @@ def __init__( projector: DualConeProjector | None = None, ) -> None: super().__init__() + self.register_buffer("_pref_vector", None, persistent=False) self.pref_vector = pref_vector self.projector = projector diff --git a/src/torchjd/aggregation/_sdmgrad.py b/src/torchjd/aggregation/_sdmgrad.py index 91de1d794..88dbec4bf 100644 --- a/src/torchjd/aggregation/_sdmgrad.py +++ b/src/torchjd/aggregation/_sdmgrad.py @@ -107,6 +107,7 @@ def __init__( self.momentum = momentum self.n_iter = n_iter self.lambda_ = lambda_ + self.register_buffer("_pref_vector", None, persistent=False) self.pref_vector = pref_vector self.register_buffer("_w", None) self._state_key: tuple[int, torch.dtype, torch.device] | None = None diff --git a/src/torchjd/aggregation/_upgrad.py b/src/torchjd/aggregation/_upgrad.py index 2e12e84d7..be585affb 100644 --- a/src/torchjd/aggregation/_upgrad.py +++ b/src/torchjd/aggregation/_upgrad.py @@ -29,6 +29,7 @@ def __init__( projector: DualConeProjector | None = None, ) -> None: super().__init__() + self.register_buffer("_pref_vector", None, persistent=False) self.pref_vector = pref_vector self.projector = projector diff --git a/tests/unit/aggregation/test_aligned_mtl.py b/tests/unit/aggregation/test_aligned_mtl.py index 3608a3eab..3fe5e24bb 100644 --- a/tests/unit/aggregation/test_aligned_mtl.py +++ b/tests/unit/aggregation/test_aligned_mtl.py @@ -66,3 +66,17 @@ def test_float64_small_eigenvalue_is_kept() -> None: J = tensor_([[1.0, 0.0], [0.0, 1e-4]], dtype=torch.float64) result = AlignedMTL()(J) assert_close(result, tensor_([5e-5, 5e-5], dtype=torch.float64)) + + +def test_to_moves_pref_vector() -> None: + A = AlignedMTL(pref_vector=tensor_([1.0, 2.0], dtype=torch.float32)).to(torch.float64) + weighting = A.gramian_weighting.weighting + assert isinstance(weighting, ConstantWeighting) + + assert A.pref_vector is not None + assert A.pref_vector.dtype == torch.float64 + assert weighting.weights.dtype == torch.float64 + buffers = list(A.buffers()) + assert any(buffer is A.pref_vector for buffer in buffers) + assert any(buffer is weighting.weights for buffer in buffers) + assert len(A.state_dict()) == 0 diff --git a/tests/unit/aggregation/test_config.py b/tests/unit/aggregation/test_config.py index fe9516404..f2f26f635 100644 --- a/tests/unit/aggregation/test_config.py +++ b/tests/unit/aggregation/test_config.py @@ -1,7 +1,7 @@ import torch from pytest import mark from torch import Tensor -from utils.tensors import ones_ +from utils.tensors import ones_, tensor_ from torchjd.aggregation import ConFIG, ConstantWeighting @@ -56,3 +56,17 @@ def test_pref_vector_setter_updates_value() -> None: assert A.pref_vector is new_pref assert isinstance(A.weighting, ConstantWeighting) assert A.weighting.weights is new_pref + + +def test_to_moves_pref_vector() -> None: + A = ConFIG(pref_vector=tensor_([1.0, 2.0], dtype=torch.float32)).to(torch.float64) + weighting = A.weighting + assert isinstance(weighting, ConstantWeighting) + + assert A.pref_vector is not None + assert A.pref_vector.dtype == torch.float64 + assert weighting.weights.dtype == torch.float64 + buffers = list(A.buffers()) + assert any(buffer is A.pref_vector for buffer in buffers) + assert any(buffer is weighting.weights for buffer in buffers) + assert len(A.state_dict()) == 0 diff --git a/tests/unit/aggregation/test_constant.py b/tests/unit/aggregation/test_constant.py index aa1332fcb..ba3af392f 100644 --- a/tests/unit/aggregation/test_constant.py +++ b/tests/unit/aggregation/test_constant.py @@ -90,3 +90,11 @@ def test_representations() -> None: A = Constant(weights=torch.tensor([1.0, 2.0], device="cpu")) assert repr(A) == "Constant(weights=tensor([1., 2.]))" assert str(A) == "Constant([1., 2.])" + + +def test_to_moves_weights() -> None: + A = Constant(tensor_([1.0, 2.0], dtype=torch.float32)).to(torch.float64) + + assert A.weighting.weights.dtype == torch.float64 + assert any(buffer is A.weighting.weights for buffer in A.buffers()) + assert len(A.state_dict()) == 0 diff --git a/tests/unit/aggregation/test_cr_mogm.py b/tests/unit/aggregation/test_cr_mogm.py index b01945b38..329561181 100644 --- a/tests/unit/aggregation/test_cr_mogm.py +++ b/tests/unit/aggregation/test_cr_mogm.py @@ -1,3 +1,4 @@ +import torch from pytest import mark, raises from torch import Tensor from torch.testing import assert_close @@ -201,3 +202,12 @@ def test_zero_columns() -> None: aggregator = WeightedAggregator(CRMOGMWeighting(MeanWeighting())) out = aggregator(tensor_([]).reshape(2, 0)) assert out.shape == (0,) + + +def test_to_moves_initial_weights() -> None: + W = CRMOGMWeighting(MeanWeighting(), initial_weights=tensor_([0.5, 0.5], dtype=torch.float32)) + W = W.to(torch.float64) + + buffers = dict(W.named_buffers()) + assert buffers["_initial_weights"].dtype == torch.float64 + assert "_initial_weights" not in W.state_dict() diff --git a/tests/unit/aggregation/test_dualproj.py b/tests/unit/aggregation/test_dualproj.py index 0f11f2d8b..34a058dbd 100644 --- a/tests/unit/aggregation/test_dualproj.py +++ b/tests/unit/aggregation/test_dualproj.py @@ -8,7 +8,7 @@ import torch from pytest import mark from torch import Tensor -from utils.tensors import ones_ +from utils.tensors import ones_, tensor_ from ._asserts import ( assert_expected_structure, @@ -89,3 +89,17 @@ def test_projector_setter_updates_value() -> None: A.projector = new_projector assert A.projector is new_projector assert A.gramian_weighting.projector is new_projector + + +def test_to_moves_pref_vector() -> None: + A = DualProj(pref_vector=tensor_([1.0, 2.0], dtype=torch.float32)).to(torch.float64) + weighting = A.gramian_weighting.weighting + assert isinstance(weighting, ConstantWeighting) + + assert A.pref_vector is not None + assert A.pref_vector.dtype == torch.float64 + assert weighting.weights.dtype == torch.float64 + buffers = list(A.buffers()) + assert any(buffer is A.pref_vector for buffer in buffers) + assert any(buffer is weighting.weights for buffer in buffers) + assert len(A.state_dict()) == 0 diff --git a/tests/unit/aggregation/test_sdmgrad.py b/tests/unit/aggregation/test_sdmgrad.py index e846b4554..1f7ee4bc5 100644 --- a/tests/unit/aggregation/test_sdmgrad.py +++ b/tests/unit/aggregation/test_sdmgrad.py @@ -170,3 +170,12 @@ def test_non_differentiable() -> None: W = SDMGradWeighting() weights = W(A) assert not weights.requires_grad + + +def test_to_moves_pref_vector() -> None: + W = SDMGradWeighting(pref_vector=tensor_([0.5, 0.5], dtype=torch.float32)).to(torch.float64) + + assert W.pref_vector is not None + assert W.pref_vector.dtype == torch.float64 + assert "_pref_vector" in dict(W.named_buffers()) + assert "_pref_vector" not in W.state_dict() diff --git a/tests/unit/aggregation/test_upgrad.py b/tests/unit/aggregation/test_upgrad.py index baa5a361d..ad9fa57dd 100644 --- a/tests/unit/aggregation/test_upgrad.py +++ b/tests/unit/aggregation/test_upgrad.py @@ -8,7 +8,7 @@ import torch from pytest import mark from torch import Tensor -from utils.tensors import ones_ +from utils.tensors import ones_, tensor_ from ._asserts import ( assert_expected_structure, @@ -95,3 +95,17 @@ def test_projector_setter_updates_value() -> None: A.projector = new_projector assert A.projector is new_projector assert A.gramian_weighting.projector is new_projector + + +def test_to_moves_pref_vector() -> None: + A = UPGrad(pref_vector=tensor_([1.0, 2.0], dtype=torch.float32)).to(torch.float64) + weighting = A.gramian_weighting.weighting + assert isinstance(weighting, ConstantWeighting) + + assert A.pref_vector is not None + assert A.pref_vector.dtype == torch.float64 + assert weighting.weights.dtype == torch.float64 + buffers = list(A.buffers()) + assert any(buffer is A.pref_vector for buffer in buffers) + assert any(buffer is weighting.weights for buffer in buffers) + assert len(A.state_dict()) == 0 From 695894d764376601c5e2c23639b9a2d477484e18 Mon Sep 17 00:00:00 2001 From: Sajal Kumar Jana Date: Thu, 5 Nov 2026 05:52:27 +0000 Subject: [PATCH 2/2] fix(aggregation): Register GradDrop leak as buffer --- CHANGELOG.md | 7 ++++--- src/torchjd/aggregation/_graddrop.py | 1 + tests/unit/aggregation/test_graddrop.py | 11 ++++++++++- 3 files changed, 15 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ca22a6fa..ba3bf197b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,9 +19,10 @@ changelog does not include internal changes that do not affect the user. gradient or a different set of gradients to average. - Fixed `.to()` not moving the tensors given to `Constant`, `ConstantWeighting`, `UPGrad`, `UPGradWeighting`, `DualProj`, `DualProjWeighting`, `AlignedMTL`, `AlignedMTLWeighting`, `ConFIG`, - `SDMGradWeighting` and `CRMOGMWeighting` (`weights`, `pref_vector` and `initial_weights`). These - tensors are now registered as non-persistent buffers, so they follow the dtype and device of the - aggregator or weighting without being added to its `state_dict`. + `SDMGradWeighting`, `CRMOGMWeighting` and `GradDrop` (`weights`, `pref_vector`, + `initial_weights` and `leak`). These tensors are now registered as non-persistent buffers, so + they follow the dtype and device of the aggregator or weighting without being added to its + `state_dict`. ## [0.18.0] - 2026-10-08 diff --git a/src/torchjd/aggregation/_graddrop.py b/src/torchjd/aggregation/_graddrop.py index 38e952327..3271284b5 100644 --- a/src/torchjd/aggregation/_graddrop.py +++ b/src/torchjd/aggregation/_graddrop.py @@ -31,6 +31,7 @@ class GradDrop(Aggregator, _NonDifferentiable): def __init__(self, f: Callable = _identity, leak: Tensor | None = None) -> None: super().__init__() self.f = f + self.register_buffer("_leak", None, persistent=False) self.leak = leak def forward(self, matrix: Matrix, /) -> Tensor: diff --git a/tests/unit/aggregation/test_graddrop.py b/tests/unit/aggregation/test_graddrop.py index d21ef690f..0afdf82a2 100644 --- a/tests/unit/aggregation/test_graddrop.py +++ b/tests/unit/aggregation/test_graddrop.py @@ -5,7 +5,7 @@ from pytest import mark, raises from torch import Tensor from utils.contexts import ExceptionContext -from utils.tensors import ones_ +from utils.tensors import ones_, tensor_ from torchjd.aggregation import GradDrop @@ -102,3 +102,12 @@ def test_leak_setter_rejects_non_1d() -> None: A = GradDrop() with raises(ValueError, match="leak"): A.leak = torch.tensor([[0.0, 1.0], [1.0, 0.0]]) + + +def test_to_moves_leak() -> None: + A = GradDrop(leak=tensor_([0.0, 1.0], dtype=torch.float32)).to(torch.float64) + + assert A.leak is not None + assert A.leak.dtype == torch.float64 + assert any(buffer is A.leak for buffer in A.buffers()) + assert len(A.state_dict()) == 0