[AutoTP] Support tied embeddings for row-parallel output-head training

deepspeed/module_inject/auto_tp.py

@@ -244,6 +244,14 @@ def __init__(self,
                     continue
                 for tied_embedding in embedding_modules_by_weight_id.get(id(module.weight), []):
                     self._tied_vocab_parallel_embedding_ids.add(id(tied_embedding))
+        # Snapshot weight ties from the original model: once an earlier sibling such as a tied
+        # embedding is sliced, its new Parameter no longer aliases the head's weight, so a check
+        # made during the replacement walk would miss the tie and silently break it.
+        self._tied_weight_partners = self._find_tied_weight_partners()
+        # Embeddings tied to a row-parallel training output head are replaced by
+        # _create_row_parallel_layer with a HiddenParallelEmbedding sharing the head's shard, so
+        # _slice_embedding must leave them alone for the same reason as the vocab-parallel case.
+        self._tied_row_parallel_embedding_ids = self._find_tied_row_parallel_embedding_ids()
         self._gathered_column_tie_fallbacks_configured = False
         self._tied_gathered_column_module_names = set()
         TensorParallel_Layer.set_keep_module_on_host(keep_module_on_host)
@@ -495,14 +503,15 @@ def _create_row_parallel_layer(self, module, spec: TPLayerSpec, name: str):
             if spec.shape is not None or spec.get_partition_dim() != 1:
                 raise NotImplementedError(
                     "Row-parallel output-head training requires an unreshaped input-dimension shard.")
-            tied = any(weight is module.weight for other in self.module.modules() if other is not module
-                       for name, weight in other.named_parameters(recurse=False) if name == 'weight')
-            if self.mp_size > 1 and tied:
-                raise NotImplementedError(
-                    "Row-parallel output-head training cannot shard a tied weight. "
-                    "Keep the tied embedding and output head replicated until coupled embedding sharding is supported."
-                )
-            return LinearAllreduceWithReplicatedInput(module, self.mp_group, name=name, tp_meta=self.tp_meta)
+            tied_embeddings = []
+            if self.mp_size > 1:
+                tied_embeddings = self._validate_row_parallel_head_ties(module)
+            row_parallel_head = LinearAllreduceWithReplicatedInput(module,
+                                                                   self.mp_group,
+                                                                   name=name,
+                                                                   tp_meta=self.tp_meta)
+            self._replace_tied_row_parallel_embeddings(row_parallel_head, tied_embeddings)
+            return row_parallel_head
 
         if spec.shape is not None:
             return SubParamLinearAllreduce(
@@ -540,6 +549,75 @@ def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
             )
         return LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output, tp_meta=self.tp_meta)
 
+    def _find_tied_weight_partners(self):
+        modules_by_weight_id = {}
+        for module in self.module.modules():
+            weight = module._parameters.get("weight")
+            if weight is not None:
+                modules_by_weight_id.setdefault(id(weight), []).append(module)
+        tied_weight_partners = {}
+        for modules in modules_by_weight_id.values():
+            if len(modules) < 2:
+                continue
+            for module in modules:
+                tied_weight_partners[id(module)] = [other for other in modules if other is not module]
+        return tied_weight_partners
+
+    def _find_tied_row_parallel_embedding_ids(self):
+        if not self.training_mode or self.partition_config is None:
+            return set()
+        model_type = self._get_model_type()
+        embedding_ids = set()
+        for name, module in self.module.named_modules():
+            if not isinstance(module, nn.Linear) or not self._is_lm_head_name(name):
+                continue
+            if self._is_vocab_parallel_lm_head(module, name):
+                continue
+            spec = self.partition_config.find_matching_spec(name + ".weight", model_type)
+            if spec is None or spec.partition_type != PartitionType.ROW:
+                continue
+            for partner in self._tied_weight_partners.get(id(module), []):
+                if isinstance(partner, nn.Embedding):
+                    embedding_ids.add(id(partner))
+        return embedding_ids
+
+    def _validate_row_parallel_head_ties(self, head):
+        partners = self._tied_weight_partners.get(id(head), [])
+        if any(not isinstance(partner, nn.Embedding) for partner in partners):
+            raise NotImplementedError("Row-parallel output-head training can only share its weight with nn.Embedding "
+                                      "modules, which are then sharded along the same hidden dimension.")
+        for partner in partners:
+            HiddenParallelEmbedding.validate_embedding(partner)
+        # Validated partners were skipped by _slice_embedding, so they still alias the head's weight here.
+        return self._find_tied_embedding_modules(head)
+
+    def _replace_tied_row_parallel_embeddings(self, row_parallel_head, tied_embeddings):
+        model_type = self._get_model_type()
+        replacements = {}
+        for embed_parent, embed_child_name, embed_module, embed_full_name in tied_embeddings:
+            spec = self.partition_config.find_matching_spec(embed_full_name + ".weight", model_type)
+            if spec is not None and spec.partition_type != PartitionType.ROW:
+                log_dist(
+                    f"AutoTP: the row-parallel output head '{row_parallel_head.name}' supersedes the configured "
+                    f"'{spec.partition_type.value}' partitioning for tied embedding '{embed_full_name}'; it shares "
+                    "the head's hidden-dimension shard instead.",
+                    ranks=[0],
+                    level=logging.WARNING)
+            if embed_module not in replacements:
+                hidden_parallel_embedding = HiddenParallelEmbedding(embed_module,
+                                                                    self.mp_group,
+                                                                    tied_row_parallel_linear=row_parallel_head,
+                                                                    name=embed_full_name,
+                                                                    tp_meta=self.tp_meta)
+                # Guards re-entry when the replacement walk later reaches this module.
+                setattr(hidden_parallel_embedding, "replaced", True)
+                replacements[embed_module] = hidden_parallel_embedding
+            setattr(embed_parent, embed_child_name, replacements[embed_module])
+            log_dist(
+                f"AutoTP: tied embedding '{embed_full_name}' shares the hidden-dimension shard of the "
+                f"row-parallel output head '{row_parallel_head.name}'",
+                ranks=[0])
+
     @staticmethod
     def _default_lm_head_patterns():
         return ("lm_head", "embed_out")
