Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

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

Open
  • Is something wrong?
#8,290 3 comments 1 reaction 0 assignees View on GitHub

Maintainers usually reply within 3 days

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
35/100
Issue type
Bug
Clarity
Mostly clear
Activity status
Active
Tech stack
python, pytorch
Domain
distributed-systems, machine-learning, testing-qa

Research direction

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.

Written by the indexing model from the issue text.

Description

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.

Dominant language
Python
Stars
43.3k
Forks
5.1k
Avg merge
4d 17h
Merged PRs (30d)
63

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from deepspeedai/DeepSpeed

All issues in deepspeedai/DeepSpeed

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.