Skip to content

Commit d7da8dc

Browse files
bartowski1182ngxson
bartowski1182
and
ngxson
authored
model : Add support for Arcee AI's upcoming AFM model (#14185)
* Add Arcee AFM support * Add draft update code * Fix linter and update URL, may still not be final * Update src/llama-model.cpp Co-authored-by: Xuan-Son Nguyen <[email protected]> * Remote accidental blank line --------- Co-authored-by: Xuan-Son Nguyen <[email protected]>
1 parent cd355ed commit d7da8dc

File tree

7 files changed

+235
-1
lines changed

7 files changed

+235
-1
lines changed

‎convert_hf_to_gguf.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2020,6 +2020,20 @@ def prepare_tensors(self):
20202020
raise ValueError(f"Unprocessed experts: {experts}")
20212021

20222022

2023+
@ModelBase.register("ArceeForCausalLM")
2024+
class ArceeModel(LlamaModel):
2025+
model_arch = gguf.MODEL_ARCH.ARCEE
2026+
2027+
def set_gguf_parameters(self):
2028+
super().set_gguf_parameters()
2029+
self._try_set_pooling_type()
2030+
rope_scaling = self.hparams.get("rope_scaling") or {}
2031+
if rope_scaling.get("rope_type", rope_scaling.get("type")) == "yarn" and "factor" in rope_scaling:
2032+
self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.YARN)
2033+
self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"])
2034+
self.gguf_writer.add_rope_scaling_orig_ctx_len(rope_scaling["original_max_position_embeddings"])
2035+
2036+
20232037
@ModelBase.register(
20242038
"LlavaForConditionalGeneration", # pixtral
20252039
"Mistral3ForConditionalGeneration", # mistral small 3.1

‎convert_hf_to_gguf_update.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,7 @@ class TOKENIZER_TYPE(IntEnum):
128128
{"name": "llama4", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct", },
129129
{"name": "pixtral", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/mistral-community/pixtral-12b", },
130130
{"name": "seed-coder", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/ByteDance-Seed/Seed-Coder-8B-Base", },
131+
{"name": "arcee", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/arcee-ai/AFM-4.5B", }, # TODO confirm final URL
131132
]
132133

133134
# some models are known to be broken upstream, so we will skip them as exceptions

‎gguf-py/gguf/constants.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -344,6 +344,7 @@ class MODEL_ARCH(IntEnum):
344344
PLM = auto()
345345
BAILINGMOE = auto()
346346
DOTS1 = auto()
347+
ARCEE = auto()
347348

348349

349350
class VISION_PROJECTOR_TYPE(IntEnum):
@@ -624,7 +625,8 @@ class MODEL_TENSOR(IntEnum):
624625
MODEL_ARCH.WAVTOKENIZER_DEC: "wavtokenizer-dec",
625626
MODEL_ARCH.PLM: "plm",
626627
MODEL_ARCH.BAILINGMOE: "bailingmoe",
627-
MODEL_ARCH.DOTS1: "dots1"
628+
MODEL_ARCH.DOTS1: "dots1",
629+
MODEL_ARCH.ARCEE: "arcee",
628630
}
629631

630632
VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = {
@@ -2070,6 +2072,21 @@ class MODEL_TENSOR(IntEnum):
20702072
MODEL_TENSOR.FFN_UP_EXP,
20712073
MODEL_TENSOR.FFN_UP_SHEXP,
20722074
],
2075+
MODEL_ARCH.ARCEE: [
2076+
MODEL_TENSOR.TOKEN_EMBD,
2077+
MODEL_TENSOR.OUTPUT_NORM,
2078+
MODEL_TENSOR.OUTPUT,
2079+
MODEL_TENSOR.ROPE_FREQS,
2080+
MODEL_TENSOR.ATTN_NORM,
2081+
MODEL_TENSOR.ATTN_Q,
2082+
MODEL_TENSOR.ATTN_K,
2083+
MODEL_TENSOR.ATTN_V,
2084+
MODEL_TENSOR.ATTN_OUT,
2085+
MODEL_TENSOR.ATTN_ROT_EMBD,
2086+
MODEL_TENSOR.FFN_NORM,
2087+
MODEL_TENSOR.FFN_DOWN,
2088+
MODEL_TENSOR.FFN_UP,
2089+
],
20732090
# TODO
20742091
}
20752092

‎src/llama-arch.cpp

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
7373
{ LLM_ARCH_PLM, "plm" },
7474
{ LLM_ARCH_BAILINGMOE, "bailingmoe" },
7575
{ LLM_ARCH_DOTS1, "dots1" },
76+
{ LLM_ARCH_ARCEE, "arcee" },
7677
{ LLM_ARCH_UNKNOWN, "(unknown)" },
7778
};
7879

@@ -244,6 +245,24 @@ static const std::map<llm_arch, std::map<llm_tensor, const char *>> LLM_TENSOR_N
244245
{ LLM_TENSOR_FFN_UP_EXPS, "blk.%d.ffn_up_exps" },
245246
},
246247
},
248+
{
249+
LLM_ARCH_ARCEE,
250+
{
251+
{ LLM_TENSOR_TOKEN_EMBD, "token_embd" },
252+
{ LLM_TENSOR_OUTPUT_NORM, "output_norm" },
253+
{ LLM_TENSOR_OUTPUT, "output" },
254+
{ LLM_TENSOR_ROPE_FREQS, "rope_freqs" },
255+
{ LLM_TENSOR_ATTN_NORM, "blk.%d.attn_norm" },
256+
{ LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" },
257+
{ LLM_TENSOR_ATTN_K, "blk.%d.attn_k" },
258+
{ LLM_TENSOR_ATTN_V, "blk.%d.attn_v" },
259+
{ LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" },
260+
{ LLM_TENSOR_ATTN_ROT_EMBD, "blk.%d.attn_rot_embd" },
261+
{ LLM_TENSOR_FFN_NORM, "blk.%d.ffn_norm" },
262+
{ LLM_TENSOR_FFN_DOWN, "blk.%d.ffn_down" },
263+
{ LLM_TENSOR_FFN_UP, "blk.%d.ffn_up" },
264+
},
265+
},
247266
{
248267
LLM_ARCH_LLAMA4,
249268
{

‎src/llama-arch.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ enum llm_arch {
7777
LLM_ARCH_PLM,
7878
LLM_ARCH_BAILINGMOE,
7979
LLM_ARCH_DOTS1,
80+
LLM_ARCH_ARCEE,
8081
LLM_ARCH_UNKNOWN,
8182
};
8283

‎src/llama-model.cpp

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -599,6 +599,16 @@ void llama_model::load_hparams(llama_model_loader & ml) {
599599
hparams.use_kq_norm = false;
600600
}
601601
} break;
602+
case LLM_ARCH_ARCEE:
603+
{
604+
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
605+
606+
// Arcee uses the same structure as Llama
607+
switch (hparams.n_layer) {
608+
case 36: type = LLM_TYPE_4B; break;
609+
default: type = LLM_TYPE_UNKNOWN;
610+
}
611+
} break;
602612
case LLM_ARCH_DECI:
603613
{
604614
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
@@ -4190,6 +4200,37 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
41904200
}
41914201
}
41924202
} break;
4203+
case LLM_ARCH_ARCEE:
4204+
{
4205+
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
4206+
4207+
// output
4208+
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
4209+
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
4210+
4211+
// if output is NULL, init from the input tok embed
4212+
if (output == NULL) {
4213+
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
4214+
}
4215+
4216+
for (int i = 0; i < n_layer; ++i) {
4217+
auto & layer = layers[i];
4218+
4219+
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
4220+
4221+
layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head}, 0);
4222+
layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}, 0);
4223+
layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}, 0);
4224+
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
4225+
4226+
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
4227+
4228+
layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
4229+
4230+
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);
4231+
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
4232+
}
4233+
} break;
41934234
default:
41944235
throw std::runtime_error("unknown architecture");
41954236
}
@@ -13411,6 +13452,141 @@ struct llm_build_dots1 : public llm_graph_context {
1341113452
}
1341213453
};
1341313454

