[AutoTP] Add support for vocab-parallel LM head and related configurations by jinyouzhi · Pull Request #8309 · deepspeedai/DeepSpeed

jinyouzhi · 2026-08-24T14:05:08Z

Motivation

Address PR-B in [RFC] Support sharding LM heads and adopting Online Softmax #8173

Changes

This pull request adds support for vocabulary-parallel language modeling (LM) heads in DeepSpeed's tensor parallelism (TP) system. It introduces a new VocabParallelLinear layer that keeps logits sharded across TP ranks for untied LM heads and installs a compatible distributed loss automatically. The changes ensure correct partitioning, validation, and integration with the rest of the DeepSpeed engine and configuration system.

Key changes include:

Vocabulary-parallel LM head support

Introduced a vocab_parallel_lm_head configuration option in TPTrainingConfig and passed it through the relevant APIs to enable this feature. [1] [2] [3] [4] [5]
Added logic in auto_tp.py to detect and correctly replace untied LM heads with the new VocabParallelLinear layer, along with validation to ensure the vocab head is not tied to the embedding. [1] [2] [3] [4]

New layer and loss integration

Implemented the VocabParallelLinear layer in layers.py, which ensures logits remain sharded and provides partitioning metadata. [1] [2]
Updated the DeepSpeed engine to automatically install a distributed causal LM loss compatible with vocab-parallel heads, and to handle model metadata/validation accordingly. [1] [2]

Miscellaneous

Made minor logic corrections and refactoring to ensure correct replacement heuristics and fallback configuration. [1] [2]
Exposed new loss functions and configuration utilities in the deepspeed.sequence module's public API.

These changes collectively enable efficient, distributed training of models with untied, sharded vocabulary projection heads and ensure compatibility with DeepSpeed's tensor parallelism ecosystem.

Limitations

TP x SP mesh are unsupported.

delock · 2026-09-04T23:29:28Z

Hi @jinyouzhi when this PR is ready for review just ping me, thanks!

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

Require an explicit sequence-parallel group when the legacy wrapper gathers losses, and document the distinct backward conventions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8f03c27ea0

ℹ️ 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".

Replace the output head in heuristic AutoTP and preserve original tied-weight evidence before embedding traversal mutates parameters. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

Resolve and validate the heuristic output head before transformer partitioning, cache the candidate for replacement, and skip tied-weight scans when the feature is disabled. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

jinyouzhi · 2026-09-08T12:26:51Z

Hi @jinyouzhi when this PR is ready for review just ping me, thanks!

Hi, this PR is ready for review now. Thank you for your patience! @delock

During development, I identified and fixed some issues related to SP in a separate PR: #8457.

delock · 2026-09-12T10:32:01Z

Hi @jinyouzhi Thanks for the PR — the loss kernel and backward-compat handling for vocab_sequence_parallel_cross_entropy look solid. Two changes requested, one suggestion:

1. License header (new file)
tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py uses the Copyright (c) Microsoft Corporation header. New files should use the plain header only:

# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team

2. vocab_parallel_lm_head silently overrides explicit lm_head specs
In _replace_with_config, the vocab-head check runs before spec lookup, so with the flag on, an explicit colwise_gather_output (or even rowwise) spec for lm_head in partition_config is silently ignored. Please raise (or at least warn) on this conflict instead, and add a test pinning the behavior.

Suggestion (non-blocking): a minimal save → load → forward-loss parity test for the vocab-sharded head would be good to have, since checkpoint paths have no coverage for the new VocabParallelLinear layout.

New files carry only the SPDX identifier and team line. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

Fail when vocab_parallel_lm_head finds no supported head instead of degrading into an ordinary gathered head with no distributed loss, and warn when the flag supersedes an explicit lm_head spec from partition_config or a HuggingFace tp_plan. Also partition the embedding with a row spec in the tied-head regression so it exercises the captured tie identities; the previous embedding_rowwise plan converts to a skip spec, which left the embedding untouched and the guard unreached. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

Save and reload a vocab-parallel LM head and compare the causal-LM loss, since no existing test exercises checkpoint paths for the VocabParallelLinear layout. The reloaded model starts from a different initialization so an unrestored checkpoint cannot pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

jinyouzhi · 2026-09-14T15:34:06Z

@delock Thank you for your detailed review! I have updated the PR according to your comments and guidance. Please take another look when you have a chance.

@jinyouzhi

