Lion 8 bit

[lucidrains]

per advice from Tim, the plan will be to closely follow the 1-state logic of RMSProp and conditionally branch out for the lion logic

status - Lion8bit is successfully training a small autoregressive transformer for character level enwik8 on my machine

some remaining todos before merging

in pytorch

def update_fn(p, grad, exp_avg, lr, wd, beta1, beta2):
    # stepweight decay

    p.data.mul_(1 - lr * wd)

    # weight update

    update = exp_avg.clone().mul_(beta1).add(grad, alpha = 1 - beta1).sign_()
    p.add_(update, alpha = -lr)

    # decay the momentum running average coefficient

    exp_avg.mul_(beta2).add_(grad, alpha = 1 - beta2)

to test

$ git clone https://github.com/lucidrains/bitsandbytes
$ cd bitsandbytes
# follow the make from source instructions above

• add tests
• add some comments
• figure out why tests break and make it pass
• consider just copy pasting lion into the repo

[lucidrains]

autoregressive enwik8 with lion8bit converged successfully (sans weight decay), even though there's some inaccuracies in kernel.cu

[lucidrains]

it may not be hitting the 8bit1state blockwise path though, where state is being smoothed before the update

[lucidrains]

weight decay looks good

[christallire]

resolves #150

[lucidrains]

@christallire welcome you to build from source and import the Lion8bit and give it a try

seems to be working great, thanks to Tim's surrounding scaffold

[lucidrains]

ok, will let this PR sit for a while to gather feedback. probably will revisit it at end of this month and see if we can get it merged

[lucidrains]

i think there may still be an issue with the blockwise version, where the momentum update with beta2 is occuring before the actual parameter updates

[TimDettmers]

Thank you for your work on this! This looks almost good to me. The only thing that is off is the order of the updates that you point out. We can get around this without rewriting the CUDA code by using the gradient variable as temporary variable to store the current momentum value + the update value for Lion (see other comments). Please add a comment to these lines to help people understand the use of the gradient as a temporary variable (adding more comments is appreciated, sorry for the mess with the undocumented code).

The enwiki test runs are great! I think this shows that it works. It would be great to add a test in test_optim.py for Lion. You need three things for that: (1) define lion str2optimizer dictionary entry, (2) define str2statenames dictionary entry, (3) add the optimizer name to the optimizer_names list for the test_optimzier8bit function. That should be all that is needed. The only question is which 32-bit baseline to use. You might want to use your own Lion 32-bit repo. I think a fair test could also be to compare against the 32-bit bnb optimizer (since you have shown that 8-bit already replicates enwiki performance). This might be simpler and quick to write up.

[lucidrains]

@TimDettmers thank you Tim for the feedback! 🙏 will address everything you brought up and poke you when it is ready later this week

[TimDettmers]

Thank you, this looks good to me. The issue with the test is expected. The error from Lion is expected to be higher at times due to its noisy update. I will merge and will have a look at the test. This is an excellent PR, thank you for all the work. I think it will be invaluable to the community!

[lucidrains]

@TimDettmers oh hey Tim! glad to see this merged

thank you for reviewing it and getting it out there! totally forgot about it, my bad