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

[delock]

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: [AutoTP] Replace tp_shard process-wide globals with per-model AutoTPMeta #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.

[delock]

Impact note: this drift currently only bites PRs that bring test_tp_plan_real_models into the modal workflow's tests/unit/v1 scope (e.g. #8241 tried; reverted there in e9612b3 until this is fixed). master's modal runs are green only because those tests still live outside the scope. Fixing this (or pinning transformers) is the prerequisite for moving the tp_plan tests under tests/unit/v1 per maintainer request.

[delock]

This issue is partially fixed, so related test can now move to GPU CI. We still need to decide whether we need to persue option 1 as a permanent fix.