From 10435529688405efb77fab9fce852eb90f2fb6eb Mon Sep 17 00:00:00 2001 From: Siddhant231xyz <87547498+Siddhant231xyz@users.noreply.github.com> Date: Sun, 13 Sep 2026 17:52:18 -0500 Subject: [PATCH 1/2] Key the tied-embedding TP rule on the model's input embedding attribute PreTrainedConfig injected a hard-coded {'embed_tokens': 'embedding_rowwise'} into base_model_tp_plan whenever tie_word_embeddings is set, which produces an invalid plan entry for models whose embedding attribute has another name (GPT-2 style 'wte', GPT-NeoX 'embed_in', ...), as noted in the TODO. Move the rule into ParallelInterface mixin plan initialization, keyed on the model's _input_embed_layer (default 'embed_tokens'), so the config stays declarative and models with differently named embeddings get vocab-parallel tied embeddings too. No behaviour change for existing models: tests/tensor_parallel + test_tensor_parallel_mixin pass (16), tp_plan tests for llama/qwen3_moe pass, and a tied-embedding Llama on a 2-rank CPU mesh still shards 15/20 params with identical logits. Related: #34789 --- src/transformers/configuration_utils.py | 7 ------- src/transformers/distributed/mixin.py | 5 ++++- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index 2f33f2c28d24..7a8310c07eae 100755 --- a/src/transformers/configuration_utils.py +++ b/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__ diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 646b8ac102cc..c46ba089478c 100644 --- a/src/transformers/distributed/mixin.py +++ b/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 {}) From 41ef36e4633f6c853913826709468de67782ad5a Mon Sep 17 00:00:00 2001 From: Siddhant231xyz <87547498+Siddhant231xyz@users.noreply.github.com> Date: Sun, 13 Sep 2026 17:56:02 -0500 Subject: [PATCH 2/2] Add tensor parallel support for GPT-J - Make GPTJAttention head-count agnostic (split/merge heads with -1) so column-sharded q/k/v projections work. - Add base_model_tp_plan to GPTJConfig (attn q/k/v colwise, out_proj rowwise, mlp fc_in colwise, fc_out rowwise) and a colwise_gather_output plan for lm_head. - Declare _input_embed_layer = 'wte' on GPTJModel so the tied input embedding is vocab-sharded. Verified on a 2-rank CPU/gloo mesh: sharded forward matches single-process forward (max abs diff 2e-7), 18/25 params sharded. tests/models/gptj: 112 passed. Related: #34789 --- src/transformers/models/gptj/configuration_gptj.py | 8 ++++++++ src/transformers/models/gptj/modeling_gptj.py | 7 +++++-- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/gptj/configuration_gptj.py b/src/transformers/models/gptj/configuration_gptj.py index 81bdf28c1a81..7fb75b05aa70 100644 --- a/src/transformers/models/gptj/configuration_gptj.py +++ b/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", diff --git a/src/transformers/models/gptj/modeling_gptj.py b/src/transformers/models/gptj/modeling_gptj.py index d43481d63c22..c50dfcafb7af 100644 --- a/src/transformers/models/gptj/modeling_gptj.py +++ b/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)