Apply GradientCheckpointingLayer to the whole repo

[qubvel]

What does this PR do?

Apply GradientCheckpointingLayer to the remaining models in the repository.

Most of the PR is pretty much similar changes for all models:

1. Add import for GradientCheckpointingLayer
2. Inherit *Layer modules from GradientCheckpointingLayer
3. Remove if/else path for gradient checkpointing, keeping the else path only
   3a) Some changes were required to make sure all tensors with gradients were passed as positional arguments.

Additionally, GradientCheckpointingLayer was modified slightly. I added handling for use_cache and past_key_values within the layer to disable them in case gradient checkpointing is enabled.

We still have to keep some redundant code, though:

Case 1.

        if self.gradient_checkpointing and self.training and use_cache:
            logger.warning_once(
                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
            )
            use_cache = False

because, later, most of the models rely on the use_cache parameter as follows

            if use_cache:
                next_decoder_cache = layer_outputs[2 if output_attentions else 1]

In case it is handled only by GradientCheckpointingLayer and not modified in the outer module, it leads to an IndexError.

Case 2.

In some cases layer parameters order doesn't allow to handle past_key_values as kwargs, e.g. GPT2

            outputs = block(
                hidden_states,
                past_key_values if not (self.gradient_checkpointing and self.training) else None,
                cache_position,
                causal_mask,
                head_mask[i],
                encoder_hidden_states,  # as a positional argument for gradient checkpointing
                encoder_attention_mask=encoder_attention_mask,
                use_cache=use_cache,
                output_attentions=output_attentions,
                **kwargs,
            )

We have to pass all params up to encoder_hidden_states as positional args (tensors that require grads have to be passed that way), so past_key_values is also passed a positional argument and resolved manually.

Alternatively, we can refactor layers' params order, but that would be a breaking change. Not that many models and mostly the old ones.

Not supported models

Also, there are a couple of exceptions where GradientCheckpointingLayer does not work. I tried to fix it, but I didn't go too far and just kept it in its original state

• zamba / zamba2
• mllama

cc @ArthurZucker @Cyrilvallez

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qubvel]

update for GradientCheckpointingLayer

[Cyrilvallez]

Big big PR, and super welcome! 🚀🤗 Can we add a common test for gradient checkpointing though? I see we don't have one yet (only in trainer) - just instantiate a small model and run a single forward with gradient checkpointing and ensure that it runs correctly would be super nice

[Cyrilvallez]

Alright, did not see the test_training_gradient_checkpointing... before, my bad! All good then! Let's merge! 🤗

[ArthurZucker]

SUper nice! Thanks