Add support for Paged Optimizers (Adam, Adamw), 8-bit optimizers, and new optimizers: LARS, LAMB and LION

[arnavgarg1]

This PR adds support for a variety of new optimizers, including paged and 8-bit variants. All of these are useful for fine-tuning. I will add better descriptions for parameters in a follow-up PR when I better understand some of the underlying papers for LAMB, LARS and LION. I will also add tests once #3578 PR lands.

What are paged optimizers?

Paged Optimizers use NVIDIA unified memory 3 feature which does automatic page-to-page transfers between the CPU and GPU for error-free GPU processing in the scenario where the GPU occasionally runs out-of-memory. The feature works like regular memory paging between CPU RAM and the disk. This allocates paged memory for the optimizer states which are then automatically evicted to CPU RAM when the GPU runs out-of-memory and paged back into GPU memory when the memory is needed in the optimizer update step.

What new optimizers are being added to Ludwig?

Here's a summary of the new optimizers that will be added with this PR:

• SGD: 8-bit support
• Adam: 8-bit, Paged Adam, Paged Adam 8-bit
• AdamW: 8-bit, Paged AdamW, Paged AdamW 8-bit
• Adagrad: 8-bit
• RMSProp: 8-bit
• LAMB, Lamb 8-bit
• LARS, LARS 8-bit
• LION, LION 8-bit, Paged LION, Paged LION 8-bit

Here's an example of how you can set these different variants:

Regular AdamW

trainer:
  optimizer:
    type: adamw

8-Bit AdamW

trainer:
  optimizer:
    type: adamw_8bit

Paged AdamW

trainer:
  optimizer:
    type: paged_adamw

Paged AdamW 8-bit

trainer:
  optimizer:
    type: paged_adamw_8bit

All of this has been made possible through a deeper integration with the bitsandbytes library directly. Note that the 8-bit and paged optimizers only work on GPU machines, and they are not compatible with Deepspeed.

[github-actions]

Unit Test Results

       6 files  ±0         6 suites  ±0   1h 40m 56s ⏱️ + 16m 20s
2 830 tests +3  2 794 ✔️ +4  12 💤 ±0  24 ❌  - 1 
2 873 runs  +3  2 828 ✔️ +4  21 💤 ±0  24 ❌  - 1 

For more details on these failures, see this check.

Results for commit 0158d03. ± Comparison against base commit 3d2ff0b.

♻️ This comment has been updated with latest results.