【vLLM 学习】Aqlm 示例

HyperAI超神经2025-04-22 15:46

vLLM 是一款专为大语言模型推理加速而设计的框架，实现了 KV 缓存内存几乎零浪费，解决了内存管理瓶颈问题。

更多 vLLM 中文文档及教程可访问 →https://vllm.hyper.ai/

复制代码

from vllm import LLM, SamplingParams
from vllm.utils import FlexibleArgumentParser


def main():

    parser = FlexibleArgumentParser(description='AQLM examples')

    parser.add_argument('--model',
 '-m',
 type=str,
                        default=None,
 help='model path, as for HF')
    parser.add_argument('--choice',
 '-c',
 type=int,
                        default=0,
 help='known good models by index, [0-4]')
    parser.add_argument('--tensor-parallel-size',
 '-t',
 type=int,
                        default=1,
 help='tensor parallel size')

    args = parser.parse_args()

    models = [
 "ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf",
 "ISTA-DASLab/Llama-2-7b-AQLM-2Bit-2x8-hf",
 "ISTA-DASLab/Llama-2-13b-AQLM-2Bit-1x16-hf",
 "ISTA-DASLab/Mixtral-8x7b-AQLM-2Bit-1x16-hf",
 "BlackSamorez/TinyLlama-1_1B-Chat-v1_0-AQLM-2Bit-1x16-hf",
 ]

    model = LLM(args.model if args.model is not None else models[args.choice],
                tensor_parallel_size=args.tensor_parallel_size)

    sampling_params = SamplingParams(max_tokens=100, temperature=0)
    outputs = model.generate("Hello my name is",
                             sampling_params=sampling_params)
 print(outputs[0].outputs[0].text)


if __name__ == '__main__':
    main()