llama: Attempt to add ModernBert

[huydt84]

I don't know whether my implementation is correct or not

[huydt84]

hparams.set_swa_pattern can't work properly with ModernBert

[huydt84]

The embedding result seems random and very low. There is something wrong with this

[CISC]

Delete the files you added in models, we don't need them, just make sure test-tokenizer-0 succeeds with the GGUF.

[ggerganov]

This should probably become:

Suggested change

	auto * inp_attn = build_attn_inp_kv_unified_iswa();
	auto * inp_attn = build_attn_inp_no_cache_iswa();

And add the corresponding mask logic in llama-graph. Special attention should be taken about how the SWA works for this model - i.e. is it symmetric or not:

# non-symmetric
token i attends to [i - n_swa, i]

# symmetric:
token i attends to [i - n_swa/2, i + n_swa/2]

[CISC]

@huydt-bti Hey, is the issue that you forgot to make this function so that swa is actually never applied?

[huydt84]

@CISC No, I made it here:

llama.cpp/src/llama-graph.cpp

Lines 349 to 352 in 454d7b7

	if (hparams.use_alibi &&
	(hparams.n_swa == 0 || (pos_diff >= -half_n_swa && pos_diff <= half_n_swa))) {
	f = -std::abs(ubatch->pos[ti] - ubatch->pos[tj]);
	} else {

I use build_attn_inp_no_cache()
#14014 (review)

[CISC]

Yes, but there's no kq_mask_swa, so is this even executed?

[huydt84]

@CISC I have just implemented it. Please check again

[ggerganov]

You have to add the new arch here:

llama.cpp/src/llama-model.cpp

Lines 13195 to 13203 in 5a8ae30

	switch (arch) {
	case LLM_ARCH_BERT:
	case LLM_ARCH_JINA_BERT_V2:
	case LLM_ARCH_NOMIC_BERT:
	case LLM_ARCH_NOMIC_BERT_MOE:
	case LLM_ARCH_WAVTOKENIZER_DEC:
	{
	res = nullptr;
	} break;

To avoid creating a memory module (a.k.a. KV cache) for these models.

[CISC]

So, since vocab is BPE you need to add modern-bert vocab handling a few places:

llama.cpp/src/llama-vocab.cpp

Line 1557 in 9f47fa5

tokenizer_pre == "roberta-bpe") {

Set correct attribute on [MASK] token, similarly to this:

llama.cpp/src/llama-vocab.cpp

Lines 2097 to 2105 in 9f47fa5

	if (false
	|| _contains_any(tokenizer_pre, {"jina-v2-de", "jina-v2-es", "jina-v2-code"})
	|| _contains_any(general_arch, {"nomic-bert-moe"})
	) {
	if (token_to_id.count("<mask>") == 0) {
	LLAMA_LOG_WARN("%s: Mask token is missing in vocab, please reconvert model!\n", __func__);
	} else {
	_set_token_attr("<mask>", LLAMA_TOKEN_ATTR_LSTRIP, true);
	}

[CISC]

The embedding result seems random and very low. There is something wrong with this

Yep, I also noticed the same with jina-reranker-v2, most likely the same issue, will investigate.

[huydt84]

We have cls, but not cls_b, so this has to be modified to handle that...

@CISC After fixing that, the result is much better now :) but it is still lower than my expectations about ModernBert. Maybe there is problem somewhere else...

[CISC]

Everything else LGTM, so pending finding the output issue.

[huydt84]

@CISC I see part of the problem! I'm masking the token inside the window, which should be outside

[huydt84]

No, the function isn't used because it belongs to llama_kv_cache_unified

[CISC]

No, this is unnecessary duplication, it should be handled like this once you add llm_graph_input/build_attn_inp_no_cache_iswa:

llama.cpp/src/llama-graph.cpp

Lines 356 to 370 in d7da8dc

	void llm_graph_input_attn_kv_unified::set_input(const llama_ubatch * ubatch) {
	if (self_kq_mask) {
	kv_state->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
	}
	}
	void llm_graph_input_attn_kv_unified_iswa::set_input(const llama_ubatch * ubatch) {
	if (self_kq_mask) {
	kv_state->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
	}
	if (self_kq_mask_swa) {
	kv_state->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);
	}
	}

[huydt84]

It is not interleaved swa, so I still prefer using build_attn_inp_no_cache. I will try to refactor llm_graph_input_attn_no_cache::set_input, but the mechanisim is different from llm_graph_input_attn_kv_unified and llm_graph_input_attn_kv_unified_iswa, since they use kv_state

[CISC]

Sure, the point is just that you can handle it similarly.

[huydt84]

I just pushed my code - but it's ugly to copy is_masked_swa and place it inside llm_graph_input_attn_no_cache::set_input. Do you have any suggestions?

[CISC]

I don't understand why you can't split the methods just like in unified, you can have a is_masked_swa implementation for no_cache and a much cleaner set_input.

[huydt84]

Personally I don't think that would be cleaner, since the new build_attn_inp_no_cache_swa will almost be the same as the current build_attn_inp_no_cache, and we have a new additional build_attn_inp_no_cache. But I will try that

[CISC]

That doesn't matter, build_attn_inp_* are tiny, the important part is that you can reuse the same set_input_kq_mask.

[CISC]

@huydt84 How is progress on this? Would be very nice to bring it to completion soon, unfortunately there have been a lot of internal changes lately, hope resolving conflicts won't be too much of a chore...

[huydt84]

@CISC my desparate attempts "to make it work" led to no results but my local branches becoming a mess now :)) I'm sorry but currenty I don't work on this anymore, so it's really helpful if other people can add support for this model, whether using the code from this PR or not (that's the reason why I still open this PR)

[CISC]

#15641