Expect AutoTP to keep a row-parallel output head tied to the embedding

[jinyouzhi]

Summary

Fixes a stale test expectation. TestTiedOutputHead::test_row_parallel_tied_head_is_refused[False/True] in tests/unit/v1/moe/test_autoep_autotp_sharded_attention.py fails on master:

AssertionError: AutoTP sharded an output head whose weight is tied to the input embedding

The test (added in #8670) expects AutoTP to refuse a row-parallel output head whose weight is tied to the input embedding. Since 3c6a5ca ("[AutoTP] Support tied embeddings for row-parallel…"), AutoTP supports this case: it replaces the embedding with a HiddenParallelEmbedding that shares the head's sharded weight. The two changes landed a day apart, so the expectation was already stale when #8670 merged.

Change (test only, no runtime code modified):

• Rename the test to test_row_parallel_tied_head_keeps_the_tie.
• With and without zero.Init, check after deepspeed.initialize that the embedding is a HiddenParallelEmbedding and lm_head.weight is embed_tokens.weight.
• Run it in the CPU multi-rank nightly (cpu-torch-multi-latest). Before, it skipped on CPU and needed 4 GPUs, so no nightly ran it: the GPU nightly (ci/torch_latest.py) has exactly 2 GPUs. The check only builds the engine, so it needs no accelerator: drop the CPU skip and mark it sequential, because that nightly selects tests under unit/v1 only through its sequential phase. It keeps 4 ranks, which that nightly admits (LOCAL_SIZE=4).

The original concern of the test still holds: under zero.Init AutoTP gathers the partitioned head weight into a new parameter, and if the embedding kept the old one the two would train as separate copies. The identity check catches that. Numerical training parity of a tied row-parallel head is already covered by TestTiedRowParallelOutputHeadTraining in tests/unit/v1/autotp/test_autotp_training.py.

Related work

No open PR fixes this test. #8829 and #8848 both list these two failures as pre-existing and leave them unchanged. I searched open PRs touching TestTiedOutputHead and tied/vocab-parallel output heads.

Tests

cd tests
# CPU, as in cpu-torch-multi-latest
DS_ACCELERATOR=cpu LOCAL_SIZE=4 DS_DISABLE_REUSE_DIST_ENV=1 pytest -p no:randomly --forked -m sequential "unit/v1/moe/test_autoep_autotp_sharded_attention.py::TestTiedOutputHead"
# XPU
CUDA_VISIBLE_DEVICES=4,5,6,7 pytest -p no:randomly -m sequential "unit/v1/moe/test_autoep_autotp_sharded_attention.py::TestTiedOutputHead"

Test	master cc7b5125d	This PR
TestTiedOutputHead[zero_init=False]	failed on XPU (AssertionError above)	passed on CPU (4 gloo ranks) and XPU (4 ranks)
TestTiedOutputHead[zero_init=True]	failed on XPU (AssertionError above)	passed on CPU (4 gloo ranks) and XPU (4 ranks)

pre-commit run --files on the changed file passes.

Hardware and software:

• CPU: Intel Xeon Platinum 8457C, 4 ranks over gloo
• XPU: 4× Intel Arc Pro B70
• Level Zero driver 1.15.39122 (compute runtime 26.27.39122.11)
• torch 2.14.0.dev20260804+xpu, transformers 5.17.0

Compatibility matrix

Test-only change; no numerics, performance or memory impact.

	Validated	Not validated
ZeRO	Stage 3, with and without zero.Init (as in the test)	Other stages (not exercised by this test)
Precision	fp32
Offload		Not exercised
Parallelism	AutoTP 2 × AutoEP 2 folding, world size 4
Accelerator	CPU, XPU	CUDA (not run locally; the GPU nightly does not select sequential tests)

AI assistance

This PR was developed with AI assistance (GitHub Copilot CLI). I reviewed every changed line and ran the tests above.

cc @pengdurice @0z5a @delock

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-18T15:31:17.319857Z	568e27b	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[delock]

Hi @jinyouzhi , did this test get skipped in nightly? I didn't see nightly failure report. Can you take a look whether this test is covered? Thanks!