llama : split a tied output projection under split mode tensor - #3
Open
Piggidragon wants to merge 3 commits into
Open
Piggidragon wants to merge 3 commits into
Piggidragon wants to merge 3 commits into
Conversation
A model with tied embeddings has no output.weight - the output projection reuses the embedding table under the table's own name. That copy reaches a meta buffer, missed pattern_output_weight and fell through to MIRRORED, so every device held the whole table and ran the whole projection. Assisted-by: Claude Opus 5
Also shorten the comment that explains the pattern. Assisted-by: Claude Opus 5
Both carry the name token_embd.weight, only the copy may be split by vocab row. Assisted-by: Claude Opus 5
This was referenced Sep 17, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
A model with tied embeddings has no
output.weight. Its output projection reuses the embeddingtable under the table's own name, and that copy is what reaches a meta buffer under
--split-mode tensor. The split-state callback matches on the name,pattern_output_weightonlymatched
output\.weight, so the copy fell through toMIRRORED: every device held the whole tableand ran the whole projection instead of its share.
The table itself stays on the input device and never reaches the meta device, so matching
token_embd.weighthere is unambiguous.Split out of GenerelSchwerz#57.
Testing
llama-bench -ngl 99 -ts 50/50 -r 3 -sm tensor, gemma-4-26B-A4B (tied embeddings), two GPUs,RTX 4070 + RTX 3060, stock clocks:
+7.9% generation. The output projection runs once per token, so prefill is unaffected, as expected.
test-llama-archs -s 1at 1, 2, 3 and 4 virtual CUDA devices: passes. Built with-DLLAMA_FATAL_WARNINGS=ON.Requirements
Assisted-by:commit trailer.🤖 Generated with Claude Code