Support for LLAMA models

[parth-chudasama]

Hi, Is there any plan to support LLAMA based models?

[aws-rhsoln]

We are working on adding support for LLAMA in an upcoming release. Will update once we have it.

[romanserg]

Hi, I am trying to apply the current implementation of the class LlamaForSampling. I am trying this code:

model = AutoModelForCausalLM.from_pretrained("openlm-research/open_llama_7b_700bt_preview")
save_pretrained_split(model, 'llama-split')
model_neuron = LlamaForSampling.from_pretrained('llama-split', batch_size=1, tp_degree=2, n_positions=256, amp='f32', unroll=None)
model_neuron.to_neuron()

and I am getting this error:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[22], line 1
----> 1 model_neuron.to_neuron()

File ~/anaconda3/envs/llm310/lib/python3.10/site-packages/transformers_neuronx/llama/model.py:76, in LlamaForSampling.to_neuron(self)
     73 new_layer.add_pre_mlp_layer_norm(layer.post_attention_layernorm.weight.detach(), None)
     75 # Note: Automatic MLP padding is safe since zeros are *only* introduced to intermediary state
---> 76 new_layer.add_parameter(mlp.gate_proj.weight.T, sharding=1, allow_pad=True)
     77 new_layer.add_parameter(mlp.up_proj.weight.T, sharding=1, allow_pad=True)
     78 new_layer.add_parameter(mlp.down_proj.weight.T, sharding=0, allow_pad=True)

File ~/anaconda3/envs/llm310/lib/python3.10/site-packages/torch/nn/modules/module.py:1269, in Module.__getattr__(self, name)
   1267     if name in modules:
   1268         return modules[name]
-> 1269 raise AttributeError("'{}' object has no attribute '{}'".format(
   1270     type(self).__name__, name))

AttributeError: 'DecoderLayer' object has no attribute 'add_parameter'

Do you have any ideas regarding this issue?

[AWSGH]

Llama is still under dev, please follow progress here: https://github.com/aws-neuron/transformers-neuronx/tree/main/src/transformers_neuronx/llama

[aws-donkrets]

Hi parth-chudasama, SDK releases > 2.14 offer support for llama2 models. Can you download one and let us know if that works for you. We plan to continue improving its accuracy and performance in subsequent releases.

[mrnikwaws]

Closing since LlamaV2 support has now been added