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.

[calpt]

Thanks a lot for digging into this and enabling compatibility w adapters, this looks great! Just a couple of small comments before we're good to merge

[calpt]

to clarify: is this fixing something specific to adapter training or just generally for the model with gradient checkpointing? If the latter, I think we shouldn't copy the class here but try to fix upstream in Transformers directly (same for Deberta-v1)

[lenglaender]

It's not specific to our library. It is a bug from HF Transformers when one calls model.enable_input_require_grads() before training. Input grads are AFAIK required when training adapters, at least HF advices to do this. This is from the doc of the enable_input_require_grads function:

Enables the gradients for the input embeddings. This is useful for fine-tuning adapter weights while keeping the model weights fixed.

I will open a PR in the HF Transformers repo this week to fix this bug for Deberta(V2). Since it's an HF Transformers bug, it doesn't work with PEFT currently either; the following script throws the same error, we encountered:

from transformers import DebertaConfig, DebertaForSequenceClassification
from peft import get_peft_model, LoraConfig, TaskType
import torch

# Create a minimal DeBERTa config
config = DebertaConfig(
    hidden_size=32,
    num_hidden_layers=5,
    num_attention_heads=4,
    intermediate_size=37,
    hidden_act="gelu",
    relative_attention=True,
    pos_att_type="p2c|c2p"
)

# PEFT model
model = DebertaForSequenceClassification(config)
peft_config = LoraConfig(task_type=TaskType.SEQ_CLS, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1)
model = get_peft_model(model, peft_config)

# Enable input gradients
model.enable_input_require_grads()

model.train()

# Move to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# Create a simple input tensor (random; doesn't matter) & forward pass
batch_size = 2
seq_length = 10
input_ids = torch.ones((batch_size, seq_length), dtype=torch.long).to(device)
attention_mask = torch.ones_like(input_ids).to(device)

outputs = model(input_ids=input_ids, attention_mask=attention_mask)

[calpt]

do we need to test both True and False here given we have other train tests for standard training?

[lenglaender]

We train it once with and once without gradient checkpointing to compare the state dicts afterwards to ensure that gradient checkpointing leads to the same weight updates. I will add a short comment to explain this.

[calpt]

Can we add a small note from where this method is overriden?

[lenglaender]

Yes, will do this 👍

[lenglaender]

Hey @calpt, can you quickly review my replies to your comments and approve the PR if everything is alright so I can merge this? We need to merge this PR first, then add all new tests to test refactoring #740, and then we can merge #763 (because Julian has already merged the test refactoring into his PR). So, currently, this PR is blocking everything we want to have in the next release.