WIP: Add Support for Gradient Checkpointing

[lenglaender]

Add Support for Gradient Checkpointing

This PR adds support for gradient checkpointing Gradient checkpointing is a technique that trades computation for memory by recomputing intermediate activations during the backward pass instead of storing them. This is particularly useful when training large models. Because we recompute values during the backpropagation, we need to preserve the original ForwardContext in this phase. I solved this by overwriting the gradient_checkpointing_enable function so that the checkpoint function receives the current ForwardContext as the backward pass context manager.

I added tests to check if it works for every model using 1) a single adapter or 2) parallel composition.

Not Yet Fixed Bugs

• Single adapter tests: For some models like BERT, the implementation doesn't work. We get an error saying one of the variables needed for gradient computation has been modified by an inplace operation. I couldn't find out where this happened. For some models, the single adapter tests run through (like for Llama, ViT, Whisper); for others not (like BERT, DistilBERT, GPT2).
  • E.g., this test works: pytest tests/test_llama.py::LlamaAdapterTest::test_prefix_tuning_gradient_checkpointing_single_adapter
  • and this does not work: pytest tests/test_bert.py::BertAdapterTest::test_prefix_tuning_gradient_checkpointing_single_adapter
• For Parallel composition, we get errors with the dimensions. I believe this is due to us changing the batch size. I have not yet had time to take a look at this.
  • E.g. this test doesn't work: pytest tests/test_llama.py::LlamaAdapterTest::test_lora_gradient_checkpointing_parallel_adapters