101 rows over two ranks gives shards of 51 and 50, so the map has to carry the per-rank extents rather than assume an even split, and restoring at TP1 merges two unequal shards into one tensor. The fixture asserts the emitted map records that split before saving, so the test fails if conversion silently falls back to the vocabulary category. Adapted from @jinyouzhi's cross-tp vocab test in deepspeedai#8309, using the column partition already in tree so it carries no dependency on that PR. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

@jinyouzhi

101 rows over two ranks gives shards of 51 and 50, so the map has to carry the per-rank extents rather than assume an even split, and restoring at TP1 merges two unequal shards into one tensor. The fixture asserts the emitted map records that split before saving, so the test fails if conversion silently falls back to the vocabulary category. Adapted from @jinyouzhi's cross-tp vocab test in deepspeedai#8309, using the column partition already in tree so it carries no dependency on that PR. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

delock · 2026-09-16T03:25:04Z

LM-head name detection is duplicated in several coupled places — should all reuse lm_head_patterns

The ("lm_head", "embed_out") check is duplicated in _create_column_parallel_layer, _create_row_parallel_layer, _configure_gathered_column_tie_fallbacks, and the new _is_lm_head_name. These sites are coupled — e.g. the vocab-head skip in _configure_gathered_column_tie_fallbacks only stays correct if its notion of "LM head" matches the legacy checks — and for coupled duplication the right fix is reuse, not a second hardcoded copy. Please make these call sites resolve the name set from lm_head_patterns (with the current pair as the default) so there's a single source of truth.

delock · 2026-09-16T06:45:55Z

_vocab_metadata_cache holds strong references to process groups — in same-process train+inference setups (RL) where groups are rebuilt, this pins them and their NCCL resources for the life of the process. Suggest WeakKeyDictionary keyed by group to avoid the leak.

@jinyouzhi

Step 4 of the staging plan in deepspeedai#8252: AutoTP layers now emit an affine map, and the converter reads it on a path a real job takes rather than one a test supplies. Follows deepspeedai#8385 (the IR, the lowering constructors, the converter branch) and deepspeedai#8477 (the scale fix). Scope here is contiguous splits and replicated parameters; the fused QKV and Yuan shared-QK layouts are next and still carry `unsupported_reason`, so they emit no map and are unaffected. ### Where a map is built, and why it is not where you would expect The producer cannot build a map from what `collect_autotp_universal_checkpoint_info` sees. The conversion metadata carries eight fields and per-rank partition sizes are not among them — those were deliberately restore-only. But the layer does have them: `_freeze_partition_sizes` resolves them via `get_shard_size_list` while the layer is built, and they are not recoverable later from a shape alone. So the map is derived at mark time, in `_set_param_uc_meta`, which already receives `partition_sizes`, `sub_param_shard_widths`, `logical_shape` and `replicated`, and has `tp_world_size` in scope. That puts the derivation in one place rather than in each of the seven `_mark_uc_metadata` implementations, and `collect_...` then gathers what the layers produced. A layout that is not describable yet returns `None`, so conversion falls back to the categories. ### The map and the categories have to coexist A checkpoint carrying an affine map also carries the existing pattern lists, so a converter predating the map still reads it — that is the additivity §6.3 promises. The consequence is that the converter which prefers the map never consults those branches, which leaves their patterns looking unused and fails a strict conversion. They are superseded, not unused, so the affine branch now marks them consumed. This is the compatibility question raised as "one boundary for review" in deepspeedai#8385; it turns out both have to be present and the reader has to account for the other. This only surfaced end to end. The producer's output was correct in isolation, and the emitted maps matched a layout already verified by a passing resume — the failure was a strict-mode assert on rank 0 during conversion, which presented as rank 1 blocking on the following barrier. ### Tests `TestAffineMapProducer` builds a real AutoTP engine at TP2 and requires the producer to emit exactly the four maps the resume fixture previously supplied by hand. Those maps are not a guess: a full train → save → convert → resume cycle reproduces uninterrupted training through them, so matching them is evidence rather than self-consistency. `affine_resume_checkpoint` now takes its layout from the producer instead of injecting one, so the four save-convert-resume cases exercise the metadata a real job writes. All four pass — TP2 → TP1 and TP2 → TP2, through both the legacy and affine paths, compared against uninterrupted training on logits, losses, gradients, FP32 weights, both Adam moments and step counters. Validated on CPU/gloo (DS_ACCELERATOR=cpu LOCAL_SIZE=4): 109 passed across the resume matrix, the producer and coverage tests, the vocab cross-TP test, the affine unit suite and tests/unit/runtime/tensor_parallel/. The five remaining failures in that file are FusedAdam JIT-compile errors on this machine and are identical on upstream master. ### Review findings Three things came out of review, all fixed here: - affine_map was added to every conversion metadata dict, changing the schema even for layouts with no map. Now emitted only where there is one, so an existing layer publishes exactly what it did before. - The producer test checked that the expected maps were present but not that nothing extra was. It now asserts the pattern set exactly. - Parameters the TP machinery leaves untouched — embeddings, norm weights — were classified TP_REPLICATED with no map, so conversion still fell back to the category. TestAffineMapCoverage now asserts the general invariant: every pattern in any category must carry a map, with only AutoTP-unsupported layouts exempt. It fails without the fix. - TestUnevenVocabCrossTpResume covers a 101-row vocabulary head saved at TP2 and restored at TP1 and TP2. The 51/50 split is the point — restoring at TP1 merges two unequal shards, where an even-split assumption would surface. The fixture asserts the emitted map records that split before saving. Adapted from @jinyouzhi's cross-tp test in deepspeedai#8309, using the column partition already in tree so it carries no dependency on that PR. affine_ir_spec.md gains §6.4 on where a map is built, and a §6.3 rule on the map and categories coexisting. Related: deepspeedai#8252, deepspeedai#8230. cc @delock --------- Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>

