ollama 支持并发之后和 vllm 相比性能如何？我们测测看

arkohut

已于 2024-05-22 14:17:01 修改

阅读量2.1w

点赞数 50

文章标签： python llama 人工智能神经网络深度学习

于 2024-05-20 22:34:47 首次发布

本文链接：https://blog.csdn.net/arkohut/article/details/139076652

版权

这是 ollama vs vllm - 开启并发之后的 ollama 和 vllm 相比怎么样？的笔记，需要完成的记录请观看视频。

上次我介绍了 ollama 支持了并发请求，使得其可以更好的利用 GPU 资源获取更大的吞吐量，也就是可以在单位时间内可以生成更多的 token 了。不过并发请求也会导致单个请求的响应时间变长。然后我在评论区看到有小伙伴提问，ollama 支持并发之后和 vllm 比效果如何呢。这里呢我就做个测试，对比下同样的模型分别采用 ollama 和 vllm 进行推理，看看他们的性能如何。

环境准备

模型我还是和上次一致，采用 llama3 8b 的模型。

虽然 vllm 已经在很久以前就支持 awq 和 gptq 的量化了，但是考虑到 ollama 那边所使用的 gguf 格式模型的量化方式和 vllm 不太一样，为了尽量控制实验变量，我这里就直接使用没有做过量化的 float16 的模型进行测试。

由于我自己没有显卡，这里我就在 autodl 上搞了一个云服务器，里面有一个 4090 用来测试。autodl 启动之后就是 jupyter 的环境，下面介绍的也都是基于 autodl 的 jupyter 进行的操作。

ollama 环境准备

ollama 本身就是一个二进制，我这里就直接从 github 下载好二进制文件之后拖拽上传到了 jupyter 里面。然后使用下面的命令进行安装：

install ollama /usr/local/bin/ollama

之后，通过下面的命令启动 ollama

OLLAMA_NUM_PARALLEL=16 OLLAMA_MODELS=/root/autodl-tmp/models ollama serve

其中 OLLAMA_NUM_PARALLEL=16 就是设置最大并发数为 16，而 OLLAMA_MODELS=/root/autodl-tmp/models 则是修改下载模型的位置到 /root/autodl-tmp/models 也就是 autodl 所提供的本地临时目录，默认有 50GB 的免费额度。

从 ollama 下载模型这个就非常简单了，我可以通过一行命令搞定：

ollama pull llama3:8b-instruct-fp16

可以看到这里就是下载了一个 fp16 版本的模型。

vllm 环境准备

然后准备 vllm 的环境，考虑到这里的 pytorch 版本稍微有点老，不能适配最新的 vllm。我就选择了稍微老一点的版本 0.3.3，当然我也有测试过最新的 vllm ，对于这个实验来说结果不会有显著的区别。

pip install vllm==0.3.3

下载模型这里有点麻烦，因为 huggingface 这里是需要审核通过才能下载模型的，虽然我这里审核通过了，但也是等了很久。好在模型权重在 modelscope 也能下得到。我还是依照官方文档对模型进行下载。

from modelscope.hub.snapshot_download import snapshot_download

model_dir = snapshot_download('LLM-Research/Meta-Llama-3-8B-Instruct', cache_dir='autodl-tmp', revision='master')

这里的我也是把模型下载到了 autodl-tmp 目录里。

然后这里我准备好了启动 vllm 服务的命令：

python -m vllm.entrypoints.openai.api_server \
    --model /root/autodl-tmp/LLM-Research/Meta-Llama-3-8B-Instruct \
    --served-model-name llama3:8b-instruct-fp16 \
    --trust-remote-code \
    --max-model-len 4096 \
    --port 11434

注意 served-model-name 这里我设置成 llama3:8b-instruct-fp16 并且将端口设置为 11434 都是为了和 ollama 那边保持一致，便于测试而已。

脚本准备

import aiohttp
import asyncio
import time
from tqdm import tqdm

async def fetch(session, url):
    """
    参数:
        session (aiohttp.ClientSession): 用于请求的会话。
        url (str): 要发送请求的 URL。
    
    返回:
        tuple: 包含完成 token 数量和请求时间。
    """
    start_time = time.time()

    # 固定请求的内容
    json_payload = {
        "model": "llama3:8b-instruct-fp16",
        "messages": [{"role": "user", "content": "Why is the sky blue?"}],
        "stream": False,
        "temperature": 0.7 # 参数使用 0.7 保证每次的结果略有区别
    }
    async with session.post(url, json=json_payload) as response:
        response_json = await response.json()
        end_time = time.time()
        request_time = end_time - start_time
        completion_tokens = response_json['usage']['completion_tokens'] # 从返回的参数里获取生成的 token 的数量
        return completion_tokens, request_time