@@ -755,6 +833,10 @@ def _get_model_type(self) -> Optional[str]:
     def _slice_embedding(self, child, name, conv_linear_layer):
         if getattr(child, "replaced", False) == True:
             return
+        if id(child) in self._tied_row_parallel_embedding_ids:
+            # Replaced by _create_row_parallel_layer's HiddenParallelEmbedding, which shares the
+            # tied head's shard; slicing a separate copy here would break the tie.
+            return
         if self.vocab_parallel_lm_head and id(child) in self._tied_vocab_parallel_embedding_ids:
             # This embedding is tied to a vocab_parallel_lm_head and gets replaced by
             # _create_vocab_parallel_layer's VocabParallelEmbedding instead, sharing that

deepspeed/module_inject/layers.py

@@ -28,7 +28,7 @@
     "TensorParallel_Layer", "LinearAllreduce", "LinearAllreduceWithReplicatedInput", "LinearLayer",
     "LmHeadLinearAllreduce", "Yuan_LinearAllreduce", "Yuan_LinearLayer", "GateUpPack_LinearLayer",
     "Conv_LinearALlreduce", "fused_LinearLayer", "conv_LinearLayer", "SubParamLinearLayer", "SubParamLinearAllreduce",
-    "VocabParallelLinear", "VocabParallelEmbedding"
+    "VocabParallelLinear", "VocabParallelEmbedding", "HiddenParallelEmbedding"
 ]
 
 DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.INFERENCE
@@ -1212,6 +1212,104 @@ def _mark_uc_metadata(self):
                                 original_shape=self._orig_weight_shape)
 
 
