DeepSpeed

3c6a5ca9 - [AutoTP] Support tied embeddings for row-parallel output-head training (#8680)

[AutoTP] Support tied embeddings for row-parallel output-head training (#8680) Follow-up to #8538 (part of #8173). ## Summary Row-parallel output-head training (#8538) rejected an `lm_head`/`embed_out` whose weight is tied to the input embedding, and that rejection could be bypassed silently (see *Bug* below). This PR fixes the bug and supports the tied case by sharding the tied `nn.Embedding` along the same hidden dimension as the row-parallel head, so both share one physical per-rank shard. ## Bug: tie silently broken when the embedding is visited first The tie was detected by scanning `self.module.modules()` *during* the replacement walk. The walk visits `embed_tokens` before `lm_head`, so when `partition_config` also matches the embedding, `_slice_embedding` swaps it for a new module with a new Parameter first. By the time `lm_head` was checked, no other module aliased its weight, the check passed, and the head was sharded while the embedding kept a separate copy — with no error. ```python AutoTPConfig(layer_specs=[ TPLayerSpec(patterns=[r".*embed_tokens\.weight$"], partition_type=PartitionType.ROW), TPLayerSpec(patterns=[r".*lm_head\.weight$"], partition_type=PartitionType.ROW), ]) # before: no error, model.lm_head.weight is model.embed_tokens.weight -> False # after: both modules share one hidden-dimension shard ``` Fixed by snapshotting weight ties from the original, pre-replacement model in `AutoTP.__init__`, following the approach the vocab-parallel path already takes for the same ordering issue (introduced in #8309, now `_tied_vocab_parallel_embedding_ids`). ## Changes - `HiddenParallelEmbedding` (`layers.py`): looks up each token's hidden-dimension slice and all-gathers the slices into the replicated activation (`GatherFromTensorParallelRegion`, so uneven hidden shards work). When tied, it reuses the row-parallel head's already-sharded weight `Parameter` and partition sizes; its universal-checkpoint metadata matches `LinearAllreduce`'s. `padding_idx`, `scale_grad_by_freq`, and Gemma3's scaled embedding are preserved; `max_norm`, `sparse`, and other custom forwards are rejected. - `AutoTP` (`auto_tp.py`): - Snapshots each module's weight-tie partners from the original model, and records embeddings tied to a row-parallel training head so `_slice_embedding` leaves them alone regardless of traversal order. - `_create_row_parallel_layer` validates all tie partners *before* partitioning the shared weight, then replaces every tied embedding alias with a `HiddenParallelEmbedding`. A tie to a non-`nn.Embedding` module still raises; a conflicting explicit embedding spec is superseded with a warning. - Docs: `autotp-training.md` describes the tied path and notes that `vocab_parallel_lm_head` is usually preferable for tied models, since a row-parallel head all-reduces full vocabulary logits. Not changed: untied embeddings that match a spec still go through `_slice_embedding`, whose un-gathered hidden slice is intentional for per-head embeddings such as T5's `relative_attention_bias`. ## Tests - Unit (`test_tp_partition_config_path.py`): tie is shared with/without an embedding spec visited first (the embedding-spec case reproduces the bug above on master); rejection for non-embedding ties and `max_norm` leaves the model untouched; Gemma3 scaled lookup matches the original module. - Integration (`TestTiedRowParallelOutputHeadTraining`, 2 GPUs): three optimizer steps against an unsharded reference comparing logits, the full tied gradient (lookup + projection contributions), and updated weights, covering even/uneven hidden (32/35), `padding_idx`, embedding-spec ordering, and ZeRO-0/2; universal-checkpoint affine map rebuilds the tied weight from uneven shards. Mutation check: detaching the embedding's weight in the lookup fails all four training cases. - `tests/unit/module_inject/` and `tests/unit/v1/autotp/`: 225 passed, 12 skipped on 3x NVIDIA GeForce RTX 5090 D. --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>