Function one_step

tests/test_optim.py:703–709

(max_unorm)

    # tight and a loose value produced identical updates. Runs on CPU, which implements
    # the 32-bit LAMB kernel.
    def one_step(max_unorm):
        torch.manual_seed(0)
        p = torch.nn.Parameter(torch.randn(128, 128))
        opt = bnb.optim.LAMB([p], lr=1e-1, max_unorm=max_unorm)
        p.grad = torch.randn(128, 128) * 5.0
        opt.step()
        return p.detach().clone()

    tight = one_step(1e-4)  # aggressive clipping
    loose = one_step(10.0)  # effectively no clipping