Repository navigation
llama : fix DFlash output head sharing - #30111
Conversation
Assisted-by: Codex
Assisted-by: Codex
Assisted-by: Codex
|
@CISC @ruixiang63 @ggerganov could you PTAL at this PR and see if it looks okay or not now ? |
ggerganov
added
the
merge ready
|
Before we merge, could you confirm whether this regression only affects Nemotron-3.5-Lightning-30B-A3B-NVFP4 with DFlash? The |
ggerganov
removed
the
merge ready
@ruixiang63 Yes for now, this regression only seems to affect Nemotron 3.5 Lightning Model. The fix makes output-head selection explicit:
This preserves the usual DFlash/DFlash2 head-sharing path and fixes Nemotron Lightning, where having its own input embeddings does not mean those embeddings should also be its output head. I verified Qwen3.8-27B + DFlash2 on the full SPEED-Bench coding subset: 80 conversations / 89 turns, using identical settings before and after the fix.
All generated outputs and per-request draft counters were identical. |
Forgot to answer this question in my earlier reply.
We don’t require every checkpoint to provide it. This PR writes the corresponding GGUF metadata through the Gemma DSpark converter; other DFlash/DFlash2 conversion paths are unchanged. If the GGUF flag is absent, a draft without its own output head defaults to borrowing the target’s head. |
| output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED); | ||
| if (output == nullptr && tok_embd != nullptr) { | ||
|
|
||
| if (output == nullptr && hparams.tie_word_embeddings) { |
There was a problem hiding this comment.
My question concerns this change in PR #29226 which appears to have introduced this issue.
Is this change necessary, or could we restore the logic from before PR #29226 o resolve the issue?
There was a problem hiding this comment.
Yes, I checked the checkpoint contents and agree we can remove the fallback. The Gemma DSpark checkpoints I inspected, already contain an explicit output head, so they don’t depend on it. I couldn’t find a checkpoint requiring the draft-embedding fallback.
I’ll simplify the fix to restore the previous behavior before #29226: use the draft’s output head when present; otherwise borrow the target’s head. I’ll also update the Gemma converter to require an explicit head and rerun the regression checks.
Assisted-by: Codex
|
@ruixiang63 @CISC @ggerganov , i have just removed the tie word embeddings fallback entirely. PTAL at this PR now. The change now uses the draft’s output head when present, otherwise borrows the target’s head. Tested the existing Gemma4 with Dspark, Nemotron Lighting with DFlash and Qwen 3.8 27B. All of them are working as expected. |
ruixiang63
added
the
merge ready
Upstream brings CUDA top-k/argsort/mmq fixes (PR ggml-org#28713, ggml-org#29953, ggml-org#30147, ggml-org#29453), a DFlash output-head fix (ggml-org#30111) and a UI fix (ggml-org#29668). Conflicts and resolution: - top-k.cu / argsort.cu: upstream reworked the top-k selection and the same radix path the fork added. Take upstream's shape-based dispatch and int64_t fixes, keep the fork workspace cap in the shared ggml_cuda_chunk_nrows (GGML_CUDA_ARGSORT_CHUNK_MB) and the fork benchmark-only GGML_CUDA_TOPK_IMPL selector used by bench-topk.py. - mmq.cu: upstream PR ggml-org#29953 fixes the same mul_mat_id padding over-read the fork patched in 1386ec2, but sizes the padding from the chosen J tile in both branches. Take upstream's version. - test-backend-ops.cpp: keep the fork GGML_TOPK_BENCH shapes and add upstream's chunk-spanning cases. Assisted-by: pi (deepseek-flash)
Overview
Fix a regression from #29226: drafts with embeddings but no output head incorrectly reused their embeddings instead of borrowing the target's head. Restrict embedding reuse to the Gemma DSpark layout.
Nemotron-3.5-Lightning-30B-A3B-NVFP4 + DFlash, three matched coding requests on #29226 (pooled accepted/proposed):
The fix matches pre-regression output and draft counters. Gemma tied-head, full-attention and SWA compatibility checks pass.
Requirements