…equence The compatibility wrapper stays available at deepspeed.sequence.cross_entropy.vocab_sequence_parallel_cross_entropy for existing callers, but new callers should use the general vocab_parallel_cross_entropy(tp_group=..., sp_group=...) API instead. Signed-off-by: iLeGend <824040212@qq.com>

_vocab_metadata_cache was a module-level dict keyed on tp_group, so it held a strong reference to any process group ever passed in and never released it. Remove it: the bare vocab_parallel_cross_entropy function now always re-validates (no cache), and VocabParallelCausalLMLoss caches its own resolved (vocab_start_index, vocab_end_index, global_vocab_size) on self after validating once, since it already owns tp_group for its whole lifetime. Reusing one loss instance for a different shard now raises instead of silently reusing stale metadata. Signed-off-by: iLeGend <824040212@qq.com>

jinyouzhi · 2026-09-16T17:05:41Z

_vocab_metadata_cache holds strong references to process groups — in same-process train+inference setups (RL) where groups are rebuilt, this pins them and their NCCL resources for the life of the process. Suggest WeakKeyDictionary keyed by group to avoid the leak.

Good catch. I removed the reference to the tp_group object from the cache entirely, rather than keeping it as a key — tp_group is now only held by VocabParallelCausalLMLoss itself (as self.tp_group), and the cache just stores the resolved start/end_idx.

I think this is safe enough for the multi-model/multi-shard/multi-loss case, since each loss instance owns its own tp_group and cache with matching lifetimes. Curious how this compares to WeakRef in your view — let me know if you'd still prefer that approach.

The ("lm_head", "embed_out") check was hardcoded separately in _replace (legacy_lm_head), _create_row_parallel_layer, _configure_gathered_column_tie_fallbacks, and _replace_last_linear_module, coupled to _is_lm_head_name without actually reusing it. Route them all through _is_lm_head_name, which now reads AutoTPConfig.lm_head_patterns when a partition_config is supplied and falls back to the legacy pair otherwise, so there is one place that defines what counts as an LM head. Signed-off-by: iLeGend <824040212@qq.com>

delock · 2026-09-17T05:00:18Z

@jinyouzhi can you fix CI failure? Thanks!

jinyouzhi · 2026-09-17T05:11:23Z

@jinyouzhi can you fix CI failure? Thanks!

OK sure

Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

Revert 53c4703 Keep _is_lm_head_name instance-bound so it can honor AutoTPConfig.lm_head_patterns, and update the stale static call in its regression test. Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

…#8538) ## Summary Implements PR-E from deepspeedai#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 deepspeedai#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: ```sh 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. Signed-off-by: 0z5a <dezhen.lu@uni-tuebingen.de> Co-authored-by: 0z5a <dezhen.lu@uni-tuebingen.de> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>

…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>

jinyouzhi mentioned this pull request Aug 24, 2026

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

Open

delock self-requested a review August 29, 2026 06:42

Achyuthan-S mentioned this pull request Aug 30, 2026

Define affine IR for checkpointing #8252

Open

Add vocab-parallel LM loss

8cf638a

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

jinyouzhi force-pushed the vocab_parallel_ce branch from ab4139a to 8cf638a Compare September 8, 2026 10:58

