[AutoTP] Support eager row-parallel output head training by 0z5a · Pull Request #8538 · deepspeedai/DeepSpeed

0z5a · 2026-09-16T09:48:41Z

Summary

Implements PR-E from #8173: eager training for explicitly row-parallel, untied lm_head / embed_out output heads. This PR is independent of the vocabulary-parallel CE work in PR-B #8309 and contains no PR-D changes.

Slice replicated activations along the last dimension using the frozen weight-shard geometry, including uneven hidden dimensions.
Reuse the existing row-parallel forward output reduction and replicated bias behavior.
Reconstruct and all-reduce the complete input gradient in backward so preceding replicated layers receive the correct gradient.
Preserve the original parameter identity, existing inference specialization, and default gathered-column training behavior.
Reject tied output weights at TP > 1, reshaped/non-input-dimension partitions, and deferred compiler collectives for this path.

Validation

Executed on jk01 with two NVIDIA L20 GPUs (46068 MiB each), in an isolated container: Python 3.12.3, PyTorch 2.13.0+cu130, CUDA 13.0, Transformers 5.14.1.

Partition/config regression: 17 passed.
Real two-rank DeepSpeed-engine integration: three configurations each completed five optimizer updates, with two accumulated micro-batches per update (SGD LR 0.05, ZeRO stage 0, FP32, clipping explicitly disabled in both paths).
Configurations cover lm_head / embed_out, bias / no bias, hidden dimensions 32 / 35, and flattened / sequence-shaped inputs.
Compared logits, CE loss, input gradients, replicated projection gradients, local head gradients, and updated weights against an independent unsharded PyTorch model after every update. Also checked the exact SGD update and that weights actually changed.
Changed-file pre-commit checks and git diff --check passed.

The integration method TestRowParallelOutputHeadTraining.test_five_optimizer_steps_match_unsharded_reference was executed through two-rank torchrun on the real devices. It can also be run through pytest:

PYTHONPATH=.:tests python3 -m pytest tests/unit/module_inject/test_tp_partition_config_path.py -q
PYTHONPATH=.:tests python3 -m pytest tests/unit/v1/autotp/test_autotp_training.py::TestRowParallelOutputHeadTraining -v -s

This validates the minimal eager training contract, not large-model convergence, mixed-precision/ZeRO checkpoint round trips, or DeepCompile. Coupled tied-embedding sharding remains outside this PR's scope.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 751e03e11d

ℹ️ 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 · 2026-09-16T09:51:23Z

        if self.conv_linear_layer:
            return Conv_LinearALlreduce(module, self.mp_group, name=name, tp_meta=self.tp_meta)
-        if name == "lm_head" or name == 'embed_out':
+        if name.rsplit('.', 1)[-1] in ("lm_head", "embed_out"):


Sign off the commit

This non-merge commit has no Signed-off-by: trailer, so it violates the repository's mandatory commit policy and will be rejected by the corresponding CI/DCO check. Recreate the commit using git commit --signoff with the configured author name and email.

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

Useful? React with 👍 / 👎.

jinyouzhi · 2026-09-18T04:24:22Z

        if self.conv_linear_layer:
            return Conv_LinearALlreduce(module, self.mp_group, name=name, tp_meta=self.tp_meta)
-        if name == "lm_head" or name == 'embed_out':
+        if name.rsplit('.', 1)[-1] in ("lm_head", "embed_out"):


DeepSpeed/deepspeed/module_inject/auto_tp.py

Line 543 in ef7d23f

def _is_lm_head_name(self, name):

Can we reuse like it?

jinyouzhi · 2026-09-18T04:24:52Z

        return None, grad_output, None


+class ScatterToTensorParallelRegion(torch.autograd.Function):


Could ScatterToTensorParallelRegion be considered the dual of GatherFromTensorParallelRegion? If so, can we use this relationship to design unit tests?

…ather duality Signed-off-by: 0z5a <dezhen.lu@uni-tuebingen.de>

0z5a · 2026-09-18T15:25:38Z

Reused _is_lm_head_name, added Scatter/Gather round-trip tests for outputs and gradients (including uneven splits), and resolved the main-branch conflicts. Two-L20 validation: 8 passed, including three five-step training comparisons. All changed-file pre-commit checks passed. Recreated the commit with the required sign-off.

jinyouzhi · 2026-09-20T08:02:03Z

LGTM

