search

THUDM / ChatGLM-6B

ChatGLM-6B: An Open Bilingual Dialogue Language Model | 开源双语对话语言模型
Apache License 2.0
40.47k stars 5.19k forks source link

Do the device_map more automaticlly for multi-GPUs #1381

Open alexhegit opened 1 year ago

alexhegit commented 1 year ago

The accelerate lib provides some functions to do the device_map automaticlly. This patch use infer_auto_device_map() to relpace the hard codes which splite the model by hand.

It works well with my two GPUs(6GB-RTX3060 + 12GB-RTX3060).