FlashLlamaForCausalLM's using name dense for its mlp submodule causes error when using LoRA adapter

[sadra-barikbin]

Hi there! 🤗

FlashLlamaForCausalLM uses name dense for its MLP submodule and when user wants to employ a LoRA adapter, get_mlp_weights skips this submodule.

text-generation-inference/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py

Line 440 in 6e32205

self.dense = LlamaMLP(

text-generation-inference/server/text_generation_server/utils/adapter.py

Lines 259 to 261 in 6e32205

	def get_mlp_weights(i, layer):
	weights = {}
	if hasattr(layer, "mlp"):

This causes error:

[rank0]: KeyError: (0, 'gate_proj')

This is not the case for FlashGemma2ForCausalLM, for example, and it works properly. When I renamed dense to mlp , llama worked as well.

[koutarou-n]

I faced the same error.
It succesfully works on release 2.3.0, but fails on 2.3.1, 2.4.0.