diff --git a/docs/source/en/perf_infer_gpu_multi.md b/docs/source/en/perf_infer_gpu_multi.md index c7fa4e4e1e64..31972298904f 100644 --- a/docs/source/en/perf_infer_gpu_multi.md +++ b/docs/source/en/perf_infer_gpu_multi.md @@ -106,7 +106,8 @@ The `ParallelInterface` class maps each strategy name you can use in a `tp_plan` ```py class ParallelInterface(GeneralInterface): _global_mapping = { - "embedding_rowwise": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + "embedding_rowwise": EmbeddingParallel(shard_dim=0), + "embedding_colwise": EmbeddingParallel(shard_dim=1), "colwise_gather_output": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), "colwise_rep": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), "colwise": ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)), @@ -134,8 +135,9 @@ The table below describes each strategy. | Strategy | Description | |---|---| -| `ColwiseParallel` | Shards a weight on its output-feature dim (`Shard(0)` for a 2D `nn.Linear` weight, `Shard(1)` for `nn.Embedding`) and shards a 1D bias. Redistributes the input to `Replicate()`, then places the output at `output_layouts`. | -| `RowwiseParallel` | Shards a weight on its input-feature dim (`Shard(-1)`, or `Shard(0)` for `nn.Embedding`) and replicates the bias, which is added after the reduction. Each rank's forward produces a `Partial()` output holding its share of the sum, and `transform_output_post_forward` redistributes that to `output_layouts`. Reducing to `Replicate()` issues an all-reduce, and reducing to `Shard(1)` issues a reduce-scatter. | +| `ColwiseParallel` | Shards a linear weight on its output-feature dim (`Shard(0)` for a 2D `nn.Linear` weight) and shards a 1D bias. Redistributes the input to `Replicate()`, then places the output at `output_layouts`. | +| `RowwiseParallel` | Shards a linear weight on its input-feature dim (`Shard(-1)`) and replicates the bias, which is added after the reduction. Each rank's forward produces a `Partial()` output holding its share of the sum, and `transform_output_post_forward` redistributes that to `output_layouts`. Reducing to `Replicate()` issues an all-reduce, and reducing to `Shard(1)` issues a reduce-scatter. | +| `EmbeddingParallel` | Shards an `nn.Embedding` with the token ids replicated. `shard_dim=0` (`embedding_rowwise`, the usual choice) shards the vocabulary and all-reduces each rank's partial lookup into the full output. `shard_dim=1` (`embedding_colwise`) shards the hidden dim and all-gathers each rank's slice of the output. | | `PackedColwiseParallel` | A variant of `ColwiseParallel` for fused weights, for example `up_proj` and `gate_proj` packed into `gate_up_proj`. Use `split_factor` when a weight packs a number of layers other than two. | | `PackedRowwiseParallel` | The row-wise counterpart, for weights packed along the final dim. Replicates 1D parameters. | | `SequenceParallel` | Replicates the module's parameters and shards its input on `sequence_dim` (defaults to `1`). Used for norms that operate per-token, such as `LayerNorm` and `RMSNorm`. | diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index 36a6c0325293..ca9bc6d2e0ea 100755 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -507,13 +507,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", - } - # Remap layer types if needed remap_legacy_layer_types(config=self) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 64a1d3bb24d0..b0fd4cc4863a 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -19,6 +19,8 @@ import warnings from typing import TYPE_CHECKING +import torch + from ..dynamic_module_utils import custom_object_save from ..utils import is_torch_greater_or_equal, logging from ..utils.hub import create_and_tag_model_card @@ -90,6 +92,9 @@ def init_parallel_plans(self) -> None: self._tp_plan.update(self.config.base_model_tp_plan or {}) self._ep_plan.update(self.config.base_model_ep_plan or {}) self._fsdp_plan.update(self.config.base_model_fsdp_plan or {}) + # Shard the input embeddings on the vocabulary unless the plan already names them. + if self._tp_plan and isinstance(getattr(self, "embed_tokens", None), torch.nn.Embedding): + self._tp_plan = {"embed_tokens": "embedding_rowwise", **self._tp_plan} for name, module in self.named_children(): if plan := getattr(module, "_ep_plan", None): diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 59c3670d594f..b4e6500ee50a 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -170,6 +170,15 @@ def context_around_forward(self, module, mesh): def transform_output_post_forward(self, module, output, mesh): return output + @staticmethod + def _from_local_last_dim_shard(local, mesh, size): + """Create a `DTensor` for a tensor sharded across its last dimension. + The local shard can have a different size from the other ranks when `size` is not divisible by the TP size.""" + shape = (*local.shape[:-1], size) + return DTensor.from_local( + local, mesh, [Shard(-1)], run_check=False, shape=shape, stride=torch.empty(shape, device="meta").stride() + ) + def install_forward(self, module, mesh): """Install pre / around / post transforms by replacing module.forward.""" original_forward = module.forward @@ -185,7 +194,9 @@ def tp_forward(*args, **kwargs): class ColwiseParallel(TensorParallelLayer): - """Column-wise: weight & bias → Shard(0) (Embedding: Shard(1)); input replicated, output Shard(-1).""" + """Column-wise for linear layers: weight & bias → Shard(0); input replicated, output Shard(-1). + + The output width (`out_features`) need not divide by the TP size: the last rank then holds a shorter slice.""" def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True): self.input_layouts = input_layouts or Replicate() @@ -197,28 +208,16 @@ def should_use_local_tensors(self, module): uses_local_inference_kernel = isinstance(module, torch.nn.Linear) and not torch.is_grad_enabled() return use_local_quantized_path or uses_local_inference_kernel - def validate_param(self, module, param, mesh, parameter_name=None): - meta = module._parameters.get(param) - gathers_output = isinstance(self.output_layouts, Replicate) - if meta is None or not gathers_output: - return - - shard_dim = 1 if isinstance(module, torch.nn.Embedding) else meta.ndim - 2 - output_size = meta.shape[shard_dim] - tp_size = mesh.size() - if output_size % tp_size != 0: - parameter_name = parameter_name or param - layer_name = parameter_name.rsplit(".", 1)[0] - raise ValueError( - f"The output size of `{layer_name}` ({output_size}) must be divisible by the tensor parallel size " - f"({tp_size}) when gathering a colwise output." - ) + def install_forward(self, module, mesh): + if not hasattr(module, "out_features"): + raise ValueError(f"{type(module).__name__} has no `out_features`: colwise shards linear layers.") + return super().install_forward(module, mesh) def shard_param(self, module, param, mesh): meta = module._parameters.get(param) if meta is None: return - placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2) + placement = Shard(meta.ndim - 2) module._parameters[param] = torch.nn.Parameter( distribute_tensor(meta, mesh, [placement], src_data_rank=None), requires_grad=meta.requires_grad, @@ -261,18 +260,19 @@ def transform_output_post_forward(self, module, output, mesh): if self.should_use_local_tensors(module) and self.use_local_output and output_is_local_shard: return output if not isinstance(output, DTensor): - output = DTensor.from_local(output, mesh, [Shard(-1)], run_check=False) + # `out_features`, not the weight's shape: a packed or quantized weight doesn't show the full width + output = self._from_local_last_dim_shard(output, mesh, module.out_features) if output.placements != (self.output_layouts,): output = output.redistribute(placements=[self.output_layouts]) return output.to_local() if self.use_local_output else output class RowwiseParallel(TensorParallelLayer): - """Row-wise: weight → Shard(1), bias → Replicate (Embedding: weight → Shard(0)). + """Row-wise for linear layers: weight → Shard(1), bias → Replicate; input sharded on the last dim. - Linear input is sharded on the last dim; Embedding input is replicated. The module - forward produces a Partial output which the boundary redistribute reduces to - output_layouts (Replicate → allreduce, Shard(1) → reduce-scatter). + The module forward produces a Partial output which the boundary redistribute reduces to + output_layouts (Replicate → allreduce, Shard(1) → reduce-scatter). The input width (`in_features`) need not + divide by the TP size. """ def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True): @@ -289,19 +289,20 @@ def shard_param(self, module, param, mesh): meta = module._parameters.get(param) if meta is None: return - if isinstance(module, torch.nn.Embedding): - placement = Shard(0) - else: - # bias is replicated (added after the row-reduce); weight shards on input dim (-1) - placement = Replicate() if param == "bias" else Shard(-1) + # bias is replicated (added after the row-reduce); weight shards on input dim (-1) + placement = Replicate() if param == "bias" else Shard(-1) module._parameters[param] = torch.nn.Parameter( distribute_tensor(meta, mesh, [placement], src_data_rank=None), requires_grad=meta.requires_grad, ) + def install_forward(self, module, mesh): + if not hasattr(module, "in_features"): + raise ValueError(f"{type(module).__name__} has no `in_features`: rowwise shards linear layers.") + return super().install_forward(module, mesh) + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): - # Embedding runtime sharding needs a replicated input; Linear needs Shard(-1). - desired = Replicate() if isinstance(module, torch.nn.Embedding) else Shard(-1) + desired = Shard(-1) x = args[0] # A local kernel can use a plain input when the previous layer already split it. # avoid a redundant Tensor -> DTensor -> Tensor round trip. @@ -309,7 +310,11 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh): if self.should_use_local_tensors(module) and input_has_desired_layout and not isinstance(x, DTensor): return args, kwargs if not isinstance(x, DTensor): - x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False) + if isinstance(self.input_layouts, Shard): + # the input is this rank's slice of the module's declared `in_features` + x = self._from_local_last_dim_shard(x, mesh, module.in_features) + else: + x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False) if x.placements != (desired,): x = x.redistribute(placements=[desired]) if self.should_use_local_tensors(module): @@ -360,6 +365,50 @@ def transform_output_post_forward(self, module, output, mesh): return output +class EmbeddingParallel(TensorParallelLayer): + """Embedding with the token ids replicated on every rank. `shard_dim=0` is vocab-parallel: each rank holds a slice + of the vocabulary, and the ranks' partial lookups are summed. `shard_dim=1` shards the hidden dim: each rank looks + up a slice of every vector, and the slices are gathered. The sharded dim need not divide by the TP size.""" + + def __init__(self, *, shard_dim: int = 0, use_local_output: bool = True): + self.shard_dim = shard_dim + self.use_local_output = use_local_output + + def should_use_local_tensors(self, module): + return getattr(module, "_hf_quantized_needs_local_tp", False) + + def shard_param(self, module, param, mesh): + meta = module._parameters.get(param) + if meta is None: + return + module._parameters[param] = torch.nn.Parameter( + distribute_tensor(meta, mesh, [Shard(self.shard_dim)], src_data_rank=None), + requires_grad=meta.requires_grad, + ) + + def transform_inputs_pre_forward(self, module, args, kwargs, mesh): + x = args[0] + if self.should_use_local_tensors(module) and not isinstance(x, DTensor): + return args, kwargs + if not isinstance(x, DTensor): + x = DTensor.from_local(x, mesh, [Replicate()], run_check=False) + elif x.placements != (Replicate(),): + x = x.redistribute(placements=[Replicate()]) + if self.should_use_local_tensors(module): + x = x.to_local() + return (x,) + args[1:], kwargs + + def transform_output_post_forward(self, module, output, mesh): + if not isinstance(output, DTensor): + if self.shard_dim == 0: + output = DTensor.from_local(output, mesh, [Partial()], run_check=False) + else: + output = self._from_local_last_dim_shard(output, mesh, module.embedding_dim) + if output.placements != (Replicate(),): + output = output.redistribute(placements=[Replicate()]) + return output.to_local() if self.use_local_output else output + + class ReplicatedWithGradAllReduce(TensorParallelLayer): """Replicated parameter whose gradient is partial. @@ -945,7 +994,8 @@ class ParallelInterface(GeneralInterface): _global_mapping = ( { - "embedding_rowwise": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), + "embedding_rowwise": EmbeddingParallel(shard_dim=0), + "embedding_colwise": EmbeddingParallel(shard_dim=1), "colwise_gather_output": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), "colwise_rep": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()), "colwise": ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)), diff --git a/src/transformers/models/doge/configuration_doge.py b/src/transformers/models/doge/configuration_doge.py index 8518d9021458..8ff3d2a4fdb1 100644 --- a/src/transformers/models/doge/configuration_doge.py +++ b/src/transformers/models/doge/configuration_doge.py @@ -61,8 +61,8 @@ class DogeConfig(PreTrainedConfig): "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", "layers.*.mlp.router_gate": "colwise_gather_output", - "layers.*.mlp.down_embed": "rowwise_split_input", - "layers.*.mlp.up_embed": "rowwise_split_input", + "layers.*.mlp.down_embed": "embedding_rowwise", + "layers.*.mlp.up_embed": "embedding_rowwise", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/doge/modular_doge.py b/src/transformers/models/doge/modular_doge.py index 897afec72f3f..eb9e83d83ac9 100644 --- a/src/transformers/models/doge/modular_doge.py +++ b/src/transformers/models/doge/modular_doge.py @@ -91,8 +91,8 @@ class DogeConfig(PreTrainedConfig): "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", "layers.*.mlp.router_gate": "colwise_gather_output", - "layers.*.mlp.down_embed": "rowwise_split_input", - "layers.*.mlp.up_embed": "rowwise_split_input", + "layers.*.mlp.down_embed": "embedding_rowwise", + "layers.*.mlp.up_embed": "embedding_rowwise", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), diff --git a/src/transformers/models/gemma4/configuration_gemma4.py b/src/transformers/models/gemma4/configuration_gemma4.py index 0b78dfc1bf1d..6cc0e2a95fd3 100644 --- a/src/transformers/models/gemma4/configuration_gemma4.py +++ b/src/transformers/models/gemma4/configuration_gemma4.py @@ -256,15 +256,15 @@ class Gemma4VisionConfig(PreTrainedConfig): model_type = "gemma4_vision" base_model_tp_plan = { - "encoder.layers.*.self_attn.q_proj": "colwise", - "encoder.layers.*.self_attn.k_proj": "colwise", - "encoder.layers.*.self_attn.v_proj": "colwise", + "encoder.layers.*.self_attn.q_proj.linear": "colwise", + "encoder.layers.*.self_attn.k_proj.linear": "colwise", + "encoder.layers.*.self_attn.v_proj.linear": "colwise", "encoder.layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", "encoder.layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "encoder.layers.*.self_attn.o_proj": "rowwise", - "encoder.layers.*.mlp.gate_proj": "colwise", - "encoder.layers.*.mlp.up_proj": "colwise", - "encoder.layers.*.mlp.down_proj": "rowwise", + "encoder.layers.*.self_attn.o_proj.linear": "rowwise", + "encoder.layers.*.mlp.gate_proj.linear": "colwise", + "encoder.layers.*.mlp.up_proj.linear": "colwise", + "encoder.layers.*.mlp.down_proj.linear": "rowwise", } default_theta = 100.0 diff --git a/src/transformers/models/inkling/configuration_inkling.py b/src/transformers/models/inkling/configuration_inkling.py index 12e7c519efa9..2415f9b6638c 100644 --- a/src/transformers/models/inkling/configuration_inkling.py +++ b/src/transformers/models/inkling/configuration_inkling.py @@ -76,7 +76,6 @@ class InklingTextConfig(PreTrainedConfig): model_type = "inkling_text" base_config_key = "text_config" base_model_tp_plan = { - "embed_tokens": "embedding_rowwise", "layers.*.mlp.experts.gate_up_proj": "colwise", "layers.*.mlp.experts.down_proj": "rowwise", "layers.*.mlp.experts": "moe_tp_experts", diff --git a/src/transformers/models/inkling/modular_inkling.py b/src/transformers/models/inkling/modular_inkling.py index 17352fb4c5f1..438fe86d0d10 100644 --- a/src/transformers/models/inkling/modular_inkling.py +++ b/src/transformers/models/inkling/modular_inkling.py @@ -108,7 +108,6 @@ class InklingTextConfig(PreTrainedConfig): model_type = "inkling_text" base_config_key = "text_config" base_model_tp_plan = { - "embed_tokens": "embedding_rowwise", "layers.*.mlp.experts.gate_up_proj": "colwise", "layers.*.mlp.experts.down_proj": "rowwise", "layers.*.mlp.experts": "moe_tp_experts", diff --git a/src/transformers/models/muse_glimmer/configuration_muse_glimmer.py b/src/transformers/models/muse_glimmer/configuration_muse_glimmer.py index ed7454bd9d06..8c423cf8fff2 100644 --- a/src/transformers/models/muse_glimmer/configuration_muse_glimmer.py +++ b/src/transformers/models/muse_glimmer/configuration_muse_glimmer.py @@ -98,7 +98,6 @@ class MuseGlimmerTextConfig(PreTrainedConfig): model_type = "muse_glimmer_text" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "embed_tokens": "embedding_rowwise", "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", diff --git a/src/transformers/models/muse_glimmer/modular_muse_glimmer.py b/src/transformers/models/muse_glimmer/modular_muse_glimmer.py index c739d331f5eb..2e3ee75b3b9f 100644 --- a/src/transformers/models/muse_glimmer/modular_muse_glimmer.py +++ b/src/transformers/models/muse_glimmer/modular_muse_glimmer.py @@ -585,7 +585,6 @@ class MuseGlimmerTextConfig(Gemma2Config, PreTrainedConfig): model_type = "muse_glimmer_text" base_model_tp_plan = { - "embed_tokens": "embedding_rowwise", "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", diff --git a/src/transformers/models/nemotron_h/configuration_nemotron_h.py b/src/transformers/models/nemotron_h/configuration_nemotron_h.py index 175cf54d2334..8727a80dd49e 100644 --- a/src/transformers/models/nemotron_h/configuration_nemotron_h.py +++ b/src/transformers/models/nemotron_h/configuration_nemotron_h.py @@ -83,6 +83,7 @@ class NemotronHConfig(PreTrainedConfig): model_type = "nemotron_h" attribute_map = {"layer_types": "layers_block_type", "num_local_experts": "n_routed_experts"} base_model_tp_plan = { + "embeddings": "embedding_rowwise", "layers.*.mixer.experts.up_proj": "colwise", "layers.*.mixer.experts.down_proj": "rowwise", "layers.*.mixer.experts": "moe_tp_experts", diff --git a/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py b/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py index 62085959be4f..9a7e85300baa 100644 --- a/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py +++ b/src/transformers/models/qwen4_exp/configuration_qwen4_exp.py @@ -96,7 +96,7 @@ class Qwen4ExpTextConfig(PreTrainedConfig): "hyper_connection_mixer.input_mix_weight_down": "rowwise_split_input", # It's extremely important to shard this embedding as its size is ~45B == 90 GiB - we shard on dim 1, as the checkpoints # for it are sharded on dim0 (this way Concatenate and tp do not work on the same dim and everything is easy) - "layers.*.ple.ple_embedding.ngram_embedding": "colwise_gather_output", + "layers.*.ple.ple_embedding.ngram_embedding": "embedding_colwise", } base_model_pp_plan = None base_model_ep_plan = { diff --git a/src/transformers/models/qwen4_exp/modular_qwen4_exp.py b/src/transformers/models/qwen4_exp/modular_qwen4_exp.py index 9a94f54073c5..e23ca547f566 100644 --- a/src/transformers/models/qwen4_exp/modular_qwen4_exp.py +++ b/src/transformers/models/qwen4_exp/modular_qwen4_exp.py @@ -141,7 +141,7 @@ class Qwen4ExpTextConfig(Qwen3_5MoeTextConfig): "hyper_connection_mixer.input_mix_weight_down": "rowwise_split_input", # It's extremely important to shard this embedding as its size is ~45B == 90 GiB - we shard on dim 1, as the checkpoints # for it are sharded on dim0 (this way Concatenate and tp do not work on the same dim and everything is easy) - "layers.*.ple.ple_embedding.ngram_embedding": "colwise_gather_output", + "layers.*.ple.ple_embedding.ngram_embedding": "embedding_colwise", } base_model_fsdp_plan = { "embed_tokens": "free_full_weight", diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py index 3dc682a34dc5..f09f50acb809 100644 --- a/tests/tensor_parallel/test_tensor_parallel.py +++ b/tests/tensor_parallel/test_tensor_parallel.py @@ -234,18 +234,6 @@ def _make_dtensor_shard_op(self, mesh, placement, param_shape, local_shape): op._axis0_size = param_shape[0] return op - def test_colwise_gather_output_rejects_indivisible_out_features(self): - model = torch.nn.Module() - model.lm_head = torch.nn.Linear(8, 99) - model.tp_plan = {"lm_head": "colwise_gather_output"} - device_mesh = self.MockDeviceMesh(world_size=2, rank=0) - - with self.assertRaises(ValueError) as context: - tensor_parallel.apply_tensor_parallelism(model, device_mesh) - - self.assertIn("lm_head", str(context.exception)) - self.assertIn("divisible", str(context.exception)) - def test_colwise_uneven_local_shapes(self): module = torch.nn.Module() module.register_parameter("weight", torch.nn.Parameter(torch.empty(10, 32))) @@ -274,18 +262,21 @@ def test_rowwise_uneven_local_shapes(self): self.assertEqual(weight_shape, (32, expected_size)) self.assertEqual(bias_shape, (10,)) - def test_embedding_uneven_local_shapes(self): - rowwise_embedding = torch.nn.Embedding(10, 10) - rowwise_placement = self._get_parameter_placements(rowwise_embedding, RowwiseParallel())["weight"] + def test_linear_styles_reject_embedding(self): + for style in ("colwise", "rowwise"): + with self.subTest(style=style): + with self.assertRaises(ValueError) as context: + ALL_PARALLEL_STYLES[style].install_forward(torch.nn.Embedding(8, 16), mesh=None) + self.assertIn("linear layers", str(context.exception)) - colwise_embedding = torch.nn.Embedding(10, 10) - colwise_placement = self._get_parameter_placements(colwise_embedding, ColwiseParallel())["weight"] + def test_embedding_uneven_local_shapes(self): + rowwise = self._get_parameter_placements(torch.nn.Embedding(10, 10), ALL_PARALLEL_STYLES["embedding_rowwise"]) + colwise = self._get_parameter_placements(torch.nn.Embedding(10, 10), ALL_PARALLEL_STYLES["embedding_colwise"]) expected_local_sizes = (4, 4, 2) for rank, expected_size in enumerate(expected_local_sizes): - rowwise_shape = self._get_local_shape((10, 10), rowwise_placement, world_size=3, rank=rank) - colwise_shape = self._get_local_shape((10, 10), colwise_placement, world_size=3, rank=rank) - + rowwise_shape = self._get_local_shape((10, 10), rowwise["weight"], world_size=3, rank=rank) + colwise_shape = self._get_local_shape((10, 10), colwise["weight"], world_size=3, rank=rank) self.assertEqual(rowwise_shape, (expected_size, 10)) self.assertEqual(colwise_shape, (10, expected_size)) @@ -314,7 +305,7 @@ def test_shard_tensor_shape_consistency(self): }, "embedding_colwise": { "module": torch.nn.Embedding(32, 16), - "style": ColwiseParallel(), + "style": ALL_PARALLEL_STYLES["embedding_colwise"], "expected_shapes": {"weight": (32, 4)}, }, } diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py index 4a8a05889908..db200f58c5ca 100644 --- a/tests/test_tensor_parallel_mixin.py +++ b/tests/test_tensor_parallel_mixin.py @@ -98,14 +98,15 @@ def get_packed_grad_shard(grad, world_size, rank, dim): total_size = grad.shape[dim] # Packed weights have 2 blocks (gate and up) block_size = total_size // 2 - shard_block_size = block_size // world_size + # every shard but the last is ceil(block_size / world_size) long, like DTensor's split + shard_block_size = -(-block_size // world_size) # Build interleaved indices indices = [] for block_idx in range(2): # gate block, then up block block_offset = block_idx * block_size - start = block_offset + rank * shard_block_size - stop = block_offset + (rank + 1) * shard_block_size + start = block_offset + min(rank * shard_block_size, block_size) + stop = block_offset + min((rank + 1) * shard_block_size, block_size) indices.extend(range(start, stop)) # Select along the sharded dimension @@ -301,10 +302,9 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol): # interleaved slicing grad = get_packed_grad_shard(grad, world_size, rank, dim) else: - # regular slicing - shard_size = grad_tp.size(dim) - start = rank * shard_size - grad = grad.narrow(dim, start, shard_size) + # regular slicing: every shard but the last is ceil(size / world_size) long + start = rank * -(-grad.size(dim) // world_size) + grad = grad.narrow(dim, start, grad_tp.size(dim)) break if not torch.allclose(grad.cpu(), grad_tp.cpu(), atol=atol, rtol=rtol): @@ -549,12 +549,10 @@ def _get_tp_model_class(self): ) def _get_tp_config(self, tie_word_embeddings: bool | None = None): - """Tiny config with `vocab_size` rounded up to a multiple of the world size, as sharded dims (typically `lm_head`) have to be split across ranks.""" + """The tester's tiny config. Its vocab stays as it is: 99 does not divide by the world size, which is + the uneven sharding of `lm_head` that real vocabs need.""" config = self.model_tester.get_config() text_config = config.get_text_config() - remainder = text_config.vocab_size % self.tensor_parallel_size - if remainder: - text_config.vocab_size += self.tensor_parallel_size - remainder if tie_word_embeddings is not None: if hasattr(text_config, "tie_word_embeddings"): text_config.tie_word_embeddings = tie_word_embeddings @@ -599,6 +597,24 @@ def _skip_if_not_supported(self, expert_parallel: bool = False): # if hasattr(config, "vision_config") and config.vision_config is not None: # self.skipTest("VLM models are not yet supported in TP tests") + def test_tp_plan_shards_input_embeddings(self): + """A model whose TP plan shards its decoder must shard its input embeddings too: replicated, they cost + `vocab x hidden` on every rank, and a tied `lm_head` sharded on the vocabulary would not match the embedding + weight it shares. A plan that only covers the `lm_head` is not a decoder plan.""" + config, _ = self.model_tester.prepare_config_and_inputs_for_common() + for model_class in self.all_model_classes: + model = model_class(copy.deepcopy(config)) + if not any(key.rsplit(".", 1)[-1] != "lm_head" for key in model.tp_plan): + continue + embeddings = model.get_input_embeddings() + if not isinstance(embeddings, torch.nn.Embedding): + continue + name = next(name for name, module in model.named_modules() if module is embeddings) + self.assertIsNotNone( + _get_parameter_plan(f"{name}.weight", model.tp_plan, is_weight=True), + f"{model_class.__name__}: the TP plan leaves `{name}` replicated; shard it with `embedding_rowwise`", + ) + def test_moe_parallel_plans_shard_experts(self): """An MoE model's expert weights, the bulk of its parameters, must be sharded by every parallel plan it defines. A plan missing them keeps every expert on every rank without changing any output, so no numerical test sees it.