+class HiddenParallelEmbedding(TensorParallel_Layer):
+    """Embedding lookup sharded along the hidden dimension (dim 1) for training.
+
+    Each rank looks up its slice of every token's embedding, and the slices are all-gathered
+    into the replicated activation that the following layers expect. When tied to a row-parallel
+    output head, whose weight is sharded along the same dimension, this module reuses that
+    layer's already-sharded weight ``Parameter`` and shard boundaries, so both modules read from
+    and accumulate gradients into the same physical tensor.
+    """
+
+    @staticmethod
+    def validate_embedding(module):
+        # max_norm renormalizes whole rows, which no rank holds, and sparse gradients cannot
+        # accumulate into a weight whose tied output head produces dense gradients.
+        if module.max_norm is not None or module.sparse:
+            raise NotImplementedError("Hidden-parallel embedding options max_norm and sparse are not supported.")
+        embedding_type = type(module)
+        if embedding_type.forward is nn.Embedding.forward:
+            return
+        if (embedding_type.__module__ == "transformers.models.gemma3.modeling_gemma3"
+                and embedding_type.__name__ == "Gemma3TextScaledWordEmbedding"):
+            return
+        raise NotImplementedError(f"Hidden-parallel sharding does not support the custom embedding forward of "
+                                  f"{embedding_type.__name__}.")
+
+    def __init__(self, module, mp_group=None, tied_row_parallel_linear=None, **kwargs):
+        super().__init__(mp_group, **kwargs)
+        self.validate_embedding(module)
+        self.weight = module.weight
+        # nn.Embedding has no bias, but the base class's extra_repr() reads self.bias.
+        self.bias = None
+        self.support_training = True
+        # Hidden sharding keeps every vocabulary row on every rank, so these apply unchanged.
+        self.padding_idx = module.padding_idx
+        self.scale_grad_by_freq = module.scale_grad_by_freq
+
+        if tied_row_parallel_linear is not None:
+            # The output head already partitioned and materialized its shard along the same
+            # dimension; share that Parameter verbatim, keeping its equivalent gather/partition hooks.
+            self._orig_weight_shape = tied_row_parallel_linear._orig_weight_shape
+            self._partition_sizes = tied_row_parallel_linear._partition_sizes
+            self.weight = tied_row_parallel_linear.weight
+        else:
+            self._orig_weight_shape = self._shape_before_zero3_partition(module.weight)
+            self._freeze_partition_sizes(self._orig_weight_shape[1])
+            if self._should_materialize_tp_partition():
+                self._tp_partition([self.weight])
+            self.config_tp_params(self.weight)
+
+        embed_scale = None
+        if type(module).forward is not nn.Embedding.forward:
+            embed_scale = module.embed_scale.to(self.weight.device)
+        self.register_buffer("embed_scale", embed_scale, persistent=False)
+        self._mark_uc_metadata()
+
+    def forward(self, input):
+        self._assert_compiled_if_deferred()
+        output = F.embedding(input,
+                             self.weight,
+                             padding_idx=self.padding_idx,
+                             scale_grad_by_freq=self.scale_grad_by_freq)
+        if self.mp_group is not None and self.tp_world_size > 1:
+            output = GatherFromTensorParallelRegion.apply(self.mp_group, output, self._partition_sizes)
+        if self.embed_scale is not None:
+            output = output * self.embed_scale.to(self.weight.dtype)
+        return output
+
+    @torch.no_grad()
+    def gather_params(self, params_list):
+        for idx, param in enumerate(params_list):
+            if param is None:
+                continue
+            if self.mp_group is None or self.tp_world_size == 1:
+                params_list[idx].data = param.data.contiguous()
+                continue
+            params_list[idx].data = self._all_gather_shards(param, self._partition_sizes, dim=1).contiguous()
+
+    @torch.no_grad()
+    def _tp_partition(self, params_list):
+        for idx, param in enumerate(params_list):
+            if param is None:
+                return
+            _partition = params_list[idx].split(self._partition_sizes, dim=1)[self.tp_index]
+            params_list[idx].data = self.move(_partition).detach()
+
+    def _mark_uc_metadata(self):
+        # Matches LinearAllreduce's weight metadata, so the label stays the same whichever side
+        # of a tied pair runs last during universal-checkpoint metadata collection.
+        self._set_param_uc_meta(self.weight,
+                                partition_type='row',
+                                partition_dim=1,
+                                logical_shape=self._orig_weight_shape,
+                                output_shape=(self._orig_weight_shape[0], ),
+                                partition_sizes=self._partition_sizes,
+                                target_partition_shape=tuple(self.weight.shape),
+                                original_shape=self._orig_weight_shape)
+
+
 class SubParamColumnParallel(LinearLayer):
     """Column-parallel layer whose shard concatenates one piece of every sub-parameter.
 

docs/_tutorials/autotp-training.md

@@ -139,15 +139,24 @@ identity rather than model configuration metadata such as `tie_word_embeddings`.
 Additional HuggingFace types such as `local_colwise` and `local_rowwise` are
 not yet handled and fall back to AutoTP preset-based partitioning.
 
-For an untied `lm_head` or `embed_out`, an explicit `row` partition rule also
+For an `lm_head` or `embed_out`, an explicit `row` partition rule also
 supports eager training with a replicated input. The head slices the input to
 match its weight shard, reduces the complete output, and reconstructs the full
 input gradient across tensor-parallel ranks. Bias stays replicated; uneven
 hidden dimensions and both flattened and sequence-shaped inputs are supported.
-The default training output-head layout remains column parallel. Explicit row
-training rejects tied weights and reshaped/non-input-dimension shards rather
-than silently breaking a parameter tie. Deferred DeepCompile collectives are
-not supported for this output-head path.
+The default training output-head layout remains column parallel. A head tied
+to its input embedding (`tie_word_embeddings`) is also supported: DeepSpeed
+shards every tied `nn.Embedding` along the same hidden dimension, so the pair
+keeps reading from and accumulating gradients into a single per-rank weight
+shard, and all-gathers each token's embedding slices into the full hidden
+activation. A conflicting explicit spec on the tied embedding is superseded
+with a warning. The embedding options `max_norm` and `sparse`, custom
+embedding forwards other than Gemma3's scaled embedding, weights tied to
+modules other than `nn.Embedding`, and reshaped/non-input-dimension shards are
+rejected rather than silently breaking a parameter tie. Deferred DeepCompile
+collectives are not supported for this output-head path. Note that a
+row-parallel head all-reduces the complete vocabulary logits; for tied models,
+`vocab_parallel_lm_head` avoids that and keeps logits sharded.
 
 If you need to override the model's built-in `tp_plan`, provide a
 `partition_config` in the DeepSpeed config -- it takes precedence.