llama : fix embeddings

[ggerganov]

ref #5655, #5783

• llama_batch.logits now also indicates if embeddings are output for that token
• fix llama_get_embeddings_ith() to return token embeddings
• add llama_get_embeddings_seq() to return sequence embeddings
• fix embedding example to work both with BERT and non-BERT models
• fix server to get the resulting embeddings correctly
• add examples/server-embd.py helper script
• server supports --pooling

TODO:

• rename llama_batch.logits to llama_batch.output (future PR)
• BERT graph does not need KV cache - remove it
• server should not queue partial prompts when computing embeddings (future PR)

[cebtenzzre]

While we're modifying embeddings.cpp - shouldn't this:

struct llama_batch batch = llama_batch_init(n_batch, 0, n_prompts);

be this instead:

struct llama_batch batch = llama_batch_init(n_batch, 0, 1);

Because only one sequence ID is assigned per token?

[tybalex]

should the embeddings from server.cpp get normalized too? embedding.cpp does that.

[iamlemec]

Regarding removing the KV cache, I think this will give a big speedup. I did a flamegraph on CUDA, and it was spending fully 50% of the time in calls to llama_kv_cell.has_seq_id for constructing the attention mask in llama_set_inputs. If you know everything is within batch though, you can construct the attention matrix from just the batch object.

[ggerganov]

@iamlemec What is the meaning of CLS in this context? I don't associate this abbreviation with anything

[cebtenzzre]

I believe this mode uses the embedding of the CLS token (which is the Bert equivalent of BOS) instead of averaging the embedding of all tokens.

[iamlemec]

Yeah, always confused me since I think CLoSe, but I guess it stands for (sentence) classification. There's something to be said for replacing it (and the command line flag) with something more expressive like START.

[cebtenzzre]

With these changes I'm getting an MSE of 1.3e-3 relative to Sentence Transformers on WikiText with nomic-embed-text-v1.f16.gguf instead of the previous 5.62e-10. Not sure what the problem is.

edit: I'm also still seeing NaNs out of the embedding example.

[ggerganov]

With these changes I'm getting an MSE of 1.3e-3 relative to Sentence Transformers on WikiText with nomic-embed-text-v1.f16.gguf instead of the previous 5.62e-10. Not sure what the problem is.

Does it use mean pooling? I think I got it wrong again - checking

[cebtenzzre]

Does it use mean pooling?

Yes.

[iamlemec]

We should have the LLAMA_POOLING_TYPE_MEAN case join the LLAMA_POOLING_TYPE_CLS case due to the output order of the averaging matrix.

[ggerganov]

Yes. I keep getting confused with the sequence-based instead of token-based embedding extraction.

Will try to modify the API to make things more clear

[cebtenzzre]

Cool, this change fixes the NaNs from embedding.cpp, and results in the MSE of nomic-embed-text-v1.f16.gguf actually being lower than before (4.71e-10 vs 5.62e-10). Also fp32 is down from 9.34e-11 to 1.18e-14.

[ggerganov]

Also fp32 is down from 9.34e-11 to 1.18e-14.

This is likely due to no longer going through the KV cache.

Do you have any performance benchmarks to see if this change improved the speed?

[ggerganov]

So I've updated the API to support both sequence and token embeddings:

llama.cpp/llama.h

Lines 657 to 671 in 79e4eed

	// Get all output token embeddings
	// shape: [n_tokens*n_embd] (1-dimensional)
	LLAMA_API float * llama_get_embeddings(struct llama_context * ctx);
	// Get the embeddings for the ith token
	// llama_get_embeddings(ctx) + i*n_embd
	// shape: [n_embd] (1-dimensional)
	LLAMA_API float * llama_get_embeddings_ith(struct llama_context * ctx, int32_t i);
	// Get the embeddings for a sequence id
	// Returns NULL if pooling_type is LLAMA_POOLING_TYPE_NONE
	// shape: [n_embd] (1-dimensional)
	LLAMA_API float * llama_get_embeddings_seq(struct llama_context * ctx, llama_seq_id seq_id);

The sequence embeddings (i.e. llama_embeddings_seq()) work only when pooling_type != LLAMA_POOLING_TYPE_NONE, while token embeddings are populated when pooling_type == LLAMA_POOLING_TYPE_NONE. In the future, we can think about adding pooling support for all models (not just BERT) - not sure if it would be useful, but it fits into the implementation logic

With this change, to make the embedding example work with both token and sequence emebeddings (i.e. with and without pooling), we do the following:

llama.cpp/examples/embedding/embedding.cpp

Lines 52 to 65 in fc9af15

	// try to get sequence embeddings - supported only when pooling_type is not NONE
	const float * embd = llama_get_embeddings_seq(ctx, batch.seq_id[i][0]);
	if (embd == NULL) {
	embd = llama_get_embeddings_ith(ctx, i);
	if (embd == NULL) {
	fprintf(stderr, "%s: failed to get embeddings for token %d\n", __func__, i);
	continue;
	}
	}
	float * out = output + batch.seq_id[i][0] * n_embd;
	normalize(embd, out, n_embd);
	}

First try to get sequence embeddings. This would return NULL if pooling is disabled and we fallback to token embeddings

Hope this finally works. Let me know if you give it a try

[iamlemec]

Working great here! Numbers look very, very similar to earlier.

I'm a little surprised to see a slight performance degredation (around 20-30% on both CPU and CUDA). We're definitely getting gains from the kv-cache attention matrix stuff when that's a factor. But I wonder if it relates to the fact that we're breaking up the final copy into sequences? Either that or one of the smaller changes in built_bert. Looking into it now.

[cebtenzzre]

Does the llama_kv_cache_clear still do anything useful?

edit: I remembered that this example is used for models with causal attention as well. I won't need the equivalent if I'm just working with embedding models.

[ggerganov]

Yes, llama_kv_cache_clear is not necessary for embedding models, but makes sense for causal attention models

[ggerganov]

I'm a little surprised to see a slight performance degredation (around 20-30% on both CPU and CUDA). We're definitely getting gains from the kv-cache attention matrix stuff when that's a factor. But I wonder if it relates to the fact that we're breaking up the final copy into sequences? Either that or one of the smaller changes in built_bert. Looking into it now.

Let me know what you find. I guess the embeddings extraction is now indeed a bit slower due to the std::map and loop over the batch. The actual computation seems to be the same speed as before (at least with Metal that is). You can bench it with llama-bench since it does not involve any embeddings extraction.

The main struggle with the API was to find a way to differentiate between token embeddings and sequence embeddings. Without the std::map there was an ambiguity in which embeddings you were extracting. There might be better way to do it though

The big benefit of not using the KV cache is that now we don't need to allocate a large KV cache memory buffer. For example, you can compute 32k token embeddings with just a large enough batch size (e.g. -c 512 -b 32768)

[iamlemec]

Oh nice, I didn't know about llama-bench. So for CPU I'm seeing ~20% performance drop, not too bad. But for CUDA I'm getting an almost 60% drop:

model	backend	ngl	test	(pre) t/s	(post) t/s
bert 109M F16	CUDA	99	pp 512	177281.98 ± 116298.41	74489.67 ± 19758.33
bert 109M F16	CUDA	0	pp 512	6824.48 ± 1151.24	5556.47 ± 536.14

It seems like the kv cache code relies more on ggml_view_3d instead of contiguous permutes, but I haven't yet found anything that yields the old performance.

[ggerganov]

Try #5891 and see if it restores the performance