…eedai#8540) ## Dependency and review scope Depends on deepspeedai#8309 (PR-B from deepspeedai#8173), currently open. This branch is based on its head `1b4ba6b162e87b10bedbb7316466bd877e2cb278`; inherited PR-B commits retain their original authorship. The final commit is the optional Liger backend work. Please review that commit for this PR's independent changes. This is a draft until PR-B lands; the master comparison currently includes the dependency changes. This addresses the optional fused CE and memory-profiling direction in deepspeedai#8173. It does not include row-parallel output-head training (deepspeedai#8538) or coupled tied-embedding sharding. ## Changes - Add opt-in `tensor_parallel.vocab_parallel_ce_backend: "liger"` alongside `vocab_parallel_lm_head: true`; default `"torch"` remains dependency-free. - Lazy-load Liger's public vocabulary-parallel CE function (`liger-kernel>=0.8.1`) and retain the existing reduction / causal-label-shifting contract. - Use identical gathered shard metadata to select Liger on every TP rank only for equal vocabulary shards, nonempty CUDA tokens, and FP32/FP16/BF16 logits. Uneven shards or unsupported layouts uniformly fall back to the PyTorch reference. - Reject repeated retained-graph backward and higher-order gradients for Liger's destructive saved-buffer backward instead of silently producing incorrect gradients. - Add optional-backend numerical and real AutoTP training integration coverage, and document the compatibility boundaries. This fuses CE only, **not the linear projection**. Local vocabulary logits are still materialized. ## Real-device validation jk01: two NVIDIA L20 GPUs (46068 MiB each), isolated container, Python 3.12.3, PyTorch 2.13.0+cu130, CUDA 13.0, Transformers 5.14.1, Triton 3.7.1, liger-kernel 0.8.2. - Partition/config and non-distributed CE regressions: **42 passed**, 7 distributed cases deselected in that regression invocation. - Two-rank numerical tests: FP32/FP16/BF16, none/sum/mean, even/uneven vocab, ignored targets, unchanged input logits, retained-graph rejection, and higher-order rejection passed. - Real DeepSpeed-engine training: vocab 34 (Liger) and vocab 35 (uneven fallback), each **five optimizer updates**, with two accumulated micro-batches per update. SGD LR 0.05, FP32, ZeRO 0, clipping explicitly disabled. Logits, loss, gradients, updated sharded head weights/biases, and replicated projection weights/biases agree with an independent unsharded PyTorch model. - Changed-file pre-commit checks and `git diff --check` passed. The distributed integration methods in `TestLigerVocabParallelCE` were executed via two-rank torchrun on the actual devices. Reproduce with pytest: ```sh PYTHONPATH=.:tests python3 -m pytest tests/unit/v1/sequence_parallelism/test_liger_vocab_parallel_cross_entropy.py -v -s ``` ## Memory profile Compared local linear + vocabulary-parallel CE forward/backward + five SGD updates, two micro-batches each, after kernel warm-up and resetting weights/allocated peaks. BF16, TP=2, 1024 tokens per rank, hidden=128, global vocabulary=65536, local vocabulary=32768. Statistics are the maximum across ranks. | Backend | Peak allocated | Incremental peak over baseline | Short-run elapsed | | --- | ---: | ---: | ---: | | PyTorch | 512247296 bytes (488.52 MiB) | 478160384 bytes | 0.0606 s | | Liger | 310922240 bytes (296.52 MiB) | 276834304 bytes | 0.0779 s | Peak allocated decreases approximately 39.3%; incremental peak decreases approximately 42.1%. Losses and final local weights agree. This is a short single-sample memory comparison, **not evidence of throughput speedup**; allocated memory is not total device or reserved memory. Four-GPU TP x SP integration, mixed-precision/ZeRO checkpoint round trips, and large-model convergence are not validated by this two-GPU run. --------- Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

deepspeedai#8680) Follow-up to deepspeedai#8538 (part of deepspeedai#8173). ## Summary Row-parallel output-head training (deepspeedai#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 deepspeedai#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>

0z5a requested review from hwchen2017, loadams, tjruwase and tohtana as code owners September 16, 2026 09:48

chatgpt-codex-connector Bot reviewed Sep 16, 2026

View reviewed changes

0z5a changed the title ~~[AutoTP] Support eager row-parallel output head training (PR-E)~~ [AutoTP] Support eager row-parallel output head training Sep 16, 2026

0z5a mentioned this pull request Sep 16, 2026

[AutoTP] Add optional Liger vocabulary-parallel cross entropy #8540

Merged

jinyouzhi mentioned this pull request Sep 17, 2026

[RFC] Support sharding LM heads and adopting Online Softmax #8173

Open

delock self-requested a review September 17, 2026 06:41

jinyouzhi reviewed Sep 18, 2026

View reviewed changes

[AutoTP] Support row-parallel output-head training and test scatter/g…

ecc25ac

…ather duality Signed-off-by: 0z5a <dezhen.lu@uni-tuebingen.de>

0z5a force-pushed the feat/autotp-row-parallel-lm-head branch from 751e03e to ecc25ac Compare September 18, 2026 15:24

Merge branch 'master' into feat/autotp-row-parallel-lm-head

a02295d

delock approved these changes Sep 20, 2026

View reviewed changes

delock enabled auto-merge September 20, 2026 08:27

delock added this pull request to the merge queue Sep 20, 2026

Merged via the queue into deepspeedai:master with commit 468a8ad Sep 20, 2026
13 checks passed

jinyouzhi mentioned this pull request Sep 27, 2026

[AutoTP] Support tied embeddings for row-parallel output-head training #8680

Merged

jinyouzhi mentioned this pull request Oct 20, 2026

AutoTP: shard an untied input embedding along the vocabulary (`vocab_… #8859

Draft

Navigation Menu

Repository navigation

[AutoTP] Support eager row-parallel output head training - #8538

[AutoTP] Support eager row-parallel output head training#8538
delock merged 2 commits into
deepspeedai:masterfrom
0z5a:feat/autotp-row-parallel-lm-head

0z5a commented Sep 16, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Sep 16, 2026

Uh oh!

jinyouzhi Sep 18, 2026

Uh oh!

jinyouzhi Sep 18, 2026

Uh oh!

0z5a commented Sep 18, 2026

Uh oh!

jinyouzhi commented Sep 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		return None, grad_output, None


		class ScatterToTensorParallelRegion(torch.autograd.Function):

Conversation

0z5a commented Sep 16, 2026

Summary

Validation

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

jinyouzhi Sep 18, 2026

Choose a reason for hiding this comment

Uh oh!

jinyouzhi Sep 18, 2026

Choose a reason for hiding this comment

Uh oh!

0z5a commented Sep 18, 2026

Uh oh!

jinyouzhi commented Sep 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants