HF transformers main injects 'embedding_rowwise' into tp_plan for tied-embedding models; AutoTP rejects the whole plan

Description

Two days ago transformers merged #47579 "TP dtensor API inference + training", which makes PretrainedConfig.__init__ inject a new tp_plan entry for tied-embedding models:

# transformers/src/transformers/configuration_utils.py (main, since 861f4c41)
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",
    }

embedding_rowwise means sharding the embedding along the vocab dimension (their new tied embed_tokens/lm_head TP semantics).

DeepSpeed's converter is a strict allowlist and rejects the whole plan on any unknown style:

# deepspeed/module_inject/tp_plan_converter.py
SUPPORTED_STYLES = {"colwise", "colwise_rep", "colwise_gather_output", "rowwise", "replicated_with_grad_allreduce"}
...
raise ValueError(f"HuggingFace tp_plan contains unsupported partition style(s): {sorted(unsupported)}. ...")

As a result, any tied-embedding HF model (Qwen2/Qwen3 small, Llama, Gemma, ...) now fails AutoTP initialization against transformers main:

ValueError: HuggingFace tp_plan contains unsupported partition style(s): ['embedding_rowwise']. Applying only the
supported entries could shard one half of a column/row pair, so the plan is rejected as a whole. Provide an explicit
'tensor_parallel.partition_config' for this model instead.

This is not tied to a particular DeepSpeed change — it surfaced in the GPU CI of #8241 simply because that PR moved tests/unit/model_parallelism into the modal GPU workflow (which tests against transformers main), and test_tp_plan_real_models.py::test_qwen2_tied_lm_head_falls_back_to_replicated (a tie_word_embeddings=True Qwen2Config) started failing there.

Repro

import torch, deepspeed
from transformers import AutoModelForCausalLM, Qwen2Config   # transformers main, >= 861f4c41 (2026-08-21)

config = Qwen2Config(vocab_size=1000, hidden_size=128, intermediate_size=256,
                     num_hidden_layers=1, num_attention_heads=4, num_key_value_heads=4,
                     tie_word_embeddings=True)
model = AutoModelForCausalLM.from_config(config)

ds_config = {"train_micro_batch_size_per_gpu": 1,
             "tensor_parallel": {"autotp_size": 2},
             "zero_optimization": {"stage": 0}}
engine = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
# -> ValueError: unsupported partition style(s): ['embedding_rowwise']

(2 GPUs; the ValueError is raised at plan-conversion time.)

Semantic note

There is a policy divergence to resolve, not just a string to whitelist:

• DeepSpeed today: tied embed_tokens/lm_head stay replicated (that is exactly what the failing test asserts: lm_head.weight is embed_tokens.weight, shape unchanged).
• transformers new intent: tied embedding should be sharded row-wise along vocab (embedding_rowwise), i.e. vocab-parallel with gather on use.

Options:

1. Map embedding_rowwise to a supported behavior (vocab-parallel embedding / rowwise with tied handling), or
2. Recognize it and intentionally fall back to the replicated-tied behavior DeepSpeed has today (skip the entry with a log line), or
3. Support it only when the model's lm_head side is also planned accordingly.

Note the transformers-side injection is unconditional for tie_word_embeddings=True, so option 2 (or a config-level way to say "keep tied replicated") is needed regardless until vocab-parallel tied embeddings are supported.

Environment

• DeepSpeed: #8241 head (606ec52) — but behavior is independent of that PR
• transformers: main @ 861f4c41ee (2026-08-21) or later; pinned releases without #47579 are unaffected
• Found via modal-torch-latest GPU CI (which builds transformers from main)

Linking for context: #8241 (CI run where this surfaced), failing test tests/unit/v1/autotp/test_tp_plan_real_models.py::test_qwen2_tied_lm_head_falls_back_to_replicated.

Start with deepspeed/module_inject/tp_plan_converter.py and inspect how unsupported Hugging Face styles are handled. Run tests/unit/v1/autotp/test_tp_plan_real_models.py::test_qwen2_tied_lm_head_falls_back_to_replicated against transformers main, then review the tied-embedding assertions and related AutoTP tests. Done means the test passes with an explicit, reviewed policy for embedding_rowwise handling.