Add ForwardContext to wrap model forward pass

src/transformers/adapters/configuration.py

@@ -207,8 +207,6 @@ def __init__(self, **kwargs):
         # TODO-V2 Save this with config?
         self.active_setup: Optional[AdapterCompositionBlock] = None
         self.skip_layers = None
-        # TODO This flag will be set & reset in every forward pass. Check if there is a better solution without state mutation.
-        self.is_parallelized = False
 
     def __contains__(self, item):
         return item in self.adapters.keys()

src/transformers/adapters/context.py

@@ -1,3 +1,4 @@
+import functools
 import threading
 
 from .composition import parse_composition, parse_heads_from_composition
@@ -41,9 +42,9 @@ def __exit__(self, type, value, traceback):
 
     @classmethod
     def get_contexts(cls):
-        if not hasattr(cls.storage, "setups"):
-            cls.storage.setups = []
-        return cls.storage.setups
+        if not hasattr(cls.storage, "contexts"):
+            cls.storage.contexts = []
+        return cls.storage.contexts
 
     @classmethod
     def get_context(cls):
@@ -58,3 +59,49 @@ def get_context_head_setup(cls):
         if context:
             return context.head_setup
         return None
+
+
+class ForwardContext:
+    """
+    Holds context information during a forward pass through a model. This class should be used via the
+    ``ForwardContext.wrap()`` method.
+
+    Note that the context is thread-local.
+    """
+
+    # thread-local storage that holds a stack of active contexts
+    storage = threading.local()
+
+    def __init__(self, model):
+        # If the model has a method ``forward_context()``, use it to create the context.
+        if hasattr(model, "forward_context"):
+            model.forward_context(self)
+
+    @classmethod
+    def wrap(cls, f):
+        """
+        Decorator method that wraps a ``forward()`` function of a model class.
+        """
+
+        @functools.wraps(f)
+        def wrapper_func(self, *args, **kwargs):
+            context = cls()
+            cls.get_contexts().append(context)
+            results = f(self, *args, **kwargs)
+            cls.get_contexts().pop()
+            return results
+
+        return wrapper_func
+
+    @classmethod
+    def get_contexts(cls):
+        if not hasattr(cls.storage, "contexts"):
+            cls.storage.contexts = []
+        return cls.storage.contexts
+
+    @classmethod
+    def get_context(cls):
+        try:
+            return cls.get_contexts()[-1]
+        except IndexError:
+            return None

src/transformers/adapters/layer.py

@@ -4,7 +4,7 @@
 from torch import nn
 
 from .composition import AdapterCompositionBlock, BatchSplit, Fuse, Parallel, Split, Stack
-from .context import AdapterSetup
+from .context import AdapterSetup, ForwardContext
 from .modeling import Adapter, BertFusion
 
 
