Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
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>
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>
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>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5453799d02
ℹ️ 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".
| and target.numel() > 0 | ||
| and _uniform_vocab_shards[(tp_group, local_vocab_size, vocab_start_index, vocab_end_index)]) | ||
| if use_liger: | ||
| loss = _liger_vocab_cross_entropy(vocab_parallel_logits, target.to(dtype=torch.long), tp_group, ignore_index) |
There was a problem hiding this comment.
When a multi-rank data-parallel job calls the public API with backend="liger" and leaves tp_group=None, the preceding metadata logic treats TP as disabled, but this call passes None to Liger's distributed group argument, where it denotes the default world process group. Liger then interprets the DP replicas' full-vocabulary logits as separate vocabulary shards, producing an incorrect global loss and gradients. Fall back to the PyTorch implementation in this case or pass an explicit singleton TP group.
Useful? React with 👍 / 👎.
| vocab_parallel_ce_backend: Literal["torch", "liger"] = "torch" | ||
| """Optional Liger CE acceleration; the PyTorch reference remains the default.""" |
There was a problem hiding this comment.
Commit 5e07eda8cb904aad0dce2b9a9df136255b432c78 is a non-merge commit but its message has no Signed-off-by trailer, so it does not satisfy this repository's commit/CI requirements and will fail DCO-style validation. Recreate the commit with the configured author identity and the required sign-off.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
3 participants
Dependency and review scope
Depends on #8309 (PR-B from #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 #8173. It does not include row-parallel output-head training (#8538) or coupled tied-embedding sharding.
Changes
tensor_parallel.vocab_parallel_ce_backend: "liger"alongsidevocab_parallel_lm_head: true; default"torch"remains dependency-free.liger-kernel>=0.8.1) and retain the existing reduction / causal-label-shifting contract.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.
git diff --checkpassed.The distributed integration methods in
TestLigerVocabParallelCEwere executed via two-rank torchrun on the actual devices. Reproduce with pytest: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.
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.