I am performing a training run and use Flash attention 2 with the llama model.

Unfortunately, I hit this error

 LlavaLlamaForCausalLM.__init__() got an unexpected keyword argument 'attn_implementation'

I have followed advice from here https://github.com/haotian-liu/LLaVA/issues/1103

and still get the error on 4.38.0,

Here is how i set it up.

model = LlamaForCausalLM( config = hparams.model_config, attn_implementation="flash_attention_2" )

I would appreciate any pointers