From d6bd0e4fb4c28909049bc45d116754e5b7382977 Mon Sep 17 00:00:00 2001 From: Zhu Lei Date: Tue, 15 Sep 2026 08:16:01 +0000 Subject: [PATCH] Fix LookSAM gv state key mismatch and float alpha handling compute_gv looked up 'old_grad' while sam_perturb_weights stores 'old_g', so gv was never written and non-k steps skipped apply_gv. Also accept Python-float alpha in apply_gv instead of calling .to(). --- src/infty/optim/geometry_reshaping/looksam.py | 8 ++- tests/optim/test_looksam_gv.py | 66 +++++++++++++++++++ 2 files changed, 71 insertions(+), 3 deletions(-) create mode 100644 tests/optim/test_looksam_gv.py diff --git a/src/infty/optim/geometry_reshaping/looksam.py b/src/infty/optim/geometry_reshaping/looksam.py index 9718c12..b602397 100644 --- a/src/infty/optim/geometry_reshaping/looksam.py +++ b/src/infty/optim/geometry_reshaping/looksam.py @@ -38,9 +38,10 @@ def sam_perturb_weights(self): def compute_gv(self): for group in self.param_groups: for p in group['params']: - if p.grad is None or 'old_grad' not in self.state[p]: + # sam_perturb_weights stores the first-order grad under 'old_g' + if p.grad is None or 'old_g' not in self.state[p]: continue - old_grad = self.state[p]['old_grad'] + old_grad = self.state[p]['old_g'] new_grad = p.grad g_grad_norm = old_grad / (old_grad.norm(p=2) + 1e-8) g_new_grad_norm = new_grad / (new_grad.norm(p=2) + 1e-8) @@ -55,7 +56,8 @@ def apply_gv(self): if p.grad is None or 'gv' not in self.state[p]: continue gv = self.state[p]['gv'] - p.grad.add_(self.alpha.to(p) * (p.grad.norm(p=2) / (gv.norm(p=2) + 1e-8)) * gv) + alpha = self.alpha.to(p) if torch.is_tensor(self.alpha) else self.alpha + p.grad.add_(alpha * (p.grad.norm(p=2) / (gv.norm(p=2) + 1e-8)) * gv) def step(self, closure=None, delay=False): if closure: diff --git a/tests/optim/test_looksam_gv.py b/tests/optim/test_looksam_gv.py new file mode 100644 index 0000000..186a6b8 --- /dev/null +++ b/tests/optim/test_looksam_gv.py @@ -0,0 +1,66 @@ +import torch + + +def test_looksam_stores_gv_after_periodic_sam_step(): + """LookSAM must persist a vertical component (`gv`) after each k-step SAM refresh. + + Regression: compute_gv looked up state['old_grad'] while sam_perturb_weights + stores state['old_g'], so gv was never written and non-k steps silently + skipped apply_gv (LookSAM degraded to plain SGD between refreshes). + """ + from infty.optim.geometry_reshaping.looksam import LookSAM + + torch.manual_seed(0) + model = torch.nn.Linear(4, 2) + base_optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + optimizer = LookSAM( + model.parameters(), + base_optimizer=base_optimizer, + model=model, + args={"rho": 0.05, "k": 2, "alpha": 0.7}, + ) + + x = torch.randn(6, 4) + y = torch.tensor([0, 1, 0, 1, 0, 1]) + criterion = torch.nn.CrossEntropyLoss() + + def closure(): + logits = model(x) + loss = criterion(logits, y) + return logits, [loss] + + optimizer.set_closure(closure) + + optimizer.step() # step 1: non-k, no gv yet + optimizer.step() # step 2: k-multiple -> SAM + compute_gv + + assert any( + "gv" in optimizer.state.get(p, {}) for p in model.parameters() + ), "LookSAM.compute_gv did not store 'gv' after a periodic SAM step" + + +def test_looksam_apply_gv_accepts_float_alpha(): + """apply_gv must accept Python-float alpha (default 0.7), not call alpha.to(p).""" + from infty.optim.geometry_reshaping.looksam import LookSAM + + torch.manual_seed(1) + model = torch.nn.Linear(4, 2) + base_optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + optimizer = LookSAM( + model.parameters(), + base_optimizer=base_optimizer, + model=model, + args={"rho": 0.05, "k": 2, "alpha": 0.7}, + ) + assert isinstance(optimizer.alpha, float) + + for p in model.parameters(): + p.grad = torch.ones_like(p) + optimizer.state[p]["gv"] = torch.randn_like(p) + + # Must not raise AttributeError: 'float' object has no attribute 'to' + optimizer.apply_gv() + + for p in model.parameters(): + assert p.grad is not None + assert torch.isfinite(p.grad).all()