Lion Optimizer With Triton Kernel

[lapp0]

Feature request

lion-pytorch implements a triton kernel https://github.com/lucidrains/lion-pytorch/blob/master/lion_pytorch/triton.py

It's implemented pretty simply in Lion.__init__, might be easy to integrate?

        if use_triton:
            from lion_pytorch.triton import update_fn as triton_update_fn
            self.update_fn = triton_update_fn

Motivation

Improve optimizer.step() performance.

Your contribution

With guidance on whether or not there is some complexity I'm missing, I'd be able to test and integrate this.

[lapp0]

Monkeypatched and benchmarked performance. Performance increase is below the margin of error.