[Feature]: 如何使用多个显卡以及Llama 13B 和70B的支持

[infinitywings]

Class | 类型

大语言模型

Feature Request | 功能请求

我直接从运行的docker镜像，但是发现chatglm 和llama2的模型都只运行在了一个gpu上，我的机器有4*A5000， 剩下的三个卡都是空着的。所以想问下本项目是否支持多卡并行，好让我把四个卡都用起来或者跑更大的模型。

[binary-husky]

是可以的，不过需要修改一些代码去改动 os.environ

[infinitywings]

是可以的，不过需要修改一些代码去改动 os.environ

我大概明白了，我自己试一下

[ggfengxuan]

是可以的，不过需要修改一些代码去改动 os.environ

我大概明白了，我自己试一下

您好，请问您解决了这个问题么？

[ggfengxuan]

是可以的，不过需要修改一些代码去改动 os.environ

我大概明白了，我自己试一下

您好，请问您解决了这个问题么？

已通过accelerate解决

[KimChow]

方便留个具体怎样操作的资料学习一下不？

[ggfengxuan]

之前Google的答案，时间太久了，记得不大清楚了，可以搜一下accelerate使用多张GPU

    def load_model_and_tokenizer(self):
        # 🏃‍♂️🏃‍♂️🏃‍♂️ 子进程执行
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
        from accelerate import dispatch_model, infer_auto_device_map
        from accelerate.utils import get_balanced_memory
        device = get_conf('LOCAL_MODEL_DEVICE')
        with ProxyNetworkActivate('Download_LLM'):
            if self._model is None:
                tokenizer = AutoTokenizer.from_pretrained("/xxx/model/internlm-chat-7b", trust_remote_code=True)
                if device=='cpu':
                    model = AutoModelForCausalLM.from_pretrained("/xxx/model/internlm-chat-7b", trust_remote_code=True).to(torch.bfloat16)
                else:
                    model = AutoModelForCausalLM.from_pretrained("xxx/model/internlm-chat-7b", trust_remote_code=True,device_map="auto").to(torch.bfloat16).cuda()
                max_memory = get_balanced_memory(
                    model,
                    max_memory=None,
                    no_split_module_classes=["InternLMDecoderLayer","InternLMMLP"],
                    dtype='float16',
                    low_zero=False,
                )

                device_map = infer_auto_device_map(
                    model,
                    max_memory=max_memory,
                    no_split_module_classes=["InternLMDecoderLayer","InternLMMLP"],
                    dtype='float16'
                )

                model = dispatch_model(model, device_map=device_map)
                model = model.eval()
        return model, tokenizer