Skip to content

Commit 436ea72

Browse files
committed
Fix: route GEMMA4 output_norm to ctx_output instead of ctx_output_split
I mistakenly added routed output_norm for GEMMA4 and GEMMA4_MTP to ctx_output_split. output_norm is a 1D vector that gains no benefit from being split across GPUs and incurs unnecessary communication overhead. Reverted it to ctx_output instead.
1 parent 92900a3 commit 436ea72

1 file changed

Lines changed: 2 additions & 2 deletions

File tree

src/llama-load-tensors.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2153,7 +2153,7 @@ bool create_tensors_helper::create_gemma4_tensors(const LLM_TN & tn) {
21532153
model.per_layer_proj_norm = create_tensor(ctx_output, tn(LLM_TENSOR_PER_LAYER_PROJ_NORM, "weight"), {n_embd_per_layer}, 0);
21542154
}
21552155

2156-
model.output_norm = create_tensor(ctx_output_split, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
2156+
model.output_norm = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
21572157

21582158
int rope_freqs_flag = 0;
21592159

@@ -2239,7 +2239,7 @@ bool create_tensors_helper::create_gemma4_mtp_tensors(const LLM_TN & tn) {
22392239
}
22402240

22412241
model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
2242-
model.output_norm = create_tensor(ctx_output_split, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
2242+
model.output_norm = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
22432243
model.output = create_tensor(ctx_output_split, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED);
22442244
auto output_extra = create_tensor(ctx_output_split, "output_extra.weight", {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED);
22452245
if (output_extra) {

0 commit comments

Comments
 (0)