Add tensor parallel support for GPT-J

src/transformers/configuration_utils.py

@@ -349,13 +349,6 @@ def __post_init__(self, **kwargs):
         if per_layer_config is not None:
             self.per_layer_config = per_layer_config
 
-        # TODO: to support models whose input embedding module is not named `embed_tokens` (e.g. GPT-NeoX's `embed_in`).
-        if getattr(self, "tie_word_embeddings", False) and self.base_model_tp_plan is not None:
-            self.base_model_tp_plan = {
-                **self.base_model_tp_plan,
-                "embed_tokens": "embedding_rowwise",
-            }
-
     def __init_subclass__(cls, *args, **kwargs):
         super().__init_subclass__(*args, **kwargs)
         cls_has_custom_init = "__init__" in cls.__dict__

src/transformers/distributed/mixin.py

@@ -66,7 +66,10 @@ def init_parallel_plans(self) -> None:
 
         if self.base_model is self:
             self._pp_plan.update(self.config.base_model_pp_plan or {})
-            self._tp_plan.update(self.config.base_model_tp_plan or {})
+            tp_plan = dict(self.config.base_model_tp_plan or {})
+            if tp_plan and getattr(self.config, "tie_word_embeddings", False):
+                tp_plan.setdefault(getattr(self, "_input_embed_layer", "embed_tokens"), "embedding_rowwise")
+            self._tp_plan.update(tp_plan)
             self._ep_plan.update(self.config.base_model_ep_plan or {})
             self._fsdp_plan.update(self.config.base_model_fsdp_plan or {})
 

src/transformers/models/gptj/configuration_gptj.py

@@ -42,6 +42,14 @@ class GPTJConfig(PreTrainedConfig):
     ```"""
 
     model_type = "gptj"
+    base_model_tp_plan = {
+        "h.*.attn.q_proj": "colwise",
+        "h.*.attn.k_proj": "colwise",
+        "h.*.attn.v_proj": "colwise",
+        "h.*.attn.out_proj": "rowwise",
+        "h.*.mlp.fc_in": "colwise",
+        "h.*.mlp.fc_out": "rowwise",
+    }
     attribute_map = {
         "max_position_embeddings": "n_positions",
         "hidden_size": "n_embd",

src/transformers/models/gptj/modeling_gptj.py

@@ -109,7 +109,7 @@ def _split_heads(self, tensor, num_attention_heads, attn_head_size, rotary):
         """
         Splits hidden dim into attn_head_size and num_attention_heads
         """
-        new_shape = tensor.size()[:-1] + (num_attention_heads, attn_head_size)
+        new_shape = tensor.size()[:-1] + (-1, attn_head_size)
         tensor = tensor.view(new_shape)
         if rotary:
             return tensor
@@ -130,7 +130,7 @@ def _merge_heads(self, tensor, num_attention_heads, attn_head_size):
             tensor = tensor.permute(0, 2, 1, 3).contiguous()
         else:
             raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}")
-        new_shape = tensor.size()[:-2] + (num_attention_heads * attn_head_size,)
+        new_shape = tensor.size()[:-2] + (-1,)
         return tensor.view(new_shape)
 
     def _attn(
@@ -431,6 +431,8 @@ def _init_weights(self, module):
 
 @auto_docstring
 class GPTJModel(GPTJPreTrainedModel):
+    _input_embed_layer = "wte"
+
     def __init__(self, config):
         super().__init__(config)
 
@@ -566,6 +568,7 @@ def forward(
 )
 class GPTJForCausalLM(GPTJPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
 
     def __init__(self, config):
         super().__init__(config)