Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix gradient checkpointing + fp16 autocast for most models #24247

Merged
merged 15 commits into from
Jun 21, 2023
Merged
6 changes: 5 additions & 1 deletion src/transformers/models/albert/modeling_albert.py
Original file line number Diff line number Diff line change
@@ -34,7 +34,11 @@
TokenClassifierOutput,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
from ...pytorch_utils import (
apply_chunking_to_forward,
find_pruneable_heads_and_indices,
prune_linear_layer,
)
from ...utils import (
ModelOutput,
add_code_sample_docstrings,
9 changes: 7 additions & 2 deletions src/transformers/models/align/modeling_align.py
Original file line number Diff line number Diff line change
@@ -30,7 +30,12 @@
BaseModelOutputWithPoolingAndNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
from ...pytorch_utils import (
apply_chunking_to_forward,
find_pruneable_heads_and_indices,
prune_linear_layer,
torch_custom_checkpointing,
)
from ...utils import (
ModelOutput,
add_start_docstrings,
@@ -1100,7 +1105,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(layer_module),
hidden_states,
attention_mask,
11 changes: 8 additions & 3 deletions src/transformers/models/altclip/modeling_altclip.py
Original file line number Diff line number Diff line change
@@ -30,7 +30,12 @@
BaseModelOutputWithPoolingAndProjection,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
from ...pytorch_utils import (
apply_chunking_to_forward,
find_pruneable_heads_and_indices,
prune_linear_layer,
torch_custom_checkpointing,
)
from ...utils import ModelOutput, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
from .configuration_altclip import AltCLIPConfig, AltCLIPTextConfig, AltCLIPVisionConfig

@@ -651,7 +656,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(layer_module),
hidden_states,
attention_mask,
@@ -965,7 +970,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(encoder_layer),
hidden_states,
attention_mask,
Original file line number Diff line number Diff line change
@@ -25,7 +25,7 @@
from ...activations import ACT2FN
from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, SequenceClassifierOutput
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer, torch_custom_checkpointing
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_audio_spectrogram_transformer import ASTConfig

@@ -343,7 +343,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(layer_module),
hidden_states,
layer_head_mask,
5 changes: 3 additions & 2 deletions src/transformers/models/autoformer/modeling_autoformer.py
Original file line number Diff line number Diff line change
@@ -34,6 +34,7 @@
Seq2SeqTSPredictionOutput,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import torch_custom_checkpointing
from ...time_series_utils import NegativeBinomialOutput, NormalOutput, StudentTOutput
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
from .configuration_autoformer import AutoformerConfig
@@ -1210,7 +1211,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(encoder_layer),
hidden_states,
attention_mask,
@@ -1428,7 +1429,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(decoder_layer),
hidden_states,
attention_mask,
5 changes: 3 additions & 2 deletions src/transformers/models/bart/modeling_bart.py
Original file line number Diff line number Diff line change
@@ -35,6 +35,7 @@
Seq2SeqSequenceClassifierOutput,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import torch_custom_checkpointing
from ...utils import (
add_code_sample_docstrings,
add_end_docstrings,
@@ -849,7 +850,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(encoder_layer),
hidden_states,
attention_mask,
@@ -1105,7 +1106,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(decoder_layer),
hidden_states,
attention_mask,
4 changes: 2 additions & 2 deletions src/transformers/models/beit/modeling_beit.py
Original file line number Diff line number Diff line change
@@ -34,7 +34,7 @@
SemanticSegmenterOutput,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import find_pruneable_heads_and_indices, meshgrid, prune_linear_layer
from ...pytorch_utils import find_pruneable_heads_and_indices, meshgrid, prune_linear_layer, torch_custom_checkpointing
from ...utils import (
add_code_sample_docstrings,
add_start_docstrings,
@@ -517,7 +517,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(layer_module),
hidden_states,
layer_head_mask,
9 changes: 7 additions & 2 deletions src/transformers/models/bert/modeling_bert.py
Original file line number Diff line number Diff line change
@@ -40,7 +40,12 @@
TokenClassifierOutput,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
from ...pytorch_utils import (
apply_chunking_to_forward,
find_pruneable_heads_and_indices,
prune_linear_layer,
torch_custom_checkpointing,
)
from ...utils import (
ModelOutput,
add_code_sample_docstrings,
@@ -598,7 +603,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(layer_module),
hidden_states,
attention_mask,
Original file line number Diff line number Diff line change
@@ -25,7 +25,12 @@
from ...activations import ACT2FN
from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
from ...pytorch_utils import (
apply_chunking_to_forward,
find_pruneable_heads_and_indices,
prune_linear_layer,
torch_custom_checkpointing,
)
from ...utils import (
add_code_sample_docstrings,
add_start_docstrings,
@@ -408,7 +413,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(layer_module),
hidden_states,
attention_mask,
4 changes: 2 additions & 2 deletions src/transformers/models/big_bird/modeling_big_bird.py
Original file line number Diff line number Diff line change
@@ -37,7 +37,7 @@
TokenClassifierOutput,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import apply_chunking_to_forward
from ...pytorch_utils import apply_chunking_to_forward, torch_custom_checkpointing
from ...utils import (
ModelOutput,
add_code_sample_docstrings,
@@ -1622,7 +1622,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(layer_module),
hidden_states,
attention_mask,
Original file line number Diff line number Diff line change
@@ -36,6 +36,7 @@
Seq2SeqSequenceClassifierOutput,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import torch_custom_checkpointing
from ...utils import (
add_code_sample_docstrings,
add_end_docstrings,
@@ -1945,7 +1946,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(encoder_layer),
hidden_states,
attention_mask,
@@ -2291,7 +2292,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(decoder_layer),
hidden_states,
attention_mask,
3 changes: 2 additions & 1 deletion src/transformers/models/biogpt/modeling_biogpt.py
Original file line number Diff line number Diff line change
@@ -32,6 +32,7 @@
TokenClassifierOutput,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import torch_custom_checkpointing
from ...utils import (
add_code_sample_docstrings,
add_start_docstrings,
@@ -594,7 +595,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(decoder_layer),
hidden_states,
attention_mask,
5 changes: 3 additions & 2 deletions src/transformers/models/blenderbot/modeling_blenderbot.py
Original file line number Diff line number Diff line change
@@ -36,6 +36,7 @@
Seq2SeqModelOutput,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import torch_custom_checkpointing
from ...utils import (
add_end_docstrings,
add_start_docstrings,
@@ -779,7 +780,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(encoder_layer),
hidden_states,
attention_mask,
@@ -1034,7 +1035,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(decoder_layer),
hidden_states,
attention_mask,
Original file line number Diff line number Diff line change
@@ -34,6 +34,7 @@
Seq2SeqModelOutput,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import torch_custom_checkpointing
from ...utils import (
add_end_docstrings,
add_start_docstrings,
@@ -777,7 +778,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(encoder_layer),
hidden_states,
attention_mask,
@@ -1031,7 +1032,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(decoder_layer),
hidden_states,
attention_mask,
3 changes: 2 additions & 1 deletion src/transformers/models/blip/modeling_blip.py
Original file line number Diff line number Diff line change
@@ -25,6 +25,7 @@
from ...activations import ACT2FN
from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import torch_custom_checkpointing
from ...utils import (
ModelOutput,
add_start_docstrings,
@@ -620,7 +621,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(encoder_layer),
hidden_states,
attention_mask,
3 changes: 2 additions & 1 deletion src/transformers/models/blip/modeling_blip_text.py
Original file line number Diff line number Diff line change
@@ -34,6 +34,7 @@
find_pruneable_heads_and_indices,
prune_linear_layer,
)
from ...pytorch_utils import torch_custom_checkpointing
from ...utils import logging
from .configuration_blip import BlipTextConfig

@@ -427,7 +428,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(layer_module),
hidden_states,
attention_mask,
11 changes: 8 additions & 3 deletions src/transformers/models/blip_2/modeling_blip_2.py
Original file line number Diff line number Diff line change
@@ -31,7 +31,12 @@
BaseModelOutputWithPoolingAndCrossAttentions,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
from ...pytorch_utils import (
apply_chunking_to_forward,
find_pruneable_heads_and_indices,
prune_linear_layer,
torch_custom_checkpointing,
)
from ...utils import (
ModelOutput,
add_start_docstrings,
@@ -492,7 +497,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(encoder_layer),
hidden_states,
attention_mask,
@@ -963,7 +968,7 @@ def custom_forward(*inputs):

return custom_forward

layer_outputs = torch.utils.checkpoint.checkpoint(
layer_outputs = torch_custom_checkpointing(
create_custom_forward(layer_module),
hidden_states,
attention_mask,
3 changes: 2 additions & 1 deletion src/transformers/models/bloom/modeling_bloom.py
Original file line number Diff line number Diff line change
@@ -33,6 +33,7 @@
TokenClassifierOutput,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import torch_custom_checkpointing
from ...utils import logging
from .configuration_bloom import BloomConfig

@@ -775,7 +776,7 @@ def custom_forward(*inputs):

return custom_forward

outputs = torch.utils.checkpoint.checkpoint(
outputs = torch_custom_checkpointing(
create_custom_forward(block),
hidden_states,
alibi,
Loading