[AutoTP] Convert HF embedding_rowwise tp_plan entries to SKIP specs

[ebarkhordar]

Problem

huggingface/transformers#47579 (on main since 861f4c41, 2026-08-21) makes PretrainedConfig.__init__ inject "embed_tokens": "embedding_rowwise" into base_model_tp_plan whenever tie_word_embeddings is true. SUPPORTED_STYLES is a strict allowlist and convert() raises on any style outside it, rejecting the whole plan, so AutoTP plan conversion now fails for every tied-embedding model (Qwen2/Qwen3, Llama, Gemma) built against transformers main.

Fix

Recognize embedding_rowwise and convert it to a SKIP spec, which is option 2 in #8290: the entry is understood and the embedding is deliberately left replicated.

What follows is the behaviour DeepSpeed already implements rather than a new policy. lm_head still converts to a gathered column spec, and _configure_gathered_column_tie_fallbacks then sees that lm_head.weight is embed_tokens.weight and leaves both modules replicated, logging that coupled vocabulary-parallel embedding is not supported yet. A tied model is therefore left in the shape it has on a transformers release without the injection, with both modules replicated and the tie intact.

The entry maps to SKIP with grad_allreduce left false, unlike replicated_with_grad_allreduce. The parameter is never split, so register_replicated_grad_hooks must not register an all-reduce for it.

Styles that are still unknown continue to reject the whole plan. test_unsupported_style_rejects_whole_plan is unchanged and still passes.

Verification

Run on CPU in a container at edaa7221, against transformers main (5.16.0.dev0) and torch 2.13.0+cpu.

• The two added tests fail on master with the reported ValueError and pass with this change.
• tests/unit/module_inject/ and tests/unit/runtime/test_tp_plan_extraction.py: 47 passed, on Python 3.11 and on 3.12.
• pre-commit run --files on the three changed files passes yapf, check-torchdist, check-license and codespell; flake8 5.0.4 exits 0 on them under Python 3.11.
• Not verified here: test_qwen2_tied_lm_head_falls_back_to_replicated, which needs 2 GPUs. That is the test HF transformers main injects 'embedding_rowwise' into tp_plan for tied-embedding models; AutoTP rejects the whole plan #8290 reports as failing and the one this change is meant to restore.

Two things worth deciding separately

Scoping this to embedding_rowwise leaves the next transformers-side style to fail the same way, since the injection is unconditional and the allowlist is deny-by-default against a vocabulary DeepSpeed does not own. A general rule for unknown styles looks like a maintainer call rather than something to settle here.

Related to that, the convert() docstring says entries with an unsupported style become SKIP specs instead of invalidating the plan, but no code path does that, and none does after this change either: an unsupported style still raises before the loop is reached. The docstring and the raise arrived together in #8204, so I have left both alone. Happy to follow up once you have picked the policy.

Refs #8290. This covers the conversion failure only, and does not implement vocabulary-parallel tied embeddings (option 1 or 3 in that issue), so I have not used a closing keyword.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 79ef659526

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Add the required Signed-off-by trailer

This is a one-parent, non-merge commit, but its commit message has no Signed-off-by: trailer, so it does not satisfy the repository's mandatory commit requirement; recreate the commit with --signoff using the configured Git identity.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.

[delock]

Hi @ebarkhordar , thanks for your timely fix to this issue. For your opens:

1. Unknown entries should lead to an explicit fail, allowing concious decision made on each new entry. Let's see whether this strategy works.
2. Yes, the docstring needs to be updated to reflect this strategy, can you update doc string in this PR? Thanks!

[ebarkhordar]

Thanks, both are settled in 0aa0c41.

1. No change. An unknown style still raises and rejects the whole plan, and test_unsupported_style_rejects_whole_plan still pins that.

2. The convert() docstring now describes that strategy instead of the SKIP-on-unsupported behaviour it claimed, which no code path has implemented since the raise landed in (1/2) Implementing Compiler Pass for AutoTP #8204. It also names the two supported styles that do become SKIP specs, embedding_rowwise and replicated_with_grad_allreduce. One more line in the same docstring was stale and I fixed it there rather than leaving it for later: the function returns None only for an empty plan, not when no entry is convertible.

The commit touches that docstring and nothing else. yapf, check-torchdist, check-license and codespell pass on the file, and flake8 5.0.4 exits 0 on it under Python 3.11.

[delock]

Hi @ebarkhordar , thanks for your timely fix. Now we can move tied tests to modal GPU CI. I have a question for you, do you think we should persue option 1 in the long run? What is pros and cons for option 1? Thanks!