Argument passing for TP degree

[VenkateshPasumarti]

Feature request

I was trying for generating text-embeddings for mistral based model using sentence transformers, but facing an issue with memory since it is trying to download the complete model in one core and throwing memory constraint issues , since mistral model requires 16GB and one neuron core is of size is 16GB. So, i wanted to activate multiple cores using an argument in order to generate using optimum neuron.

Motivation

Need to activate multiple cores and also such that i can run two models in parellel using different cores

Your contribution

Was able to run smaller models, but for larger models facing issues.

[dacorvo]

@VenkateshPasumarti thank you for your feedback. You can use the num_cores argument to increase the number of cores on which the model is deployed.

[dacorvo]

https://huggingface.co/docs/optimum-neuron/guides/export_model#exporting-llms-to-neuron

[VenkateshPasumarti]

Thanks for the reply @dacorvo , can i know how can i run different models in parallel, something like locking those cores for a certain task
ex: If i take inf2.48x large i will be having multiple cores so for first 2 cores i wanted to run one sentence transformers based embeddings model and then in next multiple cores i wanted to run a re-ranking model

[dacorvo]

@VenkateshPasumarti you can restrict the number of visible cores by using environment variables, but for that each model must run in a separate process (please refer to the AWS Neuron SDK documentation to see how).
Alternatively, you can run each model in a container, mapping only some devices to each container (you can see an example of that using docker compose in the benchmark/text-generation-inference/performance subdirectory).

[VenkateshPasumarti]

Thanks for the reply and information @dacorvo