Added support for Jais models

[grandiose-pizza]

Jais models are pretrained and fine-tuned over a curated Arabic and English text/ prompt-response pairs datasets respectively. It is trained from scratch by Core42 in partnership with MBZUAI and Cerebras systems on their Condor Galaxy. The model architecture is based on transformer-based decoder-only (GPT-3) architecture and uses SwiGLU non-linearity. It implements ALiBi position embeddings, enabling the model to extrapolate to long sequence lengths, providing improved context handling and model precision.

The research can be studied here: https://arxiv.org/pdf/2308.16149.pdf

Metrics can be found here: https://huggingface.co/core42/jais-30b-chat-v3

These are SoTA Arabic-English bilingual models. The playground can be accessed at https://arabic-gpt.ai/

[grandiose-pizza]

Not ready to be merged yet. Bugs found.Fixes to be done.

[robertgshaw2-neuralmagic]

Nice PR! Left a few cosmetic changes + some ideas for how to slightly improve the performance of the models via fusion in the MLP / a fused activation function

[7ossam81]

@robertgshaw2-neuralmagic could you please check the status of the PR ?

[esmeetu]

@grandiose-pizza Sorry for delay some days. Please merge the latest update, and you can now just read model's scale parameter from config and pass it to LogitsProcessor without custom Sampler.

[grandiose-pizza]

@grandiose-pizza Sorry for delay some days. Please merge the latest update, and you can now just read model's scale parameter from config and pass it to LogitsProcessor without custom Sampler.

Hi @esmeetu ,

I have update the code to adapt to #3233 . Tested with multi and single GPU setting. Works as expected.

However, there was a tiny miss in your PR and I have included that fix as well in this push.

I am running into an error for gpt2, I think you may have missed updating the new_tokens properly just for gpt2.py. For rest is okay. I have included this fix in this PR. Hope this is Okay:

Here is the older bug.
Your sampler forward function during a call from a model is passing three values:

vllm/vllm/model_executor/models/gpt2.py

Line 245 in 6ebd02b

next_tokens = self.sampler(self.lm_head_weight, logits,

    def sample(
        self,
        logits: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
        next_tokens = self.sampler(self.lm_head_weight, logits,
                                   sampling_metadata)
        return next_tokens

But the sampler forward function itself accepts two values:

vllm/vllm/model_executor/layers/sampler.py

Line 31 in 6ebd02b

def forward(

    def forward(
        self,
        logits: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
        assert logits is not None
        _, vocab_size = logits.shape

I am getting the following error due to this:
TypeError: forward() takes 3 positional arguments but 4 were given

[esmeetu]

Thanks for catching that! LGTM.