mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-04 02:37:27 +02:00
spec: add eagle3-v3 support for gpt-oss model (#25794)
This commit is contained in:
+13
-3
@@ -437,6 +437,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
|
||||
int32_t n_embd_dec = 0; // draft hidden size
|
||||
int32_t n_embd_enc = 0; // target_layer_ids_n * target_hidden_size
|
||||
int32_t n_embd_tgt = 0; // target model hidden size
|
||||
int32_t n_layer_tgt = 0; // target model layer count
|
||||
|
||||
const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices
|
||||
uint32_t target_layer_ids_n = 0;
|
||||
@@ -478,6 +479,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
|
||||
n_embd_tgt = llama_model_n_embd(model_tgt);
|
||||
n_embd_dec = llama_model_n_embd(model_dft);
|
||||
n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt;
|
||||
n_layer_tgt = llama_model_n_layer(model_tgt);
|
||||
|
||||
const int32_t n_b = (int32_t) llama_n_batch(ctx_dft);
|
||||
batch = llama_batch_init(/*n_tokens=*/ n_b, /*embd=*/ n_embd_dec, /*n_seq_max=*/ 1);
|
||||
@@ -510,9 +512,15 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
|
||||
}
|
||||
}
|
||||
|
||||
// turn on extraction of the target layers' input embeddings
|
||||
// turn on extraction of the target layers' hidden states
|
||||
for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
|
||||
llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
|
||||
if (target_layer_ids[k] < n_layer_tgt) {
|
||||
llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
|
||||
} else if (target_layer_ids[k] == n_layer_tgt) {
|
||||
llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ false);
|
||||
} else {
|
||||
GGML_ABORT("EAGLE3: target layer id %d exceeds target n_layer %d", target_layer_ids[k], n_layer_tgt);
|
||||
}
|
||||
}
|
||||
|
||||
// turn on extraction of the draft model's pre-norm hidden state
|
||||
@@ -600,7 +608,9 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
|
||||
features_buf.resize((size_t) n_tokens * n_embd_enc, 0.0f);
|
||||
|
||||
for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
|
||||
const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);
|
||||
const float * layer = target_layer_ids[k] < n_layer_tgt
|
||||
? llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k])
|
||||
: llama_get_embeddings_nextn(ctx_tgt);
|
||||
if (!layer) {
|
||||
GGML_ABORT("EAGLE3: target layer %d input not extracted.", target_layer_ids[k]);
|
||||
}
|
||||
|
||||
+16
-2
@@ -69,9 +69,14 @@ class LlamaModel(TextModel):
|
||||
target_config = {**target_config, **target_config["text_config"]}
|
||||
self.target_vocab_size = target_config["vocab_size"]
|
||||
|
||||
# target_layers: derived from target model layer count (low/mid/high)
|
||||
# target_layers: use the eagle3 config's explicit aux hidden-state layer ids
|
||||
# if present, else derive from the target layer count.
|
||||
target_num_layers = target_config["num_hidden_layers"]
|
||||
target_layers = [2, target_num_layers // 2, target_num_layers - 3]
|
||||
aux_layer_ids = eagle3_raw_config.get("eagle_aux_hidden_state_layer_ids")
|
||||
if aux_layer_ids:
|
||||
target_layers = aux_layer_ids
|
||||
else:
|
||||
target_layers = [2, target_num_layers // 2, target_num_layers - 3]
|
||||
logger.info(f"EAGLE-3: target_layers = {target_layers} (target model has {target_num_layers} layers)")
|
||||
self.gguf_writer.add_target_layers(target_layers)
|
||||
|
||||
@@ -90,6 +95,12 @@ class LlamaModel(TextModel):
|
||||
logger.info(f"EAGLE-3: norm_before_residual = {norm_before_residual}")
|
||||
self.gguf_writer.add_norm_before_residual(norm_before_residual)
|
||||
|
||||
# norm_before_fc: RMSNorm applied to the fused target features before the
|
||||
# fc projection (e.g. nvidia/gpt-oss-120b-Eagle3-v3)
|
||||
norm_before_fc = eagle3_raw_config.get("norm_before_fc", False)
|
||||
logger.info(f"EAGLE-3: norm_before_fc = {norm_before_fc}")
|
||||
self.gguf_writer.add_norm_before_fc(norm_before_fc)
|
||||
|
||||
def set_vocab(self):
|
||||
# eagle3: use tokenizer from target model if provided
|
||||
original_dir_model = None
|
||||
@@ -222,6 +233,9 @@ class LlamaModel(TextModel):
|
||||
if name == "fc.weight":
|
||||
yield (name, data_torch)
|
||||
return
|
||||
if name == "input_norm.weight":
|
||||
yield (self.format_tensor_name(gguf.MODEL_TENSOR.ENC_OUTPUT_NORM), data_torch)
|
||||
return
|
||||
if name == "d2t":
|
||||
# store for manual int64 handling in prepare_tensors (avoid F32 conversion)
|
||||
if not hasattr(self, '_eagle3_int_tensors'):
|
||||
|
||||
@@ -161,6 +161,7 @@ class Keys:
|
||||
TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size"
|
||||
BLOCK_SIZE = "{arch}.block_size"
|
||||
NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual"
|
||||
NORM_BEFORE_FC = "{arch}.norm_before_fc"
|
||||
|
||||
class Attention:
|
||||
HEAD_COUNT = "{arch}.attention.head_count"
|
||||
@@ -4343,6 +4344,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.FFN_DOWN,
|
||||
MODEL_TENSOR.FFN_UP,
|
||||
MODEL_TENSOR.FC,
|
||||
MODEL_TENSOR.ENC_OUTPUT_NORM,
|
||||
MODEL_TENSOR.D2T,
|
||||
],
|
||||
MODEL_ARCH.DFLASH: [
|
||||
|
||||
@@ -971,6 +971,9 @@ class GGUFWriter:
|
||||
def add_norm_before_residual(self, value: bool) -> None:
|
||||
self.add_bool(Keys.LLM.NORM_BEFORE_RESIDUAL.format(arch=self.arch), value)
|
||||
|
||||
def add_norm_before_fc(self, value: bool) -> None:
|
||||
self.add_bool(Keys.LLM.NORM_BEFORE_FC.format(arch=self.arch), value)
|
||||
|
||||
def add_attention_output_group_count(self, count: int) -> None:
|
||||
self.add_uint32(Keys.Attention.OUTPUT_GROUP_COUNT.format(arch=self.arch), count)
|
||||
|
||||
|
||||
@@ -317,6 +317,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
|
||||
{ LLM_KV_TARGET_LAYERS, "%s.target_layers" },
|
||||
{ LLM_KV_TARGET_HIDDEN_SIZE, "%s.target_hidden_size" },
|
||||
{ LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" },
|
||||
{ LLM_KV_NORM_BEFORE_FC, "%s.norm_before_fc" },
|
||||
|
||||
{ LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" },
|
||||
// sentence-transformers dense modules feature dims
|
||||
|
||||
@@ -363,6 +363,7 @@ enum llm_kv {
|
||||
LLM_KV_TARGET_LAYERS,
|
||||
LLM_KV_TARGET_HIDDEN_SIZE,
|
||||
LLM_KV_NORM_BEFORE_RESIDUAL,
|
||||
LLM_KV_NORM_BEFORE_FC,
|
||||
|
||||
LLM_KV_SHORTCONV_L_CACHE,
|
||||
|
||||
|
||||
@@ -47,6 +47,7 @@ struct llama_hparams {
|
||||
bool use_par_res;
|
||||
bool swin_norm;
|
||||
bool norm_before_residual = false;
|
||||
bool norm_before_fc = false;
|
||||
|
||||
uint32_t n_ctx_train; // context size the model was trained on
|
||||
uint32_t n_embd;
|
||||
|
||||
@@ -28,6 +28,10 @@ void llama_model_eagle3::load_arch_hparams(llama_model_loader & ml) {
|
||||
LLAMA_LOG_INFO("%s: EAGLE3gnorm_before_residual = true\n", __func__);
|
||||
}
|
||||
|
||||
// eagle3 norm_before_fc (optional, default false)
|
||||
// compatible with eagle3.1 (e.g. nvidia/gpt-oss-120b-Eagle3-v3)
|
||||
ml.get_key(LLM_KV_NORM_BEFORE_FC, hparams.norm_before_fc, false);
|
||||
|
||||
type = LLM_TYPE_UNKNOWN;
|
||||
}
|
||||
|
||||
@@ -53,6 +57,11 @@ void llama_model_eagle3::load_arch_tensors(llama_model_loader &) {
|
||||
// Feature fusion layer: projects 3 target layers to draft hidden size
|
||||
fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), {n_embd_inp, n_embd}, 0);
|
||||
|
||||
// RMSNorm on the fused target features (input to fc), only when norm_before_fc is set.
|
||||
if (hparams.norm_before_fc) {
|
||||
output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), {n_embd_inp}, 0);
|
||||
}
|
||||
|
||||
// Output layer (uses draft vocab size)
|
||||
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
|
||||
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_draft_vocab}, TENSOR_NOT_REQUIRED);
|
||||
@@ -130,6 +139,12 @@ llama_model_eagle3::graph<true>::graph(const llama_model & model, const llm_grap
|
||||
|
||||
cur = build_inp_embd_enc();
|
||||
|
||||
// RMSNorm on the fused target features before fc
|
||||
if (hparams.norm_before_fc) {
|
||||
cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1);
|
||||
cb(cur, "enc_input_norm", -1);
|
||||
}
|
||||
|
||||
// Feature fusion layer
|
||||
cur = build_lora_mm(model.fc, cur);
|
||||
cb(cur, "fc_out", -1);
|
||||
|
||||
@@ -116,7 +116,7 @@ llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_
|
||||
|
||||
cb(cur, "attn_out", il);
|
||||
}
|
||||
if (il == n_layer - 1) {
|
||||
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
|
||||
// skip computing output for unused tokens
|
||||
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
|
||||
@@ -154,6 +154,12 @@ llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_
|
||||
}
|
||||
cur = inpL;
|
||||
|
||||
res->t_h_nextn = cur;
|
||||
|
||||
if (!cparams.embeddings_nextn_masked && inp_out_ids) {
|
||||
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