DeepSpeed

68faa67c - Expect AutoTP to keep a row-parallel output head tied to the embedding (#8855)

Expect AutoTP to keep a row-parallel output head tied to the embedding (#8855) ## 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 3c6a5ca95 ("[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 --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>