When gradient checkpointing is enabled, flash_attn_kwargs cannot be passed into the decoder_layer

[efsotr]

System Info

transformers 4.47.1

Who can help?

@ArthurZucker

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

https://github.com/huggingface/transformers/blob/241c04d36867259cdf11dbb4e9d9a60f9cb65ebc/src/transformers/models/llama/modeling_llama.py#L896C1-L931C54

for decoder_layer in self.layers[: self.config.num_hidden_layers]:
    if output_hidden_states:
        all_hidden_states += (hidden_states,)


    if self.gradient_checkpointing and self.training:
        layer_outputs = self._gradient_checkpointing_func(
            decoder_layer.__call__,
            hidden_states,
            causal_mask,
            position_ids,
            past_key_values,
            output_attentions,
            use_cache,
            cache_position,
            position_embeddings,
        )
    else:
        layer_outputs = decoder_layer(
            hidden_states,
            attention_mask=causal_mask,
            position_ids=position_ids,
            past_key_value=past_key_values,
            output_attentions=output_attentions,
            use_cache=use_cache,
            cache_position=cache_position,
            position_embeddings=position_embeddings,
            **flash_attn_kwargs,
        )

Expected behavior

x

[ArthurZucker]

Hey! Yep, this is a known issue and it is rather a limitation from the way we implement gradient checkpointing AFAIK.
You can try this branch: #34987 with GradientCheckpointLayer should be working fine! Tell me if that helps?