Skip to content

GitLab

Unverified Commit 7811bf7e authored by Karim Foda's avatar Karim Foda Committed by GitHub
Browse files

Fix gradient checkpointing bug in gptneox (#21815) 

* Fix gradient checkpointing bug in gptneox

* Remove use_cache block
parent 0c7f93f5
Hide whitespace changes
Inline Side-by-side
Showing with 7 additions and 5 deletions
src/transformers/models/gpt_neox/modeling_gpt_neox.py
View file @ 7811bf7e
...@@ -500,6 +500,13 @@ class GPTNeoXModel(GPTNeoXPreTrainedModel): ...@@ -500,6 +500,13 @@ class GPTNeoXModel(GPTNeoXPreTrainedModel):
hidden_states = inputs_embeds 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 presents = () if use_cache else None
all_attentions = () if output_attentions else None all_attentions = () if output_attentions else None
all_hidden_states = () if output_hidden_states else None all_hidden_states = () if output_hidden_states else None
...@@ -508,11 +515,6 @@ class GPTNeoXModel(GPTNeoXPreTrainedModel): ...@@ -508,11 +515,6 @@ class GPTNeoXModel(GPTNeoXPreTrainedModel):
all_hidden_states = all_hidden_states + (hidden_states,) all_hidden_states = all_hidden_states + (hidden_states,)
if self.gradient_checkpointing and self.training: 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 create_custom_forward(module):
def custom_forward(*inputs): def custom_forward(*inputs):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment