llava support about transformers-neuronx HOT 3 OPEN

Thank you for reaching out about using Neuron for LLaVa.

We recommend that you enable this model using neuronx-distributed (NxD). NxD is PyTorch based library intended to provide developers with the ability to enable model parallelism/sharding techniques themselves. In order to apply tensor parallelism to a model, the weights and compute must be split across NeuronCores. Intermediary calculations are synchronized across NeuronCores using collective operations. Typically this is done by applying different parallel layers to each portion of the network in the form of RowParallelLinear, ColumnParallelLinear and other layers (See: Parallel Layers Documentation):

• Embedding (example)
• Attention (example)
• MLP (example)
• Language model head (example)

Please check the Llama example inference model for a reference on applying tensor parallelism to a model using NxD. You are welcome to submit a PR with your contribution when you implement the model.

from transformers-neuronx.

Hi @hannanjgaws

Thanks for the response 👍 just a few comments:

Seems that the llava model is a combination of a LLMmodel + visual model + intermediate module (like any of these models).

In the case of llava-1.6-mistral-* as it name says, it uses mistral model, and seeing the implementation and the config, it also uses "clip_model" for the vision encoder part.

So a pair of questions:

• The approach to compile it, for Inf2 instances, should be to replace the inner layers with the corresponding parrallel layers before doing a trace with a sample input, is'nt it?

• After it compiled, when loading the .pt file, should I use model.load_state_dict(hf_model_state_dict) to load the weights isn't it?

Another interesting approach could be to have the "clip_model" model in transformers-neuronx and to load it in conjuntion of https://github.com/aws-neuron/transformers-neuronx/tree/main/src/transformers_neuronx/mistral + the intermediate module

from transformers-neuronx.

Hi, I would like to follow-up on this, do we have any pointers on how to compile and inference multi-modal like this on Inf2?

from transformers-neuronx.