13455+
struct llm_build_arcee : public llm_graph_context {
13456+
llm_build_arcee(const llama_model & model, const llm_graph_params & params, ggml_cgraph * gf) : llm_graph_context(params) {
13457+
const int64_t n_embd_head = hparams.n_embd_head_v;
13458+
13459+
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
13460+
GGML_ASSERT(n_embd_head == hparams.n_rot);
13461+
13462+
ggml_tensor * cur;
13463+
ggml_tensor * inpL;
13464+
13465+
inpL = build_inp_embd(model.tok_embd);
13466+
13467+
// inp_pos - contains the positions
13468+
ggml_tensor * inp_pos = build_inp_pos();
13469+
13470+
auto * inp_attn = build_attn_inp_kv_unified();
13471+
13472+
const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
13473+
13474+
for (int il = 0; il < n_layer; ++il) {
13475+
ggml_tensor * inpSA = inpL;
13476+
13477+
// norm
13478+
cur = build_norm(inpL,
13479+
model.layers[il].attn_norm, NULL,
13480+
LLM_NORM_RMS, il);
13481+
cb(cur, "attn_norm", il);
13482+
13483+
// self-attention
13484+
{
13485+
// rope freq factors for llama3; may return nullptr for llama2 and other models
13486+
ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
13487+
13488+
// compute Q and K and RoPE them
13489+
ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur);
13490+
cb(Qcur, "Qcur", il);
13491+
if (model.layers[il].bq) {
13492+
Qcur = ggml_add(ctx0, Qcur, model.layers[il].bq);
13493+
cb(Qcur, "Qcur", il);
13494+
}
13495+
13496+
ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur);
13497+
cb(Kcur, "Kcur", il);
13498+
if (model.layers[il].bk) {
13499+
Kcur = ggml_add(ctx0, Kcur, model.layers[il].bk);
13500+
cb(Kcur, "Kcur", il);
13501+
}
13502+
13503+
ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur);
13504+
cb(Vcur, "Vcur", il);
13505+
if (model.layers[il].bv) {
13506+
Vcur = ggml_add(ctx0, Vcur, model.layers[il].bv);
13507+
cb(Vcur, "Vcur", il);
13508+
}
13509+
13510+
Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
13511+
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
13512+
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
13513+
13514+
Qcur = ggml_rope_ext(
13515+
ctx0, Qcur, inp_pos, rope_factors,
13516+
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
13517+
ext_factor, attn_factor, beta_fast, beta_slow
13518+
);
13519+
13520+
Kcur = ggml_rope_ext(
13521+
ctx0, Kcur, inp_pos, rope_factors,
13522+
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
13523+
ext_factor, attn_factor, beta_fast, beta_slow
13524+
);
13525+
13526+
cb(Qcur, "Qcur", il);
13527+
cb(Kcur, "Kcur", il);
13528+
cb(Vcur, "Vcur", il);
13529+
13530+
cur = build_attn(inp_attn, gf,
13531+
model.layers[il].wo, model.layers[il].bo,
13532+
Qcur, Kcur, Vcur, nullptr, nullptr, kq_scale, il);
13533+
cb(cur, "attn_out", il);
13534+
}
13535+
13536+
if (il == n_layer - 1) {
13537+
// skip computing output for unused tokens
13538+
ggml_tensor * inp_out_ids = build_inp_out_ids();
13539+
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
13540+
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
13541+
}
13542+
13543+
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
13544+
cb(ffn_inp, "ffn_inp", il);
13545+
13546+
// feed-forward network
13547+
// ARCEE uses relu^2 instead of silu
13548+
cur = build_norm(ffn_inp,
13549+
model.layers[il].ffn_norm, NULL,
13550+
LLM_NORM_RMS, il);
13551+
cb(cur, "ffn_norm", il);
13552+
13553+
cur = build_ffn(cur,
13554+
model.layers[il].ffn_up, NULL, NULL,
13555+
NULL, NULL, NULL,
13556+
model.layers[il].ffn_down, NULL, NULL,
13557+
NULL,
13558+
LLM_FFN_RELU_SQR, LLM_FFN_SEQ, il);
13559+
cb(cur, "ffn_out", il);
13560+
13561+
cur = ggml_add(ctx0, cur, ffn_inp);
13562+
cb(cur, "ffn_out", il);
13563+
13564+
cur = build_cvec(cur, il);
13565+
cb(cur, "l_out", il);
13566+
13567+
// input for next layer
13568+
inpL = cur;
13569+
}
13570+
13571+
cur = inpL;
13572+
13573+
cur = build_norm(cur,
13574+
model.output_norm, NULL,
13575+
LLM_NORM_RMS, -1);
13576+
13577+
cb(cur, "result_norm", -1);
13578+
res->t_embd = cur;
13579+
13580+
// lm_head
13581+
cur = build_lora_mm(model.output, cur);
13582+
13583+
cb(cur, "result_output", -1);
13584+
res->t_logits = cur;
13585+
13586+
ggml_build_forward_expand(gf, cur);
13587+
}
13588+
};
13589+
1341413590
llama_memory_i * llama_model::create_memory(const llama_memory_params & params, llama_cparams & cparams) const {
1341513591
llama_memory_i * res;
1341613592

@@ -13753,6 +13929,10 @@ llm_graph_result_ptr llama_model::build_graph(
1375313929
{
1375413930
llm = std::make_unique<llm_build_dots1>(*this, params, gf);
1375513931
} break;
13932+
case LLM_ARCH_ARCEE:
13933+
{
13934+
llm = std::make_unique<llm_build_arcee>(*this, params, gf);
13935+
} break;
1375613936
default:
1375713937
GGML_ABORT("fatal error");
1375813938
}
@@ -13902,6 +14082,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
1390214082
case LLM_ARCH_GRANITE_MOE:
1390314083
case LLM_ARCH_CHAMELEON:
1390414084
case LLM_ARCH_BAILINGMOE:
14085+
case LLM_ARCH_ARCEE:
1390514086
return LLAMA_ROPE_TYPE_NORM;
1390614087

1390714088
// the pairs of head values are offset by n_rot/2

‎src/llama-vocab.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1987,6 +1987,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
19871987
|| t.first == "<|eom_id|>"
19881988
|| t.first == "<EOT>"
19891989
|| t.first == "_<EOT>"
1990+
|| t.first == "<|end_of_text|>"
19901991
) {
19911992
special_eog_ids.insert(t.second);
19921993
if ((id_to_token[t.second].attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {

0 commit comments

Comments
 (0)