Adding adapter support for NeoX

src/transformers/__init__.py

@@ -2596,6 +2596,8 @@
  "GPT2AdapterModel",
  "GPT2ModelWithHeads",
  "GPTJAdapterModel",
+ "GPTNeoXAdapterModel",
+ "GPTNeoXModelWithHeads",
  "HoulsbyConfig",
  "HoulsbyInvConfig",
  "IA3Config",
@@ -2942,6 +2944,15 @@
  "TFGPTJPreTrainedModel",
  ]
  )
+ _import_structure["models.gpt_neox"].extend(
+ [
+ "TFGPTNeoXForCausalLM",
+ "TFGPTNeoXForQuestionAnswering",
+ "TFGPTNeoXForSequenceClassification",
+ "TFGPTNeoXModel",
+ "TFGPTNeoXPreTrainedModel",
+ ]
+ )
  _import_structure["models.groupvit"].extend(
  [
  "TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST",
@@ -3456,6 +3467,8 @@
  ["FlaxGPTNeoForCausalLM", "FlaxGPTNeoModel", "FlaxGPTNeoPreTrainedModel"]
  )
  _import_structure["models.gptj"].extend(["FlaxGPTJForCausalLM", "FlaxGPTJModel", "FlaxGPTJPreTrainedModel"])
+ _import_structure["models.gpt_neox"].extend(["FlaxGPTNeoXForCausalLM", "FlaxGPTNeoXModel", "FlaxGPTNeoXPreTrainedModel"])
+
  _import_structure["models.longt5"].extend(
  ["FlaxLongT5ForConditionalGeneration", "FlaxLongT5Model", "FlaxLongT5PreTrainedModel"]
  )
@@ -5699,6 +5712,8 @@
  ForwardContext,
  GPT2AdapterModel,
  GPT2ModelWithHeads,
+ GPTNeoXAdapterModel,
+ GPTNeoXModelWithHeads,
  GPTJAdapterModel,
  HoulsbyConfig,
  HoulsbyInvConfig,
@@ -6010,6 +6025,13 @@
  TFGPTJModel,
  TFGPTJPreTrainedModel,
  )
+ from .models.gpt_neox import (
+ TFGPTNeoXForCausalLM,
+ TFGPTNeoXForQuestionAnswering,
+ TFGPTNeoXForSequenceClassification,
+ TFGPTNeoXModel,
+ TFGPTNeoXPreTrainedModel,
+ )
  from .models.groupvit import (
  TF_GROUPVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
  TFGroupViTModel,
@@ -6399,6 +6421,7 @@
  from .models.encoder_decoder import FlaxEncoderDecoderModel
  from .models.gpt2 import FlaxGPT2LMHeadModel, FlaxGPT2Model, FlaxGPT2PreTrainedModel
  from .models.gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel, FlaxGPTNeoPreTrainedModel
+ from .models.gpt_neox import FlaxGPTNeoXForCausalLM, FlaxGPTNeoXModel, FlaxGPTNeoXPreTrainedModel
  from .models.gptj import FlaxGPTJForCausalLM, FlaxGPTJModel, FlaxGPTJPreTrainedModel
  from .models.longt5 import FlaxLongT5ForConditionalGeneration, FlaxLongT5Model, FlaxLongT5PreTrainedModel
  from .models.marian import FlaxMarianModel, FlaxMarianMTModel, FlaxMarianPreTrainedModel

src/transformers/adapters/__init__.py

@@ -113,6 +113,10 @@
  "GPT2AdapterModel",
  "GPT2ModelWithHeads",
  ],