@@ -317,11 +317,12 @@ def adapter_parallel(self, adapter_setup: Parallel, hidden_states, input_tensor,
         # We assume that all adapters have the same config
         adapter_config = self.config.adapters.get(adapter_setup.first())
 
-        if not self.config.adapters.is_parallelized:
+        context = ForwardContext.get_context()
+        if not context.adapters_parallelized:
             orig_batch_size = input_tensor.shape[0]
             input_tensor = input_tensor.repeat(self.config.adapters.active_setup.parallel_channels, 1, 1)
             hidden_states = hidden_states.repeat(self.config.adapters.active_setup.parallel_channels, 1, 1)
-            self.config.adapters.is_parallelized = True
+            context.adapters_parallelized = True
         else:
             # The base model should handle replication of input.
             # Therefore, we assume the (replicated) input batch to be divisible by the number of parallel channels.

src/transformers/adapters/model_mixin.py

@@ -12,7 +12,7 @@
 from ..models.auto.tokenization_auto import AutoTokenizer
 from .composition import AdapterCompositionBlock, Fuse, Stack, parse_composition
 from .configuration import AdapterConfig, AdapterFusionConfig, ModelAdaptersConfig, get_adapter_config_hash
-from .context import AdapterSetup
+from .context import AdapterSetup, ForwardContext
 from .hub_mixin import PushAdapterToHubMixin
 from .layer import AdapterLayer
 from .loading import AdapterFusionLoader, AdapterLoader, PredictionHeadLoader, WeightsLoader
@@ -563,16 +563,16 @@ def freeze_model(self, freeze=True):
             param.requires_grad = not freeze
         self.model_freezed = freeze
 
-    def pre_transformer_forward(self):
+    def forward_context(self, context: ForwardContext):
         """
-        This method should be called by every adapter-implementing model at the very beginning of the forward() method.
+        This method is called by the ``ForwardContext`` at the beginning of the forward pass.
         """
         # some warnings if we don't use available adapters
         active_adapters = self.active_adapters or AdapterSetup.get_context()
         if not active_adapters and self.has_adapters():
             logger.warning("There are adapters available but none are activated for the forward pass.")
 
-        self.config.adapters.is_parallelized = False
+        context.adapters_parallelized = False
 
     def load_embeddings(self, path: str, name: str):
         """

src/transformers/models/bart/modeling_bart.py

@@ -26,6 +26,7 @@
 
 from ...activations import ACT2FN
 from ...adapters.composition import adjust_tensors_for_parallel
+from ...adapters.context import ForwardContext
 from ...adapters.model_mixin import InvertibleAdaptersMixin, ModelWithHeadsAdaptersMixin
 from ...adapters.models.bart import (
     BartDecoderLayerAdaptersMixin,
@@ -1163,6 +1164,7 @@ def get_decoder(self):
         output_type=Seq2SeqModelOutput,
         config_class=_CONFIG_FOR_DOC,
     )
+    @ForwardContext.wrap
     def forward(
         self,
         input_ids=None,
@@ -1195,7 +1197,6 @@ def forward(
         )
         use_cache = use_cache if use_cache is not None else self.config.use_cache
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        self.pre_transformer_forward()
 
         if encoder_outputs is None:
             encoder_outputs = self.encoder(
@@ -1741,8 +1742,8 @@ def __init__(self, config):
 
         self._init_adapter_modules()
 
+    @ForwardContext.wrap
     def forward(self, *args, **kwargs):
-        self.pre_transformer_forward()
 
         return self.decoder(*args, **kwargs)
 

src/transformers/models/bert/modeling_bert.py

@@ -30,7 +30,7 @@
 
 from ...activations import ACT2FN
 from ...adapters.composition import adjust_tensors_for_parallel
-from ...adapters.context import AdapterSetup
+from ...adapters.context import AdapterSetup, ForwardContext
 from ...adapters.model_mixin import ModelWithHeadsAdaptersMixin
 from ...adapters.models.bert import (
     BertModelAdaptersMixin,
@@ -912,6 +912,7 @@ class PreTrainedModel
         output_type=BaseModelOutputWithPoolingAndCrossAttentions,
         config_class=_CONFIG_FOR_DOC,
     )
+    @ForwardContext.wrap
     def forward(
         self,
         input_ids=None,
@@ -953,7 +954,6 @@ def forward(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        self.pre_transformer_forward()
 
         if self.config.is_decoder:
             use_cache = use_cache if use_cache is not None else self.config.use_cache

src/transformers/models/distilbert/modeling_distilbert.py

@@ -28,6 +28,7 @@
 
 from ...activations import gelu
 from ...adapters.composition import adjust_tensors_for_parallel
+from ...adapters.context import ForwardContext
 from ...adapters.model_mixin import ModelWithHeadsAdaptersMixin
 from ...adapters.models.distilbert import (
     DistilBertModelAdaptersMixin,
@@ -527,6 +528,7 @@ class PreTrainedModel
         output_type=BaseModelOutput,
         config_class=_CONFIG_FOR_DOC,
     )
+    @ForwardContext.wrap
     def forward(
         self,
         input_ids=None,
@@ -542,7 +544,6 @@ def forward(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        self.pre_transformer_forward()
 
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")

src/transformers/models/encoder_decoder/modeling_encoder_decoder.py

@@ -20,6 +20,7 @@
 import torch
 from torch.nn import CrossEntropyLoss
 
+from ...adapters.context import ForwardContext
 from ...adapters.models.encoder_decoder import EncoderDecoderModelAdaptersMixin
 from ...configuration_utils import PretrainedConfig
 from ...file_utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
@@ -408,6 +409,7 @@ def from_encoder_decoder_pretrained(
 
     @add_start_docstrings_to_model_forward(ENCODER_DECODER_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
+    @ForwardContext.wrap
     def forward(
         self,
         input_ids=None,
@@ -462,9 +464,6 @@ def forward(
             argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
         }
 
-        if self.config.adapters:
-            self.pre_transformer_forward()
-
         if encoder_outputs is None:
             encoder_outputs = self.encoder(
                 input_ids=input_ids,

src/transformers/models/gpt2/modeling_gpt2.py

@@ -35,6 +35,7 @@
 
 from ...activations import ACT2FN
 from ...adapters.composition import adjust_tensors_for_parallel
+from ...adapters.context import ForwardContext
 from ...adapters.model_mixin import ModelWithHeadsAdaptersMixin
 from ...adapters.models.gpt2 import GPT2DecoderBlockAdaptersMixin, GPT2ModelAdapterMixin, GPT2ModelHeadsMixin
 from ...file_utils import (
@@ -742,6 +743,7 @@ def _prune_heads(self, heads_to_prune):
         output_type=BaseModelOutputWithPastAndCrossAttentions,
         config_class=_CONFIG_FOR_DOC,
     )
+    @ForwardContext.wrap
     def forward(
         self,
         input_ids=None,
@@ -758,8 +760,6 @@ def forward(
         output_hidden_states=None,
         return_dict=None,
     ):
-        self.pre_transformer_forward()
-
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states

src/transformers/models/mbart/modeling_mbart.py

@@ -25,6 +25,7 @@
 
 from ...activations import ACT2FN
 from ...adapters.composition import adjust_tensors_for_parallel
+from ...adapters.context import ForwardContext
 from ...adapters.model_mixin import InvertibleAdaptersMixin, ModelWithHeadsAdaptersMixin
 from ...adapters.models.bart import (
     BartDecoderLayerAdaptersMixin,
@@ -1166,6 +1167,7 @@ def get_decoder(self):
         output_type=Seq2SeqModelOutput,
         config_class=_CONFIG_FOR_DOC,
     )
+    @ForwardContext.wrap
     def forward(
         self,
         input_ids=None,
@@ -1190,7 +1192,6 @@ def forward(
         )
         use_cache = use_cache if use_cache is not None else self.config.use_cache
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        self.pre_transformer_forward()
 
         # different to other models, MBart automatically creates decoder_input_ids from
         # input_ids if no decoder_input_ids are provided
@@ -1748,8 +1749,8 @@ def __init__(self, config):
 
         self._init_adapter_modules()
 
+    @ForwardContext.wrap
     def forward(self, *args, **kwargs):
-        self.pre_transformer_forward()
 
         return self.decoder(*args, **kwargs)
 

src/transformers/models/roberta/modeling_roberta.py

@@ -25,7 +25,7 @@
 
 from ...activations import ACT2FN, gelu
 from ...adapters.composition import adjust_tensors_for_parallel
-from ...adapters.context import AdapterSetup
+from ...adapters.context import AdapterSetup, ForwardContext
 from ...adapters.model_mixin import ModelWithHeadsAdaptersMixin
 from ...adapters.models.bert import (
     BertModelAdaptersMixin,
@@ -764,6 +764,7 @@ class PreTrainedModel
         output_type=BaseModelOutputWithPoolingAndCrossAttentions,
         config_class=_CONFIG_FOR_DOC,
     )
+    @ForwardContext.wrap
     # Copied from transformers.models.bert.modeling_bert.BertModel.forward
     def forward(
         self,
@@ -806,7 +807,6 @@ def forward(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        self.pre_transformer_forward()
 
         if self.config.is_decoder:
             use_cache = use_cache if use_cache is not None else self.config.use_cache

src/transformers/models/t5/modeling_t5.py

@@ -27,6 +27,7 @@
 
 from ...activations import ACT2FN
 from ...adapters.composition import adjust_tensors_for_parallel
+from ...adapters.context import ForwardContext
 from ...adapters.model_mixin import InvertibleAdaptersMixin, ModelWithHeadsAdaptersMixin
 from ...adapters.models.t5 import (
     T5CrossAttentionLayerAdaptersMixin,
@@ -1354,6 +1355,7 @@ class PreTrainedModel
 
     @add_start_docstrings_to_model_forward(T5_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=Seq2SeqModelOutput, config_class=_CONFIG_FOR_DOC)
+    @ForwardContext.wrap
     def forward(
         self,
         input_ids=None,
@@ -1391,7 +1393,6 @@ def forward(
         """
         use_cache = use_cache if use_cache is not None else self.config.use_cache
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        self.pre_transformer_forward()
 
         # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask
         if head_mask is not None and decoder_head_mask is None:
@@ -1554,6 +1555,7 @@ def get_decoder(self):
 
     @add_start_docstrings_to_model_forward(T5_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
+    @ForwardContext.wrap
     def forward(
         self,
         input_ids=None,
@@ -1603,7 +1605,6 @@ def forward(
         """
         use_cache = use_cache if use_cache is not None else self.config.use_cache
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        self.pre_transformer_forward()
 
         # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask
         if head_mask is not None and decoder_head_mask is None: