Dflash support for nemotron-3.5 (#26905)
* conversion: skip untrained DFlash embeddings * Add Nemotron DFlash support * Add DFlash NVFP4 support * Address review comments * add missing output_s for nvfp4 * Include change for keeping residual for last layer also if requested in future dflash models * Update conversion/qwen.py Defensive check, not needed Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * Fixing bug introduced by merge conflict --------- Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>
This commit is contained in:
co-authored by
Sigbjørn Skjæret
parent
6e62ba5384
commit
cc078b45b6
+16
-10
@@ -79,6 +79,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
|
||||
|
||||
const int64_t n_embd_inp = hparams.n_embd_inp_enc();
|
||||
|
||||
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
|
||||
// DSpark = DFlash + a semi-autoregressive Markov head and Confidence head
|
||||
//
|
||||
// TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4)
|
||||
@@ -97,6 +98,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
|
||||
}
|
||||
|
||||
fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0);
|
||||
fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
|
||||
output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
|
||||
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm
|
||||
|
||||
@@ -205,7 +207,7 @@ template <>
|
||||
llama_model_dflash::graph<true>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
|
||||
ggml_tensor * cur = build_inp_embd_enc();
|
||||
|
||||
cur = build_lora_mm(model.fc, cur);
|
||||
cur = build_lora_mm(model.fc, cur, model.fc_s);
|
||||
cb(cur, "fc_out", -1);
|
||||
|
||||
cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1);
|
||||
@@ -460,9 +462,9 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
|
||||
cb(cur, "ffn_norm", il);
|
||||
|
||||
cur = build_ffn(cur,
|
||||
layer.ffn_up, NULL, NULL,
|
||||
layer.ffn_gate, NULL, NULL,
|
||||
layer.ffn_down, NULL, NULL,
|
||||
layer.ffn_up, NULL, layer.ffn_up_s,
|
||||
layer.ffn_gate, NULL, layer.ffn_gate_s,
|
||||
layer.ffn_down, NULL, layer.ffn_down_s,
|
||||
NULL,
|
||||
LLM_FFN_SILU, LLM_FFN_PAR, il);
|
||||
cb(cur, "ffn_out", il);
|
||||
@@ -479,15 +481,17 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
|
||||
res->t_embd = cur;
|
||||
|
||||
// lm_head from the target model (shared via ctx_other)
|
||||
auto * output = model.output;
|
||||
auto * output = model.output;
|
||||
auto * output_s = model.output_s;
|
||||
if (output == nullptr) {
|
||||
GGML_ASSERT(cparams.ctx_other != nullptr);
|
||||
const auto * model_other = llama_get_model(cparams.ctx_other);
|
||||
GGML_ASSERT(model_other->output != nullptr && "DFlash decoder requires the target model's output projection");
|
||||
output = model_other->output;
|
||||
output = model_other->output;
|
||||
output_s = model_other->output_s;
|
||||
}
|
||||
|
||||
cur = build_lora_mm(output, cur);
|
||||
cur = build_lora_mm(output, cur, output_s);
|
||||
cb(cur, "result_output", -1);
|
||||
res->t_logits = cur;
|
||||
|
||||
@@ -655,15 +659,17 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
|
||||
cb(cur, "result_norm", -1);
|
||||
|
||||
// lm_head from the target model (shared via ctx_other)
|
||||
auto * output = model.output;
|
||||
auto * output = model.output;
|
||||
auto * output_s = model.output_s;
|
||||
if (output == nullptr) {
|
||||
GGML_ASSERT(cparams.ctx_other != nullptr);
|
||||
const auto * model_other = llama_get_model(cparams.ctx_other);
|
||||
GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection");
|
||||
output = model_other->output;
|
||||
output = model_other->output;
|
||||
output_s = model_other->output_s;
|
||||
}
|
||||
|
||||
cur = build_lora_mm(output, cur);
|
||||
cur = build_lora_mm(output, cur, output_s);
|
||||
cb(cur, "result_output", -1);
|
||||
res->t_logits = cur;
|
||||
|
||||
|
||||
@@ -177,8 +177,11 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
|
||||
auto * inp = build_inp_mem_hybrid();
|
||||
|
||||
ggml_tensor * inp_out_ids = build_inp_out_ids();
|
||||
const bool extract_final_inp = (size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer];
|
||||
|
||||
for (int il = 0; il < n_layer; ++il) {
|
||||
res->t_layer_inp[il] = inpL;
|
||||
|
||||
struct ggml_tensor * inpSA = inpL;
|
||||
|
||||
// norm
|
||||
@@ -195,7 +198,7 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
|
||||
cur = build_ffn_layer(cur, model, il);
|
||||
}
|
||||
|
||||
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
|
||||
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked && !extract_final_inp) {
|
||||
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
|
||||
}
|
||||
@@ -209,6 +212,13 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
|
||||
}
|
||||
|
||||
cur = inpL;
|
||||
if (extract_final_inp) {
|
||||
res->t_layer_inp[n_layer] = cur;
|
||||
|
||||
if (inp_out_ids && cparams.embeddings_nextn_masked) {
|
||||
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
}
|
||||
}
|
||||
|
||||
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user