Skip to content

fix(optim): reject nesterov and dampening in SGD and LARS - #2085

Open
ConnorMoss02 wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
ConnorMoss02:fix/sgd-lars-reject-unsupported-nesterov-dampening
Open

ConnorMoss02 wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
ConnorMoss02:fix/sgd-lars-reject-unsupported-nesterov-dampening

Conversation

@ConnorMoss02

Copy link
Copy Markdown

Problem

SGD, SGD8bit, SGD32bit, LARS, LARS8bit and LARS32bit accept nesterov and dampening, but neither reaches any update path:

  • nesterov is never referenced after the constructor signature.
  • dampening is passed to the base optimizer as betas[1], which no momentum path reads.

So both are silently ignored. bnb.optim.SGD(..., nesterov=True) is bit-identical to plain momentum, which is a different update rule:

torch.optim.SGD  nesterov False vs True                   max|diff| = 0.666526
bnb.optim.SGD    nesterov False vs True                   max|diff| = 0.000000
bnb SGD(nesterov=True) vs torch SGD(nesterov=False)       max|diff| = 0.000000   <- identical
torch.optim.SGD  dampening 0 vs 0.5                       max|diff| = 0.848891
bnb.optim.SGD    dampening 0 vs 0.5                       max|diff| = 0.000000

Both are 0.000000 for all six classes. docs/source/reference/optim/sgd.mdx also stated that bitsandbytes supports Nesterov momentum.

Fix

Notes

  • No public API changes. The guards only reject values that were already silently ignored, and no code in the repo constructs these classes with them.
  • PytorchLARS is unchanged: it implements nesterov itself and already validates it (lars.py, "Nesterov momentum requires a momentum and zero dampening").
  • This rejects rather than implements. Implementing Nesterov would require changes across all four backends.

Tests (tests/test_optim.py)

  • test_sgd_lars_reject_nesterov, test_sgd_lars_reject_dampening — parametrized over all six classes; defaults still construct.
  • test_pytorch_lars_still_supports_nesterov — behavioral, guards the boundary of this change.

pre-commit run --all-files passes. 12 of the 13 new cases fail with the source reverted; the PytorchLARS case passes either way, as it should. On this machine tests/test_optim.py -k "momentum or lars or sgd" gives 54 passed / 9 failed against 41 passed / 9 failed on main: the same 9 are test_optimizer8bit, failing with NotImplementedError: The operator 'bitsandbytes::optimizer_update_8bit_blockwise' is not implemented because there is no CUDA device here.

Relates to #1261

All six SGD and LARS classes accepted `nesterov` and `dampening`, and the SGD
docs advertised Nesterov momentum, but neither argument reached any update
path. `nesterov` was never referenced after the signature and `dampening` was
passed as betas[1], which no momentum kernel reads, so both were silently
ignored.

Reject non-default values instead, mirroring the existing Adam8bit and
AdamW8bit guards, and correct the docs sentence. PytorchLARS implements both
and is unchanged.
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.

1 participant