PytorchLARS.step() crashes with UnboundLocalError on its own default momentum=0

[caiotheodoro]

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 value

Ran 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.