Add bitsandbytes optimizer support

[aicrumb]

Uses bitsandbytes adam optimizer instead of torch, adds very simple gradient accumulation, finetuning only bias/layernorms (tested, works very well and is very fast) and allows for different precisions easier. (sorry for the absolute mass of commits, I was making them on github.com)

[jon-tow]

Hey @Crumb! Just getting back to this. I think it's best to separate the 3 new features into their own PRs to avoid unwieldy commit histories (if there's an issue in one feature we'd have to revert the other features as well).

Let's make this a PR focused on adding the bitsandbytes optimizer support. I believe there might some more thought that needs to go into adding grad accumulation; it's possible accelerate handles this in subtly different ways for deepspeed integration (which we rely heavily on) and it should be properly tested.

Hope that's alright with you 🙏

(Also, no worries about the many commits! We squash them before merging)

[LouisCastricato]

lets make the bitsandbytes pr

[jon-tow]

Details on latest changes

• We opt to not convert nn.Embedding layers into bnb.nn.StableEmbedding layers since trlx is essentially a fine-tuning library; according to Tim Dettmers:

  "StableEmbedding layer is only required if the model was pretrained with the StableEmbedding layer."

  See relevant discussion here.

  • Instead, we take the more general approach of forcing nn.Embedding layers into 32-bit precision following advice from Tim Dettmers:

    "for pretrained models, the best would be to use the 32-bit optimized embedding layer (bnb.nn.Embedding) and no layer norm if the pretrained model was not trained with a [StableEmbedding] layer norm."

    Again, relevant discussion here.

Reports

• In practice, we can see ~13% memory savings for GPT-J when trained with 8 unfrozen layers. Note that large memory savings occur only when a large portion of the model is unfrozen (num_layers_unfrozen is approx. some large fraction of total model layers).
  • There seems to be a slight penalty for using 8-bit optimizers but nothing significantly divergent - can possibly be overcome with improved hyper-parameter tuning.
  • Relevant wandb report: here

CC @LouisCastricato

[LouisCastricato]

Looks good to me, let's merge.

[maxreciprocate]

hey @jon-tow, your wandb report (https://wandb.ai/carperai/trlx/reports/trlx-Add-bitsandbytes-optimizer-support-133--VmlldzozMjU5OTQx) doesn't render for me for some reason (there are no data in charts as if the run set is empty), is it again some wandb's shenanigans or just me 😴

[LouisCastricato]

It was rendering yesterday. I think wandb is just being weird.

[jon-tow]

@reciprocated those were jon-tow shenanigans 😅 Yesterday I discovered deleting a run also deletes its data from reports (which makes sense in hindsight). I re-ran them from my personal wandb account - try this link: https://wandb.ai/jon-tow/trlx/reports/trlx-Add-bitsandbytes-optimizer-support-133--VmlldzozMjY1MzI1

[maxreciprocate]

@jon-tow that link gives 404 for me 😓 you have to share a "magic" link instead

[jon-tow]

@reciprocated This oughta do it: https://wandb.ai/jon-tow/trlx/reports/trlx-Add-bitsandbytes-optimizer-support-133--VmlldzozMjY1MzI1?accessToken=4t40dgib8rzs1qb6kqgsx8yuqhp5csfc2dsmxpz0qdwv2justj4yer44qwwtg7qp

[maxreciprocate]

@jon-tow Thanks 🙏