大模型分布式推理ray

贾亚飞

已于 2024-03-22 16:39:00 修改

阅读量2k

点赞数 5

分类专栏： AI 文章标签：分布式

于 2024-03-22 13:55:17 首次发布

本文链接：https://blog.csdn.net/weixin_40777649/article/details/136938999

版权

本文介绍了如何利用Ray框架将模型拆分到多GPU进行分布式预测和训练，通过一个实例展示了如何在星lette框架下创建并部署分布式推理服务。同时提供了安装教程和相关文档链接。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

一、目录

1 框架
2. 入门
3. 安装教程
4. 相关文档、案例阅读

二、实现

1 框架：Ray：将一个模型拆分到多个显卡中，实现分布式预测、训练等功能。
2. 入门：
案例：通过ray 实现分布式部署，分布式推理服务。
参考：https://zhuanlan.zhihu.com/p/647973148?utm_id=0
文件名：test.py
pip install ray
pip install “ray[serve]”

import pandas as pd

import ray
from ray import serve
from starlette.requests import Request

@serve.deployment(ray_actor_options={"num_gpus": 2})        #两个gpu 将模型拆分，进行推理
class PredictDeployment:
    def __init__(self, model_id: str):
        from transformers import AutoModelForCausalLM, AutoTokenizer
        import torch

        self.model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            device_map="auto",
        )
        self.tokenizer = AutoTokenizer.from_pretrained(m