Add GPTQ support for Gemma

[TechxGenus]

The latest version of AutoGPTQ already supports Gemma. Add some minor modifications can provide support for vllm.

My test code is as follow:

from vllm import LLM, SamplingParams

sampling_params = SamplingParams(temperature=0, max_tokens=256)
llm = LLM(model="TechxGenus/gemma-2b-GPTQ")

outputs = llm.generate(["def min(arr):\n"], sampling_params)
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(prompt + generated_text)

The output result is:

def min(arr):
    min = arr[0]
    for i in range(1, len(arr)):
        if arr[i] < min:
            min = arr[i]
    return min


def max(arr):
    max = arr[0]
    for i in range(1, len(arr)):
        if arr[i] > max:
            max = arr[i]
    return max


def main():
    arr = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    print(min(arr))
    print(max(arr))


if __name__ == "__main__":
    main()

[simon-mo]

thank you for the PR. can you confirm the original model still work?

[TechxGenus]

Sure! The test code is as follows:

from vllm import LLM, SamplingParams

sampling_params = SamplingParams(temperature=0, max_tokens=256)
llm = LLM(model="google/gemma-2b")

outputs = llm.generate(["def min(arr):\n"], sampling_params)
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(prompt + generated_text)

The output result is:

def min(arr):
    min = arr[0]
    for i in range(1, len(arr)):
        if arr[i] < min:
            min = arr[i]
    return min

def max(arr):
    max = arr[0]
    for i in range(1, len(arr)):
        if arr[i] > max:
            max = arr[i]
    return max

def main():
    arr = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    print(min(arr))
    print(max(arr))

if __name__ == '__main__':
    main()

[esmeetu]

Quick question: Why this check is needed?

[TechxGenus]

Gemma's input and output weights are shared and there is no separate lm_head parameter.
There will be no problem when loading the original model, but when loading the quantized model, lm_head is not included in params_dict. I skip it to load the model normally.

[esmeetu]

But gptq output is right when removing this check.🤔

[TechxGenus]

After reinstalling from the latest main branch of transformers and AutoGPTQ, I found that removing them works fine. Thanks for testing!

[esmeetu]

Could you merge this check with the one above and place it outside the for loop?

[TechxGenus]

I copy them from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/llama.py#L377-L379 and https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/llama.py#L385-L387.
Modifying the for-else structure may cause excessive modifications. I recommend reusing existing code.

[esmeetu]

@TechxGenus I think we can do this optimization on Gemma model file. Just put this check onto for (param_name, shard_name, shard_id) in stacked_params_mapping: without changing for-else code part.

[TechxGenus]

@TechxGenus I think we can do this optimization on Gemma model file. Just put this check onto for (param_name, shard_name, shard_id) in stacked_params_mapping: without changing for-else code part.

The quantized versions of gemma-2b and gemma-7b work fine, but will cause problems for future models with attention_bias set to true.

[WoosukKwon]

QQ: Why is this needed?

[TechxGenus]

Gemma's input and output weights are shared and there is no separate lm_head parameter.
There will be no problem when loading the original model, but when loading the quantized model, lm_head is not included in params_dict. I skip it to load the model normally.

[TechxGenus]

After reinstalling from the latest main branch of transformers and AutoGPTQ, I found that removing them works fine. Thanks for testing!

[esmeetu]

Thanks for your contribution!