search

huggingface / optimum-neuron

Easy, fast and very cheap training and inference on AWS Trainium and Inferentia chips.
Apache License 2.0
176 stars 53 forks source link

Add support for phi-2 #508

Closed 5cp closed 3 months ago

5cp commented 4 months ago

Feature request

phi-2 is one of the most popular models and should be supported by Optimum Neuron. I was able to trace the model using torch_neuronx.trace() using a slightly modified version of the BERT tutorial and did not encounter any errors:

2024-03-07T05:44:03Z Compiler status PASS

CPU paraphrase logits: [[-1.0661223 0.8461684]]
Neuron paraphrase logits: [[-1.0701762 0.8514435]]
CPU not-paraphrase logits: [[-1.111944 0.84410715]]
Neuron not-paraphrase logits: [[-1.1247162 0.8532015]]

Motivation

phi-2 was identified as one of the top models. It has also been requested by end users.

Your contribution

Confirmed that model can be traced by Neuron as-is.

JingyaHuang commented 4 months ago

Hi @5cp, I will open a PR to add support for phi family models on the following tasks: feature-extraction, text-classification, token-classification.

For the causal task, I wonder if phi is supported by transformers_neuronx @dacorvo.

dacorvo commented 3 months ago

@JingyaHuang no, it is not supported.