jinyouzhi and others added 2 commits September 8, 2026 11:11

Fix vocab-parallel integration

bf8b7e2

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

Preserve sequence loss gradients

8f03c27

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

jinyouzhi marked this pull request as ready for review September 8, 2026 11:51

jinyouzhi requested review from hwchen2017, loadams, tjruwase and tohtana as code owners September 8, 2026 11:51

Document sequence loss gradients

97078ce

Require an explicit sequence-parallel group when the legacy wrapper gathers losses, and document the distinct backward conventions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

chatgpt-codex-connector Bot reviewed Sep 8, 2026

View reviewed changes

Comment thread deepspeed/runtime/engine.py

Comment thread deepspeed/module_inject/auto_tp.py

Comment thread deepspeed/runtime/tensor_parallel/config.py

jinyouzhi and others added 2 commits September 8, 2026 12:03

jinyouzhi and others added 4 commits September 14, 2026 14:56

Use the plain license header in the new test

1764363

New files carry only the SPDX identifier and team line. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

Merge branch 'master' into vocab_parallel_ce

2998b2c

Achyuthan-S mentioned this pull request Sep 15, 2026

Emit affine maps from AutoTP layers #8519

Merged

Merge branch 'master' into vocab_parallel_ce

1b4ba6b

This was referenced Sep 16, 2026

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

Merged

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

Merged

jinyouzhi added 2 commits September 16, 2026 23:58

delock approved these changes Sep 17, 2026

View reviewed changes

Merge branch 'master' into vocab_parallel_ce

ac48c57

delock enabled auto-merge September 17, 2026 06:03

jinyouzhi and others added 2 commits September 17, 2026 06:06

fix TypeError in AutoTP._is_lm_head_name

53c4703

Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

Merge branch 'master' into vocab_parallel_ce

13a5ef7

delock disabled auto-merge September 17, 2026 06:21

jinyouzhi added 2 commits September 17, 2026 06:48

Refactor the fix TypeError in AutoTP._is_lm_head_name

e57cbf8

Revert 53c4703 Keep _is_lm_head_name instance-bound so it can honor AutoTPConfig.lm_head_patterns, and update the stale static call in its regression test. Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

Merge branch 'master' into vocab_parallel_ce

77ce8c6

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

Merged via the queue into deepspeedai:master with commit ef7d23f Sep 18, 2026
13 of 14 checks passed

jinyouzhi mentioned this pull request Sep 18, 2026

ZeRO-3 applies Muon's Newton-Schulz once per micro-batch, so gradient accumulation changes the optimizer #8443

Closed

tohtana mentioned this pull request Sep 27, 2026

DeepSpeed - Q3 Roadmap #8104

Open

10 of 22 tasks

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 10, 2026

DeepSpeed - Q4 2026 Roadmap #8672

Open

33 tasks

pengdurice mentioned this pull request Oct 19, 2026

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

Draft

Navigation Menu

Repository navigation

[AutoTP] Add support for vocab-parallel LM head and related configurations - #8309

[AutoTP] Add support for vocab-parallel LM head and related configurations#8309
delock merged 19 commits into
deepspeedai:masterfrom
jinyouzhi:vocab_parallel_ce

jinyouzhi commented Aug 24, 2026 •

edited

Loading

Uh oh!

delock commented Sep 4, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jinyouzhi commented Sep 8, 2026

Uh oh!

delock commented Sep 12, 2026 •

edited

Loading

Uh oh!

jinyouzhi commented Sep 14, 2026

Uh oh!

delock commented Sep 16, 2026

Uh oh!

delock commented Sep 16, 2026

Uh oh!

jinyouzhi commented Sep 16, 2026

Uh oh!

delock commented Sep 17, 2026

Uh oh!

jinyouzhi commented Sep 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

jinyouzhi commented Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Changes

Vocabulary-parallel LM head support

New layer and loss integration

Miscellaneous

Limitations

Uh oh!

delock commented Sep 4, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jinyouzhi commented Sep 8, 2026

Uh oh!

delock commented Sep 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jinyouzhi commented Sep 14, 2026

Uh oh!

delock commented Sep 16, 2026

Uh oh!

delock commented Sep 16, 2026

Uh oh!

jinyouzhi commented Sep 16, 2026

Uh oh!

delock commented Sep 17, 2026

Uh oh!

jinyouzhi commented Sep 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jinyouzhi commented Aug 24, 2026 •

edited

Loading

delock commented Sep 12, 2026 •

edited

Loading