diff --git a/deepspeed/module_inject/tp_plan_converter.py b/deepspeed/module_inject/tp_plan_converter.py index 17579b6788a2..03cabb3a28a6 100644 --- a/deepspeed/module_inject/tp_plan_converter.py +++ b/deepspeed/module_inject/tp_plan_converter.py @@ -6,8 +6,11 @@ from typing import List, Dict, Optional from .autotp_config import TPLayerSpec, PartitionType -SUPPORTED_STYLES = {"colwise", "colwise_rep", "colwise_gather_output", "rowwise", "replicated_with_grad_allreduce"} +SUPPORTED_STYLES = { + "colwise", "colwise_rep", "colwise_gather_output", "rowwise", "replicated_with_grad_allreduce", "embedding_rowwise" +} # `colwise_rep` was renamed to `colwise_gather_output` in huggingface/transformers#42809. +# `embedding_rowwise` is injected for every tied-embedding model since huggingface/transformers#47579. class TPPlanConverter: @@ -17,15 +20,21 @@ class TPPlanConverter: def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]: """Convert HF tp_plan to DeepSpeed layer specs. - Entries whose style is not supported are converted to SKIP specs instead of invalidating - the whole plan. Discarding the plan used to send models like Llama4 or Qwen3 down the - heuristic path, which shards by name patterns alone and has no notion of the modules the - plan deliberately excluded — on Llama4 it wraps the MoE router, whose forward returns a - tuple, and breaks the model. A SKIP spec keeps such layers untouched on purpose while the - supported entries are still applied. - - Returns None only when no entry is convertible, so the caller can fall back to the - existing AutoTP path for models whose plan gives us nothing to work with. + A style outside SUPPORTED_STYLES raises ValueError and invalidates the whole plan, so a + style newly introduced upstream fails loudly and gets a deliberate decision rather than + being dropped on the quiet. Converting only the recognized entries could shard one half + of a column/row pair, and silently discarding the plan would send models like Llama4 or + Qwen3 down the heuristic path, which shards by name patterns alone and has no notion of + the modules the plan deliberately excluded. On Llama4 that path wraps the MoE router, + whose forward returns a tuple, and breaks the model. + + Supported styles that must stay whole become SKIP specs, which keeps those layers + untouched on purpose while the rest of the plan is still applied: `embedding_rowwise`, + because vocabulary-parallel embeddings are not supported yet, and + `replicated_with_grad_allreduce`, which additionally sums the gradient across the group. + + Returns None only for an empty plan, so the caller can fall back to the existing AutoTP + path for models that give us nothing to work with. """ if not hf_tp_plan: return None @@ -50,6 +59,10 @@ def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]: gather_output = partition_style != "colwise" elif partition_style == "rowwise": partition_type = PartitionType.ROW + elif partition_style == "embedding_rowwise": + # Vocabulary-parallel embeddings are not supported yet, so the embedding stays whole and + # the gathered-column tie fallback keeps any LM head tied to it replicated as well. + partition_type = PartitionType.SKIP else: # replicated_with_grad_allreduce, the only other supported style # The parameter stays whole; only its gradient needs summing across the group. partition_type = PartitionType.SKIP 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 c86dc71afa18..76fbb6d28e20 100644 --- a/tests/unit/module_inject/test_tp_partition_config_path.py +++ b/tests/unit/module_inject/test_tp_partition_config_path.py @@ -13,6 +13,7 @@ from deepspeed.module_inject.auto_tp import AutoTP, AutoTPConfig, PartitionType, TPLayerSpec from deepspeed.module_inject.layers import LinearLayer +from deepspeed.module_inject.tp_plan_converter import TPPlanConverter class SubAttn(nn.Module): @@ -176,6 +177,28 @@ def test_gathered_lm_head_falls_back_for_runtime_parameter_tie(): assert model.lm_head.weight is model.embed_tokens.weight +def test_tied_embedding_plan_leaves_lm_head_and_embedding_replicated(): + model = OutputModel(tied=True) + specs = TPPlanConverter.convert({"embed_tokens": "embedding_rowwise", "lm_head": "colwise_gather_output"}) + + autotp = AutoTP( + module=model, + all_reduce_linears=[], + prefix="", + state_dict=None, + linear_layer_setting=None, + orig_layer_impl=None, + partition_config=AutoTPConfig(layer_specs=specs), + ) + autotp.set_tensor_parallel_config(1, None) + autotp.update_linear_policies() + autotp._replace_module(model) + + assert isinstance(model.embed_tokens, nn.Embedding) + assert isinstance(model.lm_head, nn.Linear) + assert model.lm_head.weight is model.embed_tokens.weight + + def test_gathered_lm_head_uses_column_parallel_layer_when_output_dim_is_uneven(): model = OutputModel(tied=False) model.lm_head = nn.Linear(32, 101, bias=False) diff --git a/tests/unit/module_inject/test_tp_plan_converter.py b/tests/unit/module_inject/test_tp_plan_converter.py index f8446d60980d..7ca50aeaf892 100644 --- a/tests/unit/module_inject/test_tp_plan_converter.py +++ b/tests/unit/module_inject/test_tp_plan_converter.py @@ -126,6 +126,28 @@ def test_unsupported_style_rejects_whole_plan(self): with pytest.raises(ValueError, match="unsupported partition style"): TPPlanConverter.convert(hf_plan) + def test_tied_embedding_style_converts_to_skip(self): + """`embedding_rowwise` is injected for every tied-embedding model, so rejecting it would + strand those models on the heuristic path. Tied embeddings stay replicated here, and the + entry carries no gradient all-reduce because the parameter is never split.""" + hf_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "embed_tokens": "embedding_rowwise", + "lm_head": "colwise_gather_output", + } + specs = TPPlanConverter.convert(hf_plan) + + assert len(specs) == 4 + + embed_spec = [s for s in specs if "embed_tokens" in s.patterns[0]][0] + lm_head_spec = [s for s in specs if "lm_head" in s.patterns[0]][0] + + assert embed_spec.partition_type == PartitionType.SKIP + assert not embed_spec.grad_allreduce + assert lm_head_spec.partition_type == PartitionType.COLUMN + assert lm_head_spec.gather_output + def test_alternate_prefixes(self): """Test tp_plan with non-layers prefix""" hf_plan = {