From 777f4d3795bf3e2e48772beb7e4e28c112661caa Mon Sep 17 00:00:00 2001 From: "Florian (Feuermagier)" Date: Mon, 24 Feb 2025 15:51:43 +0100 Subject: [PATCH 1/6] sample params when closure is passed to step --- ivon/_ivon.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ivon/_ivon.py b/ivon/_ivon.py index 5dfb6cf..6e39ef4 100644 --- a/ivon/_ivon.py +++ b/ivon/_ivon.py @@ -174,7 +174,8 @@ def step(self, closure: ClosureType = None) -> Optional[Tensor]: losses = [] for _ in range(self.mc_samples): with torch.enable_grad(): - loss = closure() + with self.sampled_params(): + loss = closure() losses.append(loss) loss = sum(losses) / self.mc_samples if self.sync and dist.is_initialized(): # explicit sync From dc7c18d37f5204bd386c43af8656db808dab3614 Mon Sep 17 00:00:00 2001 From: "Florian (Feuermagier)" Date: Mon, 24 Feb 2025 17:04:10 +0100 Subject: [PATCH 2/6] move to device on demand --- ivon/_ivon.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/ivon/_ivon.py b/ivon/_ivon.py index 6e39ef4..ab3ba81 100644 --- a/ivon/_ivon.py +++ b/ivon/_ivon.py @@ -198,8 +198,17 @@ def _sample_params(self) -> Tuple[Tensor, Tensor]: offset = 0 for group in self.param_groups: gnumel = group["numel"] + + for p in group["params"]: + if p is None: + continue + group_device = p.device + break + + group["hess"] = group["hess"].to(group_device) + noise_sample = ( - torch.randn(gnumel, device=self._device, dtype=self._dtype) + torch.randn(gnumel, device=group_device, dtype=self._dtype) / ( group["ess"] * (group["hess"] + group["weight_decay"]) ).sqrt() From 50761ad73b863b9f0bc9f72f65302e71307341f5 Mon Sep 17 00:00:00 2001 From: "Florian (Feuermagier)" Date: Mon, 24 Feb 2025 17:11:44 +0100 Subject: [PATCH 3/6] fix --- ivon/_ivon.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ivon/_ivon.py b/ivon/_ivon.py index ab3ba81..2ee0446 100644 --- a/ivon/_ivon.py +++ b/ivon/_ivon.py @@ -174,7 +174,7 @@ def step(self, closure: ClosureType = None) -> Optional[Tensor]: losses = [] for _ in range(self.mc_samples): with torch.enable_grad(): - with self.sampled_params(): + with self.sampled_params(train=True): loss = closure() losses.append(loss) loss = sum(losses) / self.mc_samples From 347f02933129f201c75a38ce6b7e1b2f20c961ab Mon Sep 17 00:00:00 2001 From: "Florian (Feuermagier)" Date: Mon, 24 Feb 2025 17:21:33 +0100 Subject: [PATCH 4/6] move momentum to device --- ivon/_ivon.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/ivon/_ivon.py b/ivon/_ivon.py index 2ee0446..efdabad 100644 --- a/ivon/_ivon.py +++ b/ivon/_ivon.py @@ -238,6 +238,12 @@ def _update(self): offset = 0 for group in self.param_groups: + for p in group["params"]: + if p is None: + continue + group_device = p.device + break + lr = group["lr"] b1 = group["beta1"] b2 = group["beta2"] @@ -247,6 +253,8 @@ def _update(self): [p.flatten() for p in group["params"] if p is not None], 0 ) + group["momentum"] = group["momentum"].to(group_device) + group["momentum"] = self._new_momentum( self.state["avg_grad"][pg_slice], group["momentum"], b1 ) From 72e46654fe75d0a64fddd30c2bdfc417bab21ffd Mon Sep 17 00:00:00 2001 From: "Florian (Feuermagier)" Date: Sun, 2 Mar 2025 22:51:35 +0100 Subject: [PATCH 5/6] always restore param average --- ivon/_ivon.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/ivon/_ivon.py b/ivon/_ivon.py index efdabad..84207ad 100644 --- a/ivon/_ivon.py +++ b/ivon/_ivon.py @@ -128,8 +128,10 @@ def _init_buffers(self): @contextmanager def sampled_params(self, train: bool = False): param_avg, noise = self._sample_params() - yield - self._restore_param_average(train, param_avg, noise) + try: + yield + finally: + self._restore_param_average(train, param_avg, noise) def _restore_param_average( self, train: bool, param_avg: Tensor, noise: Tensor From 3257031c2f7732f5332c7eac0af83225cceefd2b Mon Sep 17 00:00:00 2001 From: Florian Seligmann Date: Thu, 10 Apr 2025 12:46:21 +0900 Subject: [PATCH 6/6] proper device handling --- ivon/_ivon.py | 55 ++++++++++++++++++++++++++++++--------------------- 1 file changed, 33 insertions(+), 22 deletions(-) diff --git a/ivon/_ivon.py b/ivon/_ivon.py index 84207ad..4397e9e 100644 --- a/ivon/_ivon.py +++ b/ivon/_ivon.py @@ -74,7 +74,7 @@ def __init__( self.mc_samples = mc_samples self.hess_approx = hess_approx self.sync = sync - self._numel, self._device, self._dtype = self._get_param_configs() + self._numel, self._dtype = self._get_param_configs() self.current_step = 0 self.debias = debias self.rescale_lr = rescale_lr @@ -90,14 +90,8 @@ def _get_param_configs(self): pg["numel"] = sum(p.numel() for p in pg["params"] if p is not None) all_params += [p for p in pg["params"] if p is not None] if len(all_params) == 0: - return 0, torch.device("cpu"), torch.get_default_dtype() - devices = {p.device for p in all_params} - if len(devices) > 1: - raise ValueError( - "Parameters are on different devices: " - f"{[str(d) for d in devices]}" - ) - device = next(iter(devices)) + return 0, torch.get_default_dtype() + dtypes = {p.dtype for p in all_params} if len(dtypes) > 1: raise ValueError( @@ -106,7 +100,7 @@ def _get_param_configs(self): ) dtype = next(iter(dtypes)) total = sum(pg["numel"] for pg in self.param_groups) - return total, device, dtype + return total, dtype def _reset_samples(self): self.state['count'] = 0 @@ -118,11 +112,13 @@ def _init_buffers(self): for group in self.param_groups: hess_init, numel = group["hess_init"], group["numel"] + group_device = IVON._find_group_device(group, check_same_device=True) + group["momentum"] = torch.zeros( - numel, device=self._device, dtype=self._dtype + numel, device=group_device, dtype=self._dtype ) group["hess"] = torch.zeros( - numel, device=self._device, dtype=self._dtype + numel, device=group_device, dtype=self._dtype ).add(torch.as_tensor(hess_init)) @contextmanager @@ -201,11 +197,7 @@ def _sample_params(self) -> Tuple[Tensor, Tensor]: for group in self.param_groups: gnumel = group["numel"] - for p in group["params"]: - if p is None: - continue - group_device = p.device - break + group_device = IVON._find_group_device(group, check_same_device=False) group["hess"] = group["hess"].to(group_device) @@ -240,11 +232,7 @@ def _update(self): offset = 0 for group in self.param_groups: - for p in group["params"]: - if p is None: - continue - group_device = p.device - break + group_device = IVON._find_group_device(group, check_same_device=True) lr = group["lr"] b1 = group["beta1"] @@ -327,3 +315,26 @@ def _new_param_averages( min=-clip_radius, max=clip_radius, ) + + @staticmethod + def _find_group_device(group, check_same_device=True): + group_device = None + for p in group["params"]: + if p is None: + continue + + if group_device is None: + group_device = p.device + elif group_device != p.device: + raise ValueError( + "Parameters are on different devices: " + f"{group_device} and {p.device}" + ) + + if not check_same_device: + return group_device + + if group_device is None: + return torch.device("cpu") + + return group_device