Add support for gradient checkpointing for LLM fine-tuning

[arnavgarg1]

This PR adds support in the finetune trainer to optionally enabled gradient_checkpointing.

What is gradient checkpointing?

Gradient checkpointing works by recomputing the activations of the model during the backward pass, rather than storing them in memory during the forward pass. This is a tradeoff between compute and memory, as the activations need to be recomputed during the backward pass, but the memory footprint is reduced. This is set to false by default because it is not always beneficial to use gradient checkpointing, and it can sometimes slow down training.

How can you use gradient checkpointing in the config?

Gradient checkpointing is disabled by default. To enable it, you can simply set enable_gradient_checkpointing to True.

trainer:
    enable_gradient_checkpointing: true

When should I enable gradient checkpointing?

This is useful when training very large models that run into out of memory errors very quickly during training. It is particularly helpful when doing non-quantization based training (adapter based or full fine-tuning).

[github-actions]

Unit Test Results

  6 files  ±0    6 suites  ±0   39m 8s ⏱️ +45s
31 tests ±0  26 ✔️ ±0    5 💤 ±0  0 ❌ ±0 
82 runs  ±0  66 ✔️ ±0  16 💤 ±0  0 ❌ ±0 

Results for commit 1b3fe04. ± Comparison against base commit 6198693.

[justinxzhao]

Nice description! Do you have a sense for which LLMs support gradient checkpointing and which ones don't?

[arnavgarg1]

Nice description! Do you have a sense for which LLMs support gradient checkpointing and which ones don't?

I think almost all models coming from the transformer package support it, but I have the additional check here just in case there is an edge case that hasn't been handled.