Potential bugs for the llama inference example - Githubissues

aws-neuron / neuronx-distributed

MIT No Attribution

47 stars 9 forks source link

Potential bugs for the llama inference example #30

Open yinsong1986 opened 2 months ago

yinsong1986 commented 2 months ago

Hi team,

Sharing some observation for potential bugs as below:

https://github.com/aws-neuron/neuronx-distributed/blob/4f954715f39b2cc9e628ded79274957401bea086/examples/inference/llama2/neuron_modeling_llama.py#L283-L312 this seems to be duplicated. The first class definition should be changed to ResBlock
https://github.com/aws-neuron/neuronx-distributed/blob/main/examples/inference/modules/model_base.py#L453-L486 these lines seems to be not relevant, and should be deleted and replaced with return [res] + updated_kv_cache

Thank you!

jyang-aws commented 2 months ago

@yinsong1986 thanks for reporting! Our inference team will take a look.