async def bound_fetch(sem, session, url, pbar):
    # 使用信号量 sem 来限制并发请求的数量，确保不会超过最大并发请求数
    async with sem:
        result = await fetch(session, url)
        pbar.update(1)
        return result

async def run(load_url, max_concurrent_requests, total_requests):
    """
    通过发送多个并发请求来运行基准测试。
    
    参数:
        load_url (str): 要发送请求的URL。
        max_concurrent_requests (int): 最大并发请求数。
        total_requests (int): 要发送的总请求数。
    
    返回:
        tuple: 包含完成 token 总数列表和响应时间列表。
    """
    # 创建 Semaphore 来限制并发请求的数量
    sem = asyncio.Semaphore(max_concurrent_requests)
    
    # 创建一个异步的HTTP会话
    async with aiohttp.ClientSession() as session:
        tasks = []
        
        # 创建一个进度条来可视化请求的进度
        with tqdm(total=total_requests) as pbar:
            # 循环创建任务，直到达到总请求数
            for _ in range(total_requests):
                # 为每个请求创建一个任务，确保它遵守信号量的限制
                task = asyncio.ensure_future(bound_fetch(sem, session, load_url, pbar))
                tasks.append(task)  # 将任务添加到任务列表中
            
            # 等待所有任务完成并收集它们的结果
            results = await asyncio.gather(*tasks)
        
        # 计算所有结果中的完成token总数
        completion_tokens = sum(result[0] for result in results)
        
        # 从所有结果中提取响应时间
        response_times = [result[1] for result in results]
        
        # 返回完成token的总数和响应时间的列表
        return completion_tokens, response_times

if __name__ == '__main__':
    import sys

    if len(sys.argv) != 3:
        print("Usage: python bench.py <C> <N>")
        sys.exit(1)

    C = int(sys.argv[1])  # 最大并发数
    N = int(sys.argv[2])  # 请求总数

    # vllm 和 ollama 都兼容了 openai 的 api 让测试变得更简单了
    url = 'http://localhost:11434/v1/chat/completions'

    start_time = time.time()
    completion_tokens, response_times = asyncio.run(run(url, C, N))
    end_time = time.time()

    # 计算总时间
    total_time = end_time - start_time
    # 计算每个请求的平均时间
    avg_time_per_request = sum(response_times) / len(response_times)
    # 计算每秒生成的 token 数量
    tokens_per_second = completion_tokens / total_time

    print(f'Performance Results:')
    print(f'  Total requests            : {N}')
    print(f'  Max concurrent requests   : {C}')
    print(f'  Total time                : {total_time:.2f} seconds')
    print(f'  Average time per request  : {avg_time_per_request:.2f} seconds')
    print(f'  Tokens per second         : {tokens_per_second:.2f}')

通过 Semaphore 对最大的并发量做了控制。
llama 和 vllm 都兼容了 openai 的 api 格式，因此直接拼接 openai 的 api 要求的 request body 即可。
最后，记录请求处理时间以及返回结果里生成的 token 数量。
以平均请求时长和平均每秒生成 token 数量作为关键指标。

最后准备一个这样的 bash 脚本记录并发 1 - 16 的情况下脚本的结果。

python bench.py 1 4
python bench.py 2 8
python bench.py 4 16
python bench.py 8 32
python bench.py 16 64

然后按照上文提到的 ollama 和 vllm 的启动脚本，分别启动 ollama 和 vllm 并运行 bash 脚本即可获取测试结果。

补充信息

有小伙伴反馈是不是同样的提示词作测试会影响结果，这里我看了下感觉问题不大：

提示词很短，就是缓存了也不会带来很大性能提升
temperature保证结果都有所区别，那么后续的结果也不会一直，缓存的意义也不大

不过，反正我自己的环境还在，我就直接对原来代码做了调整，支持了随机化的提示词：

import aiohttp
import asyncio
import time
from tqdm import tqdm

import random

