DeepSpeed

20ea4543 - [AutoTP] Convert HF embedding_rowwise tp_plan entries to SKIP specs (#8294)

[AutoTP] Convert HF embedding_rowwise tp_plan entries to SKIP specs (#8294) ## 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 #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. --------- Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>