"...git@developer.sourcefind.cn:chenpangpang/transformers.git" did not exist on "8594dd80dd49b388fe2ba29c68e609a8d3ace9f6"
Unverified Commit 4edfd2d4 authored by saswatmeher's avatar saswatmeher Committed by GitHub
Browse files

Fix Gradient checkpointing bug BigBird (#21882)


Co-authored-by: default avatarsaswatmeher <saswatmeher@cse.iitb.ac.in>
parent 269b0549
...@@ -1592,6 +1592,13 @@ class BigBirdEncoder(nn.Module): ...@@ -1592,6 +1592,13 @@ class BigBirdEncoder(nn.Module):
all_self_attentions = () if output_attentions else None all_self_attentions = () if output_attentions else None
all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
if self.gradient_checkpointing and self.training:
if use_cache:
logger.warning_once(
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
)
use_cache = False
next_decoder_cache = () if use_cache else None next_decoder_cache = () if use_cache else None
for i, layer_module in enumerate(self.layer): for i, layer_module in enumerate(self.layer):
...@@ -1602,11 +1609,6 @@ class BigBirdEncoder(nn.Module): ...@@ -1602,11 +1609,6 @@ class BigBirdEncoder(nn.Module):
past_key_value = past_key_values[i] if past_key_values is not None else None past_key_value = past_key_values[i] if past_key_values is not None else None
if self.gradient_checkpointing and self.training: if self.gradient_checkpointing and self.training:
if use_cache:
logger.warning_once(
"`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