From 2ab76c2c4f1d095a69aa04b8feef8833a2fa03b0 Mon Sep 17 00:00:00 2001 From: Yun Dai Date: Thu, 5 Oct 2023 01:18:27 -0700 Subject: [PATCH] [Falcon] Set `use_cache=False` before creating `presents` which relies on `use_cache` (#26328) * Set `presents=None` when `use_cache` is set to False for activation ckpt * Update modeling_falcon.py * fix black --- src/transformers/models/falcon/modeling_falcon.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/transformers/models/falcon/modeling_falcon.py b/src/transformers/models/falcon/modeling_falcon.py index 5c3de182b2a..db8546ffa59 100644 --- a/src/transformers/models/falcon/modeling_falcon.py +++ b/src/transformers/models/falcon/modeling_falcon.py @@ -1094,6 +1094,12 @@ class FalconModel(FalconPreTrainedModel): hidden_states = inputs_embeds + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False presents = () if use_cache else None all_self_attentions = () if output_attentions else None all_hidden_states = () if output_hidden_states else None @@ -1137,11 +1143,6 @@ class FalconModel(FalconPreTrainedModel): all_hidden_states = all_hidden_states + (hidden_states,) if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False def create_custom_forward(module): def custom_forward(*inputs):