[Feature] Add BAdam algorithm

[Ledzy]

What does this PR do?

This PR incorporates the BAdam algorithm to the repository. One can now use BAdam by setting argument "--use_badam". It enables the full parameter finetuning of Llama 2-7b within 24GB RAM under mixed precision training, while using only half time compared with LoRA.

• Paper link: https://arxiv.org/abs/2404.02827.
• Code repo: https://github.com/Ledzy/BAdam.

A sample running script is available at "examples/extras/badam/sft.sh". The script attains 0.8591 eval_loss after 900 steps within training time of 3 hours (including evaluation time) on a single RTX-3090. Note that one needs to run pip install badam before running the script.

Notes

• The gradient_checkpointing_enable in "llmtuner/model/patcher.py" is modified to check if any parameter of the checkpointed layer is trainable. If so, it sets the input's requires_grad to be True. The modification enables checkpointed layer to be trainable when the input's requires_grad is False. Such a modification is essential for the acceleration of BAdam (otherwise, the backward will always go to the first layer, due to model.enable_input_require_grads(), which reduces the training time by half.
• BAdam uses mixed precision training, where the fp32 master weight are created during the creation and update of the optimizer; the model should be loaded in fp16 format. The file "llmtuner/model/adapter.py" is modified accordingly.
• Currently, BAdam's implementation doesn't support distributed training. The DDP and ZeRO implementation is scheduled to be support in future versions of package badam.

Before submitting

• Did you read the contributor guideline?

[hiyouga]

thanks for adding this brilliant algorithm to LLaMA Factory, please take a look at my comments, especially the implementation of gradient checkpointing.

[hiyouga]

Some necessary changes are made, it can be merged now

[leekum2018]

How to select Badam in the web interface, instead of the command line?