From 53c67928e107e9e4671cbae6268d3f2d1b94d6ce Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Sun, 27 Sep 2026 10:16:33 +0000 Subject: [PATCH 1/3] [AutoTP] Detect tied weights before replacement for row-parallel output-head training The row-parallel output-head tie check scanned the module tree during the replacement walk. When a tied embedding matched a partition_config spec, it was sliced into a new Parameter before lm_head was visited, so the check no longer saw the tie and the head was sharded while the embedding kept a separate copy. Snapshot weight ties from the original model in __init__ and consult that snapshot instead. Signed-off-by: Jin, Youzhi --- deepspeed/module_inject/auto_tp.py | 20 ++++++++++++++++--- .../test_tp_partition_config_path.py | 19 ++++++++++++++---- 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py index ad71ee1c6319..6b6f5f92774c 100644 --- a/deepspeed/module_inject/auto_tp.py +++ b/deepspeed/module_inject/auto_tp.py @@ -244,6 +244,10 @@ 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_module_ids = self._find_tied_weight_module_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,9 +499,7 @@ 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: + if self.mp_size > 1 and id(module) in self._tied_weight_module_ids: 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." @@ -540,6 +542,18 @@ 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_module_ids(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_module_ids = set() + for modules in modules_by_weight_id.values(): + if len(modules) > 1: + tied_module_ids.update(id(module) for module in modules) + return tied_module_ids + @staticmethod def _default_lm_head_patterns(): return ("lm_head", "embed_out") diff --git a/tests/unit/module_inject/test_tp_partition_config_path.py b/tests/unit/module_inject/test_tp_partition_config_path.py index 5ea3eecd91bf..bba41790834d 100644 --- a/tests/unit/module_inject/test_tp_partition_config_path.py +++ b/tests/unit/module_inject/test_tp_partition_config_path.py @@ -212,10 +212,11 @@ def _build_legacy_lm_head_autotp(model, training_mode=False): return autotp -def _build_row_output_head_autotp(model, head="lm_head", training_mode=False, mp_size=1): - config = AutoTPConfig(layer_specs=[ - TPLayerSpec(patterns=[rf".*{head}\.weight$"], partition_type=PartitionType.ROW), - ]) +def _build_row_output_head_autotp(model, head="lm_head", training_mode=False, mp_size=1, embedding_spec=False): + layer_specs = [TPLayerSpec(patterns=[rf".*{head}\.weight$"], partition_type=PartitionType.ROW)] + if embedding_spec: + layer_specs.append(TPLayerSpec(patterns=[r".*embed_tokens\.weight$"], partition_type=PartitionType.ROW)) + config = AutoTPConfig(layer_specs=layer_specs) autotp = AutoTP( module=model, all_reduce_linears=(), @@ -582,6 +583,16 @@ def test_explicit_row_parallel_lm_head_training_rejects_tied_weight(): assert model.embed_tokens.weight is weight +def test_explicit_row_parallel_lm_head_training_rejects_tie_after_embedding_is_sliced(): + # The embedding is visited and sliced before lm_head, so the tie must still be rejected + # rather than leaving the embedding and head on separate, silently untied weights. + model = OutputModel(tied=True) + autotp = _build_row_output_head_autotp(model, training_mode=True, mp_size=2, embedding_spec=True) + + with pytest.raises(NotImplementedError, match="cannot shard a tied weight"): + autotp._replace_module(model) + + @pytest.mark.parametrize("head", ["lm_head", "embed_out"]) def test_explicit_row_parallel_output_head_training_accepts_tp_greater_than_one(head): model = OutputModel(tied=False) From f96d8bb67c75061aaaeb3931a645697f227597f9 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Sun, 27 Sep 2026 10:50:35 +0000 Subject: [PATCH 2/3] [AutoTP] Support tied embeddings for row-parallel output-head training Row-parallel output-head training rejected a weight tied to the input embedding. Shard every tied nn.Embedding along the same hidden dimension with a new HiddenParallelEmbedding that reuses the head's weight shard and all-gathers each token's slices into the replicated activation. Tied embeddings are skipped by _slice_embedding so traversal order cannot break the tie, and all tie partners are validated before the shared weight is partitioned. Signed-off-by: Jin, Youzhi --- deepspeed/module_inject/auto_tp.py | 92 +++++++++++-- deepspeed/module_inject/layers.py | 100 +++++++++++++- docs/_tutorials/autotp-training.md | 19 ++- .../test_tp_partition_config_path.py | 53 +++++-- tests/unit/v1/autotp/test_autotp_training.py | 130 +++++++++++++++++- 5 files changed, 364 insertions(+), 30 deletions(-) diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py index 6b6f5f92774c..42bbe7fc804b 100644 --- a/deepspeed/module_inject/auto_tp.py +++ b/deepspeed/module_inject/auto_tp.py @@ -247,7 +247,11 @@ def __init__(self, # 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_module_ids = self._find_tied_weight_module_ids() + 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) @@ -499,12 +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.") - if self.mp_size > 1 and id(module) in self._tied_weight_module_ids: - 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( @@ -542,17 +549,74 @@ 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_module_ids(self): + 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_module_ids = set() + tied_weight_partners = {} for modules in modules_by_weight_id.values(): - if len(modules) > 1: - tied_module_ids.update(id(module) for module in modules) - return tied_module_ids + 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(): @@ -769,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 diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py index 1267c241fc4f..0780f8a571ea 100644 --- a/deepspeed/module_inject/layers.py +++ b/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. diff --git a/docs/_tutorials/autotp-training.md b/docs/_tutorials/autotp-training.md index 24b8d6cad751..adc127f562da 100644 --- a/docs/_tutorials/autotp-training.md +++ b/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. diff --git a/tests/unit/module_inject/test_tp_partition_config_path.py b/tests/unit/module_inject/test_tp_partition_config_path.py index bba41790834d..58dfe44dd9c0 100644 --- a/tests/unit/module_inject/test_tp_partition_config_path.py +++ b/tests/unit/module_inject/test_tp_partition_config_path.py @@ -11,12 +11,14 @@ import logging import pytest +import torch import torch.nn as nn from transformers import PreTrainedModel, PretrainedConfig from deepspeed.module_inject.auto_tp import AutoTP, AutoTPConfig, PartitionType, TPLayerSpec -from deepspeed.module_inject.layers import (LinearAllreduce, LinearLayer, LmHeadLinearAllreduce, VocabParallelLinear, - VocabParallelEmbedding, set_autotp_mode) +from deepspeed.module_inject.layers import (HiddenParallelEmbedding, LinearAllreduce, LinearLayer, + LmHeadLinearAllreduce, VocabParallelLinear, VocabParallelEmbedding, + set_autotp_mode) from deepspeed.module_inject.tp_plan_converter import TPPlanConverter from deepspeed.utils import logger as ds_logger from deepspeed.sequence.cross_entropy import (VocabParallelCausalLMLoss, configure_vocab_parallel_loss, @@ -573,24 +575,53 @@ def test_explicit_row_parallel_lm_head_is_not_overridden_by_its_name(): assert not isinstance(model.lm_head, LmHeadLinearAllreduce) -def test_explicit_row_parallel_lm_head_training_rejects_tied_weight(): +@pytest.mark.parametrize("embedding_spec", [False, True]) +def test_explicit_row_parallel_lm_head_training_shares_tied_weight_with_embedding(embedding_spec): + # With embedding_spec the embedding is visited before lm_head and matches a spec of its own; + # it must still end up on the head's Parameter rather than a separately sliced copy. model = OutputModel(tied=True) weight = model.embed_tokens.weight + _build_row_output_head_autotp(model, training_mode=True, mp_size=2, + embedding_spec=embedding_spec)._replace_module(model) - with pytest.raises(NotImplementedError, match="cannot shard a tied weight"): - _build_row_output_head_autotp(model, training_mode=True, mp_size=2)._replace_module(model) + assert isinstance(model.lm_head, LinearAllreduce) + assert isinstance(model.embed_tokens, HiddenParallelEmbedding) assert model.lm_head.weight is weight assert model.embed_tokens.weight is weight -def test_explicit_row_parallel_lm_head_training_rejects_tie_after_embedding_is_sliced(): - # The embedding is visited and sliced before lm_head, so the tie must still be rejected - # rather than leaving the embedding and head on separate, silently untied weights. +def test_explicit_row_parallel_lm_head_training_rejects_weight_tied_to_non_embedding(): + model = OutputModel(tied=False) + model.mirror = nn.Linear(32, 100, bias=False) + model.mirror.weight = model.lm_head.weight + + with pytest.raises(NotImplementedError, match="can only share its weight with nn.Embedding"): + _build_row_output_head_autotp(model, training_mode=True, mp_size=2)._replace_module(model) + assert isinstance(model.lm_head, nn.Linear) + + +def test_explicit_row_parallel_lm_head_training_rejects_unsupported_tied_embedding(): model = OutputModel(tied=True) - autotp = _build_row_output_head_autotp(model, training_mode=True, mp_size=2, embedding_spec=True) + model.embed_tokens.max_norm = 1.0 + weight = model.embed_tokens.weight - with pytest.raises(NotImplementedError, match="cannot shard a tied weight"): - autotp._replace_module(model) + with pytest.raises(NotImplementedError, match="max_norm"): + _build_row_output_head_autotp(model, training_mode=True, mp_size=2)._replace_module(model) + # Validation runs before the head partitions the shared weight. + assert isinstance(model.lm_head, nn.Linear) + assert model.lm_head.weight is weight + assert model.embed_tokens.weight is weight + + +def test_hidden_parallel_embedding_keeps_gemma3_scaled_lookup(): + gemma3 = pytest.importorskip("transformers.models.gemma3.modeling_gemma3") + embedding = gemma3.Gemma3TextScaledWordEmbedding(11, 8, padding_idx=0, embed_scale=8**0.5) + input_ids = torch.tensor([[0, 3, 10, 3]]) + expected = embedding(input_ids) + + hidden_parallel_embedding = HiddenParallelEmbedding(embedding) + + torch.testing.assert_close(hidden_parallel_embedding(input_ids), expected) @pytest.mark.parametrize("head", ["lm_head", "embed_out"]) diff --git a/tests/unit/v1/autotp/test_autotp_training.py b/tests/unit/v1/autotp/test_autotp_training.py index b98f1ef91cd4..59624c0dd1c7 100644 --- a/tests/unit/v1/autotp/test_autotp_training.py +++ b/tests/unit/v1/autotp/test_autotp_training.py @@ -25,7 +25,7 @@ from deepspeed.module_inject.auto_tp import AutoTP from deepspeed.module_inject.layers import (LinearAllreduce, LinearLayer, VocabParallelLinear, VocabParallelEmbedding, set_autotp_mode, is_autotp_training_mode, GatherFromTensorParallelRegion, - ScatterToTensorParallelRegion) + ScatterToTensorParallelRegion, HiddenParallelEmbedding) from deepspeed.module_inject.tp_shard import get_shard_size_list from unit.checkpoint.common import compare_lr_scheduler_states, compare_optimizer_states import os @@ -1290,6 +1290,134 @@ def test_five_optimizer_steps_match_unsharded_reference(self, hidden_dim, input_ reset_tp_model_init_state() +class TestTiedRowParallelOutputHeadTraining(DistributedTest): + """Catch a tied embedding and row-parallel head drifting apart or losing either gradient.""" + world_size = 2 + reuse_dist_env = False + + @pytest.mark.parametrize("hidden_dim,padding_idx,embedding_spec,zero_stage", [ + (32, None, False, 0), + (35, 3, False, 0), + (35, 3, True, 0), + (32, 3, False, 2), + ]) + def test_training_matches_unsharded_reference(self, hidden_dim, padding_idx, embedding_spec, zero_stage): + skip_on_device() + reset_tp_model_init_state() + torch.manual_seed(8173) + vocab_size = 67 + device = get_accelerator().current_device_name() + model = TiedRowParallelOutputTrainingModel(hidden_dim, vocab_size, padding_idx).to(device) + reference = deepcopy(model) + learning_rate = 0.05 + optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate) + reference_optimizer = torch.optim.SGD(reference.parameters(), lr=learning_rate) + layer_specs = [{"patterns": [r".*lm_head\.weight$"], "partition_type": "row"}] + if embedding_spec: + layer_specs.append({"patterns": [r".*embed_tokens\.weight$"], "partition_type": "row"}) + config = { + "train_micro_batch_size_per_gpu": 1, + "gradient_clipping": 0.0, + "zero_allow_untested_optimizer": True, + "zero_optimization": { + "stage": zero_stage + }, + "tensor_parallel": { + "autotp_size": self.world_size, + "partition_config": { + "use_default_specs": False, + "layer_specs": layer_specs, + }, + }, + } + try: + engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config) + head = engine.module.lm_head + assert isinstance(engine.module.embed_tokens, HiddenParallelEmbedding) + assert head.weight is engine.module.embed_tokens.weight + initial_weight = reference.lm_head.weight.detach().clone() + for step in range(3): + torch.manual_seed(100 + step) + input_ids = torch.randint(vocab_size, (2, 5), device=device) + if padding_idx is not None: + input_ids[0, 0] = padding_idx + labels = torch.randint(vocab_size, (2, 5), device=device) + reference_logits = reference(input_ids) + tp_logits = engine(input_ids) + torch.testing.assert_close(tp_logits, reference_logits, atol=1e-6, rtol=1e-5) + reference_loss = nn.functional.cross_entropy(reference_logits.reshape(-1, vocab_size), + labels.reshape(-1)) + tp_loss = nn.functional.cross_entropy(tp_logits.reshape(-1, vocab_size), labels.reshape(-1)) + reference_loss.backward() + engine.backward(tp_loss) + if zero_stage == 0: + # The tied gradient is the sum of the lookup and projection contributions. + full_grad = head.weight.grad.detach().clone() + head.gather_params([full_grad]) + torch.testing.assert_close(full_grad, reference.lm_head.weight.grad, atol=1e-6, rtol=1e-5) + engine.step() + reference_optimizer.step() + reference_optimizer.zero_grad() + full_weight = head.weight.detach().clone() + head.gather_params([full_weight]) + torch.testing.assert_close(full_weight, reference.lm_head.weight, atol=1e-6, rtol=1e-5) + torch.testing.assert_close(engine.module.projection.weight, + reference.projection.weight, + atol=1e-6, + rtol=1e-5) + assert not torch.equal(reference.lm_head.weight, initial_weight) + finally: + reset_tp_model_init_state() + + def test_universal_checkpoint_reconstructs_tied_weight(self): + skip_on_device() + reset_tp_model_init_state() + torch.manual_seed(8173) + hidden_dim = 35 + model = TiedRowParallelOutputTrainingModel(hidden_dim, 67, None) + original_weight = model.lm_head.weight.detach().clone() + config = { + "train_micro_batch_size_per_gpu": 1, + "tensor_parallel": { + "autotp_size": self.world_size, + "partition_config": { + "use_default_specs": False, + "layer_specs": [{ + "patterns": [r".*lm_head\.weight$"], + "partition_type": "row" + }], + }, + }, + } + try: + engine, _, _, _ = deepspeed.initialize(model=model, + optimizer=torch.optim.SGD(model.parameters(), lr=0.01), + config=config) + uc_info = getattr(engine.module, UNIVERSAL_CHECKPOINT_INFO) + affine_map = ParamAffineMap.from_dict(uc_info[AFFINE_MAP][AFFINE_MAP_PARAMS][r"^embed_tokens\.weight$"]) + shards = [None] * self.world_size + dist.all_gather_object(shards, + engine.module.lm_head.weight.detach().cpu(), + group=groups.get_tensor_model_parallel_group()) + assert shards[0].shape[1] != shards[1].shape[1] + torch.testing.assert_close(affine_map.rebuild(dict(enumerate(shards))), original_weight) + finally: + reset_tp_model_init_state() + + +class TiedRowParallelOutputTrainingModel(nn.Module): + + def __init__(self, hidden_dim, vocab_size, padding_idx): + super().__init__() + self.embed_tokens = nn.Embedding(vocab_size, hidden_dim, padding_idx=padding_idx) + self.projection = nn.Linear(hidden_dim, hidden_dim) + self.lm_head = nn.Linear(hidden_dim, vocab_size, bias=False) + self.lm_head.weight = self.embed_tokens.weight + + def forward(self, input_ids): + return self.lm_head(torch.tanh(self.projection(self.embed_tokens(input_ids)))) + + # @pytest.mark.sequential class TestParamsGather(DistributedTest): world_size = 4 From 12ae31433ae00c038e29a7afeed5120fc7ae5352 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Thu, 1 Oct 2026 15:36:00 +0000 Subject: [PATCH 3/3] [AutoTP] Cover head-first visit order for tied row-parallel output head Signed-off-by: Jin, Youzhi --- .../module_inject/test_tp_partition_config_path.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/unit/module_inject/test_tp_partition_config_path.py b/tests/unit/module_inject/test_tp_partition_config_path.py index 58dfe44dd9c0..5b6df2641f4f 100644 --- a/tests/unit/module_inject/test_tp_partition_config_path.py +++ b/tests/unit/module_inject/test_tp_partition_config_path.py @@ -575,12 +575,19 @@ def test_explicit_row_parallel_lm_head_is_not_overridden_by_its_name(): assert not isinstance(model.lm_head, LmHeadLinearAllreduce) +@pytest.mark.parametrize("head_first", [False, True]) @pytest.mark.parametrize("embedding_spec", [False, True]) -def test_explicit_row_parallel_lm_head_training_shares_tied_weight_with_embedding(embedding_spec): - # With embedding_spec the embedding is visited before lm_head and matches a spec of its own; - # it must still end up on the head's Parameter rather than a separately sliced copy. +def test_explicit_row_parallel_lm_head_training_shares_tied_weight_with_embedding(embedding_spec, head_first): + # With embedding_spec the embedding matches a spec of its own; it must still end up on the head's + # Parameter rather than a separately sliced copy, and must not be re-wrapped when the walk reaches + # it after the head has already replaced it. model = OutputModel(tied=True) weight = model.embed_tokens.weight + if head_first: + embed_tokens = model.embed_tokens + del model.embed_tokens + model.embed_tokens = embed_tokens + assert list(model._modules) == ["lm_head", "embed_tokens"] _build_row_output_head_autotp(model, training_mode=True, mp_size=2, embedding_spec=embedding_spec)._replace_module(model)