lightllm icon indicating copy to clipboard operation
lightllm copied to clipboard

[Feature]请帮忙提供load_from_weight_dict(weight_dict)接口。

Open bingo787 opened this issue 1 year ago • 10 comments

需求背景: TGI适配lightllm,多卡加载模型的时候,用到几张卡就会有几个进程,并且每个进程都会完整的加载整个模型到内存中来。 当模型文件太大,比如65B以上的模型,使用8卡加载的话就会需要8*130G的内存,这显然是不合理的,会导致OOM。

解决办法: 可在lightllm中帮忙提供load_from_weight_dict(weight_dict) 接口。TGI层传入权重词典,一边加载一边释放内存,才能解决此问题。

bingo787 avatar Jan 04 '24 03:01 bingo787

@bingo787 你如果存储文件的权重是分片足够小,理论上应该不会OOM。或者把权重文件切细小一点也可以。

hiworldwzj avatar Jan 04 '24 03:01 hiworldwzj

@bingo787 你如果存储文件的权重是分片足够小,理论上应该不会OOM。或者把权重文件切细小一点也可以。

跟权重的分片没关系,TGI的8个进程是独立的,每个进程都要读全量权重到内存(变量)中,然后透传给Lightllm。 我们验证过8卡跑65B的模型,需要8*130G的内存,一般服务器没有那么多内存,一定会发生OOM。 你们可以跑一下TGI+Lightllm的组合,问题必现。

bingo787 avatar Jan 04 '24 07:01 bingo787

@bingo787 可以改造一下加载部分的逻辑呀,因为lightllm本身的加载实际上是按照分片的权重文件,一个一个读取的,所以实际上8个进程读取的时候,每个进程同时只会读取一个小片文件,处理完当前文件再处理下一个,理论上如果权重分片足够细,就可以避免OOM。或许是和TGI的配合部分出了问题,可以改造一下相关的代码。

hiworldwzj avatar Jan 07 '24 05:01 hiworldwzj

@bingo787 我们尝试过把TGI 跟LightLLM 加载,你们不需要再TGI里在进行这个权重的加载了,直接走默认LightLLM 的参数加载也是分片读取的,我尝试过在128G的内存环境下8卡读取70B的 llama模型,不会出现OOM 现象

shihaobai avatar Jan 23 '24 08:01 shihaobai

@bingo787 我们尝试过把TGI 跟LightLLM 加载,你们不需要再TGI里在进行这个权重的加载了,直接走默认LightLLM 的参数加载也是分片读取的,我尝试过在128G的内存环境下8卡读取70B的 llama模型,不会出现OOM 现象

感谢回复!按照你说的流程确实没有问题, 因为直接把模型的本地目录透传给lightllm了,真正加载权重的任务全部交给lightllm了。 但是我们想一下另一种场景: 我需要在TGI里面对权重做一些客制化的动作(比如模型文件解密),这种情况lightllm如果不提供load_from_weight_dict之类的接口, 好像没法处理?或者有更好的方式吗?

bingo787 avatar Jan 23 '24 09:01 bingo787

@bingo787 这种情况确实可以加一个load_from_weight_dict,你们可以https://github.com/ModelTC/lightllm/blob/4a9824b6b248f4561584b8a48ae126a0c8f5b000/lightllm/common/basemodel/basemodel.py#L79快速在这个位置修改,实现想要的功能,后续我们也加一个支持

shihaobai avatar Jan 23 '24 10:01 shihaobai

@shihaobai 我提了个PR,https://github.com/ModelTC/lightllm/pull/318 帮忙review下。

bingo787 avatar Jan 25 '24 07:01 bingo787

ok,方便的话可以根据意见简单修改一下

shihaobai avatar Jan 25 '24 08:01 shihaobai

ok,方便的话可以根据意见简单修改一下

我没有看到review comment

bingo787 avatar Jan 26 '24 13:01 bingo787

@hiworldwzj 当创建对象的时候,weight_dict传入空,肯定不能verfiy_load吧。PR318 我觉得是合理的,或者你有什么修改建议吗?

bingo787 avatar Jan 30 '24 03:01 bingo787