+ "models.gpt_neox": [
+ "GPTNeoXAdapterModel",
+ "GPTNeoXModelWithHeads",
+ ],
  "models.gptj": ["GPTJAdapterModel"],
  "models.mbart": [
  "MBartAdapterModel",
@@ -217,6 +221,7 @@
  from .models.debertaV2 import DebertaV2AdapterModel
  from .models.distilbert import DistilBertAdapterModel, DistilBertModelWithHeads
  from .models.gpt2 import GPT2AdapterModel, GPT2ModelWithHeads
+ from .models.gpt_neox import GPTNeoXAdapterModel, GPTNeoXModelWithHeads
  from .models.gptj import GPTJAdapterModel
  from .models.mbart import MBartAdapterModel, MBartModelWithHeads
  from .models.roberta import RobertaAdapterModel, RobertaModelWithHeads

src/transformers/adapters/head_utils.py

@@ -381,6 +381,13 @@
  },
  "layers": [None, "classifier"],
  },
+ #GPT-NeoX
+ "GPTNeoXForCausalLM": {
+ "config": {
+ "head_type": "causal_lm",
+ },
+ "layers": ["embed_out"],
+ },
  # GPT-J
  "GPTJForSequenceClassification": {
  "config": {

src/transformers/adapters/mixins/gpt_neox.py

@@ -0,0 +1,32 @@
+from typing import Iterable, Tuple
+
+import torch.nn as nn
+
+from ..layer import AdapterLayer
+from ..model_mixin import (
+ EmbeddingAdaptersMixin,
+ EmbeddingAdaptersWrapperMixin,
+ InvertibleAdaptersMixin,
+ ModelAdaptersMixin,
+ ModelWithHeadsAdaptersMixin,
+)
+
+
+class GPTNeoXDecoderBlockAdaptersMixin:
+ """Adds adapters to the TransformerBlock module of DistilBert."""
+
+ def _init_adapter_modules(self):
+ self.attention_adapters = AdapterLayer("mh_adapter", self.config)
+ self.output_adapters = AdapterLayer("output_adapter", self.config)
+ self.attention_adapters._init_adapter_modules()
+ self.output_adapters._init_adapter_modules()
+
+
+class GPTNeoXModelAdapterMixin(EmbeddingAdaptersMixin, InvertibleAdaptersMixin, ModelAdaptersMixin):
+ def iter_layers(self) -> Iterable[Tuple[int, nn.Module]]:
+ for i, layer in enumerate(self.base_model.layers):
+ yield i, layer
+
+
+class GPTNeoXModelWithHeadsAdaptersMixin(EmbeddingAdaptersWrapperMixin, ModelWithHeadsAdaptersMixin):
+ pass

src/transformers/adapters/models/auto/adapter_model.py

@@ -20,6 +20,7 @@
  ("bart", "BartAdapterModel"),
  ("mbart", "MBartAdapterModel"),
  ("gpt2", "GPT2AdapterModel"),
+ ("gpt_neox", "GPTNeoXAdapterModel"),
  ("gptj", "GPTJAdapterModel"),
  ("t5", "T5AdapterModel"),
  ("vit", "ViTAdapterModel"),
@@ -34,6 +35,7 @@
  ("bart", "BartModelWithHeads"),
  ("mbart", "MBartModelWithHeads"),
  ("gpt2", "GPT2ModelWithHeads"),
+ ("gpt_neox", "GPTNeoXModelWithHeads"),
  ("t5", "T5ModelWithHeads"),
  ]
 )

src/transformers/adapters/models/gpt_neox/__init__.py

@@ -0,0 +1,42 @@
+# flake8: noqa
+# There's no way to ignore "F401 '...' imported but unused" warnings in this
+# module, but to preserve other warnings. So, don't check this module at all.
+
+# Copyright 2020 The Adapter-Hub Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING
+
+from ....utils import _LazyModule
+
+
+_import_structure = {
+ "adapter_model": [
+ "GPTNeoXAdapterModel",
+ "GPTNeoXModelWithHeads"
+ ],
+}
+
+
+if TYPE_CHECKING:
+ from .adapter_model import GPTNeoXAdapterModel, GPTNeoXModelWithHeads
+
+else:
+ import sys
+
+ sys.modules[__name__] = _LazyModule(
+ __name__,
+ globals()["__file__"],
+ _import_structure,
+ )