questions = [
    "Why is the sky blue?", "Why do we dream?", "Why is the ocean salty?", "Why do leaves change color?",
    "Why do birds sing?", "Why do we have seasons?", "Why do stars twinkle?", "Why do we yawn?",
    "Why is the sun hot?", "Why do cats purr?", "Why do dogs bark?", "Why do fish swim?",
    "Why do we have fingerprints?", "Why do we sneeze?", "Why do we have eyebrows?", "Why do we have hair?",
    "Why do we have nails?", "Why do we have teeth?", "Why do we have bones?", "Why do we have muscles?",
    "Why do we have blood?", "Why do we have a heart?", "Why do we have lungs?", "Why do we have a brain?",
    "Why do we have skin?", "Why do we have ears?", "Why do we have eyes?", "Why do we have a nose?",
    "Why do we have a mouth?", "Why do we have a tongue?", "Why do we have a stomach?", "Why do we have intestines?",
    "Why do we have a liver?", "Why do we have kidneys?", "Why do we have a bladder?", "Why do we have a pancreas?",
    "Why do we have a spleen?", "Why do we have a gallbladder?", "Why do we have a thyroid?", "Why do we have adrenal glands?",
    "Why do we have a pituitary gland?", "Why do we have a hypothalamus?", "Why do we have a thymus?", "Why do we have lymph nodes?",
    "Why do we have a spinal cord?", "Why do we have nerves?", "Why do we have a circulatory system?", "Why do we have a respiratory system?",
    "Why do we have a digestive system?", "Why do we have an immune system?"
]

async def fetch(session, url):
    """
    参数:
        session (aiohttp.ClientSession): 用于请求的会话。
        url (str): 要发送请求的 URL。
    
    返回:
        tuple: 包含完成 token 数量和请求时间。
    """
    start_time = time.time()

    # 随机选择一个问题
    question = random.choice(questions) # <--- 这两个必须注释一个

    # 固定问题                                 
    # question = questions[0]             # <--- 这两个必须注释一个

    # 请求的内容
    json_payload = {
        "model": "llama3:8b-instruct-fp16",
        "messages": [{"role": "user", "content": question}],
        "stream": False,
        "temperature": 0.7 # 参数使用 0.7 保证每次的结果略有区别
    }
    async with session.post(url, json=json_payload) as response:
        response_json = await response.json()
        end_time = time.time()
        request_time = end_time - start_time
        completion_tokens = response_json['usage']['completion_tokens'] # 从返回的参数里获取生成的 token 的数量
        return completion_tokens, request_time

async def bound_fetch(sem, session, url, pbar):
    # 使用信号量 sem 来限制并发请求的数量，确保不会超过最大并发请求数
    async with sem:
        result = await fetch(session, url)
        pbar.update(1)
        return result

async def run(load_url, max_concurrent_requests, total_requests):
    """
    通过发送多个并发请求来运行基准测试。
    
    参数:
        load_url (str): 要发送请求的URL。
        max_concurrent_requests (int): 最大并发请求数。
        total_requests (int): 要发送的总请求数。
    
    返回:
        tuple: 包含完成 token 总数列表和响应时间列表。
    """
    # 创建 Semaphore 来限制并发请求的数量
    sem = asyncio.Semaphore(max_concurrent_requests)
    
    # 创建一个异步的HTTP会话
    async with aiohttp.ClientSession() as session:
        tasks = []
        
        # 创建一个进度条来可视化请求的进度
        with tqdm(total=total_requests) as pbar:
            # 循环创建任务，直到达到总请求数
            for _ in range(total_requests):
                # 为每个请求创建一个任务，确保它遵守信号量的限制
                task = asyncio.ensure_future(bound_fetch(sem, session, load_url, pbar))
                tasks.append(task)  # 将任务添加到任务列表中
            
            # 等待所有任务完成并收集它们的结果
            results = await asyncio.gather(*tasks)
        
        # 计算所有结果中的完成token总数
        completion_tokens = sum(result[0] for result in results)
        
        # 从所有结果中提取响应时间
        response_times = [result[1] for result in results]
        
        # 返回完成token的总数和响应时间的列表
        return completion_tokens, response_times

if __name__ == '__main__':
    import sys

    if len(sys.argv) != 3:
        print("Usage: python bench.py <C> <N>")
        sys.exit(1)

    C = int(sys.argv[1])  # 最大并发数
    N = int(sys.argv[2])  # 请求总数

    # vllm 和 ollama 都兼容了 openai 的 api 让测试变得更简单了
    url = 'http://localhost:11434/v1/chat/completions'

    start_time = time.time()
    completion_tokens, response_times = asyncio.run(run(url, C, N))
    end_time = time.time()

    # 计算总时间
    total_time = end_time - start_time
    # 计算每个请求的平均时间
    avg_time_per_request = sum(response_times) / len(response_times)
    # 计算每秒生成的 token 数量
    tokens_per_second = completion_tokens / total_time

    print(f'Performance Results:')
    print(f'  Total requests            : {N}')
    print(f'  Max concurrent requests   : {C}')
    print(f'  Total time                : {total_time:.2f} seconds')
    print(f'  Average time per request  : {avg_time_per_request:.2f} seconds')
    print(f'  Tokens per second         : {tokens_per_second:.2f}')