bnb optimizers could use bnb.nn.StableEmbedding instead of torch.nn.Embedding

[mtasic85]

According to bnb documentation here:

https://huggingface.co/docs/bitsandbytes/main/optimizers
https://huggingface.co/docs/bitsandbytes/main/explanations/optimizers#stable-embedding-layer

This line could alter between bnb.nn.StableEmbedding and torch.nn.Embedding, or allow it to be configurable in config file:

litgpt/litgpt/model.py

Line 28 in a8aa4ba

wte=nn.Embedding(config.padded_vocab_size, config.n_embd),

There are also other places in code where torch.nn.Embedding is used.

[rasbt]

Thanks for the note and good point, I didn't know about this.

One challenge I see with configuring it in the config file is that it's used to model creation. But one can later optionally run with --quantize bnb.nf4 or not. So, ideally, that swap should only take place upon calling the inference/training functions and leave the original model as is.

[rasbt]

Upon reading a bit more, this would only be required for training (due to the optimizer choice). I added it in #1770