raise error on FlashAttnKwargs + grad ckpt

Author: garrett361

src/transformers/models/bamba/modeling_bamba.py

@@ -1302,6 +1302,11 @@ def forward(
                 all_hidden_states += (hidden_states,)
 
             if self.gradient_checkpointing and self.training:
+                if "cu_seq_lens_q" in "flash_attn_kwargs":
+                    raise NotImplementedError(
+                        "Padding-free training with FlashAttentionKwargs and gradient checkpointing"
+                        " not currently supported."
+                    )
                 layer_outputs = self._gradient_checkpointing_func(
                     decoder_layer.__call__,
                     hidden_states,

src/transformers/models/bamba/modular_bamba.py

@@ -1050,6 +1050,11 @@ def forward(
                 all_hidden_states += (hidden_states,)
 
             if self.gradient_checkpointing and self.training:
+                if "cu_seq_lens_q" in "flash_attn_kwargs":
+                    raise NotImplementedError(
+                        "Padding-free training with FlashAttentionKwargs and gradient checkpointing"
+                        " not currently supported."
+                    )
                 layer_outputs = self._gradient_checkpointing_func(
                     decoder_layer.__call__,
                     hidden_states,