Support for LLAMA models about transformers-neuronx HOT 5 CLOSED

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()

---------------------------------------------------------------------------
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'