Skip to content

Fix LookSAM gv state key mismatch and float alpha handling - #19

Open
kabishou11 wants to merge 1 commit into
THUDM:mainfrom
kabishou11:fix/looksam-gv-state-key
Open

kabishou11 wants to merge 1 commit into
THUDM:mainfrom
kabishou11:fix/looksam-gv-state-key

Conversation

@kabishou11

Copy link
Copy Markdown

Summary

Fixes #18.

LookSAM's periodic vertical-component path was dead on main:

  1. sam_perturb_weights stores the first-order gradient under state[p]['old_g'], but compute_gv looked up old_grad, so gv was never written and non-k steps skipped apply_gv (silent degrade to plain SGD between refreshes).
  2. Once that path runs, default float alpha=0.7 hit self.alpha.to(p) and raised AttributeError.

Changes

  • Read old_g in compute_gv (with a short comment tying it to sam_perturb_weights).
  • Accept Python-float or Tensor alpha in apply_gv.
  • Add unit tests that fail on unpatched main and pass after the fix.

Test plan

  • pytest tests/optim/test_looksam_gv.py tests/optim/test_geometry_reshaping.py -v
  • CI green on this PR

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().

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

LookSAM never builds gv: state key mismatch (old_grad vs old_g) and float alpha.to crash

1 participant