Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
bingo787 opened this issue Jan 4, 2024 · 10 comments
Open

Comments

@bingo787
Copy link
Contributor

bingo787 commented Jan 4, 2024

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

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

@hiworldwzj
Copy link
Collaborator

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

@bingo787
Copy link
Contributor Author

bingo787 commented Jan 4, 2024

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

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

@hiworldwzj
Copy link
Collaborator

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

@shihaobai
Copy link
Collaborator

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

@bingo787
Copy link
Contributor Author

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

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

bingo787 pushed a commit to bingo787/lightllm that referenced this issue Jan 25, 2024
@bingo787
Copy link
Contributor Author

@shihaobai 我提了个PR,#318
帮忙review下。

@shihaobai
Copy link
Collaborator

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

@bingo787
Copy link
Contributor Author

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

我没有看到review comment

@bingo787
Copy link
Contributor Author

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

3 participants