PytorchLARS.step() crashes with UnboundLocalError on its own default momentum=0
PytorchLARS.step() (bitsandbytes/optim/lars.py) only assigns update inside the if momentum != 0: block:
if momentum != 0:
buf = state.get("momentum_buffer", None)
if buf is None:
buf = torch.clone(d_p).detach()
state["momentum_buffer"] = buf
else:
buf.mul_(momentum).add_(d_p, alpha=1 - dampening)
if nesterov:
update = d_p + buf * momentum
else:
update = buf
update_scale = 1.0
if max_unorm > 0.0:
...
p.add_(update, alpha=-lr * update_scale)There's no else branch defining update for momentum == 0, and momentum defaults to 0 in __init__. So the class crashes on its own default construction:
import torch
import bitsandbytes as bnb
p = torch.nn.Parameter(torch.randn(8))
opt = bnb.optim.PytorchLARS([p]) # momentum=0, the default
p.grad = torch.randn_like(p)
opt.step()UnboundLocalError: cannot access local variable 'update' where it is not associated with a valueRan on main at 833649043474794b8fe7a4136e0c40faf077b2e0, torch 2.14.0, Python 3.13.12, CPU (no CUDA/MPS involved — this is pure Python control flow, backend-independent).
PytorchLARS has no test coverage in tests/test_optim.py and no other references anywhere in the codebase, which is presumably why this hasn't surfaced — the bitsandbytes-native LARS/LARS8bit/LARS32bit classes in the same file explicitly guard against this (if momentum == 0: raise NotImplementedError("LARS without momentum is not supported!") in each __init__), but PytorchLARS.__init__ has no such check.
Fix: either add the same momentum == 0 guard to PytorchLARS.__init__ that the other three classes already have, or add an else: update = d_p branch (plain SGD update, matching the semantics of torch.optim.SGD with momentum disabled) if momentum-free operation is meant to be supported.
Source: bitsandbytes-foundation/bitsandbytes