多卡部署时缺少utils模块/Missing utils module while running on multiple gpus

[Chenhong-Zhang]

It was recommended that we can use following code to run on multiple gpus:

from utils import load_model_on_gpus
model = load_model_on_gpus("THUDM/chatglm-6b", num_gpus=6)

But I did not find a utils module in this repository.

I tried simply copying the utils module from ChatGLM1 but it doesn't work. I'm guessing the structure was changed in ChatGLM3. So I modified the function load_model_on_gpus.

So I changed the device_map settings. And now it works:

def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike], num_gpus: int = 2,
                       device_map: Optional[Dict[str, int]] = None, **kwargs) -> Module:
    if num_gpus < 2 and device_map is None:
        model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half().cuda()
    else:
        from accelerate import dispatch_model

        model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half()
        

        model = dispatch_model(model, device_map="auto")

    return model

Basically I changed device_map to "auto".

Please let me know if I got anything wrong? And I would appreciate it if the developers could provide their official version of utils in ChatGLM3.