[Question] Why llama model doesn't support input_emmbedings, what's the consideration here?

[sunfvrise]

❓ General Questions

In llama.py, Model LlamaForCausalLM#forward() doesn't support input_emmbedings, what's the consideration here?

[yzh119]

The embedding exists in LlamaModel class:

mlc-llm/mlc_llm/relax_model/llama.py

Lines 436 to 438 in e366975

	self.embed_tokens = Embedding(
	config.vocab_size, config.hidden_size, dtype=config.dtype
	)

[sunfvrise]

yes,i see. And actually many llm model can accept both input_ids & input_emmbedings. So that developers can customize the input emmbeding by prompt tuning, which can enhance the performance of the model under specific domain.
In current code, if i want to do some customization in the emmbedings, i must modify the self.embed_tokens which type is relax.Var. Can you give some suggestions here?

[yzh119]

Maybe related to #408 ?

[Kathryn-cat]

Hi, now llama model supports input embeddings in the prefilling stage. Currently we have exposed the API PrefillWithEmbedStep() in C++ and prefill_with_embed() in Python which accepts input embeddings instead of input messages.