sophgo/LLM-TPU

GPU memory allocation failure

Closed this issue · 2 comments

我在对LLM-TPU进行二次开发的时候(适配一个小众的国产模型)
显存分配报错炸掉了
报错代码大致如下:

[BMRT][load_bmodel:1573] INFO:Loading bmodel from [../../bmodels/rwkv6-1b5_bf16_1dev.bmodel]. Thanks for your patience...
[BMRT][load_bmodel:1501] INFO:pre net num: 0, load net num: 1
[bmlib_memory][error] bm_alloc_gmem failed, dev_id = 0, size = 0xc0fb4000

但我导出的模型体积,明显小于官方提供的实例文件中的llama模型文件
我的模型体积:

du -shb ../../bmodels/rwkv6-1b5_bf16_1dev.bmodel
3244769616      ../../bmodels/rwkv6-1b5_bf16_1dev.bmodel

llama模型体积:

du -shb ../../../LLM-TPU/bmodels/llama2-7b_int4_1dev.bmodel 
3888418068      ../../../LLM-TPU/bmodels/llama2-7b_int4_1dev.bmodel

我的工程大部分都参考了llama2的cpp工程,我是用我的工程加载llama模型不会出错,但加载我自己导出的模型产生报错(该模型已经在windows docker经过了model_runner.py验证)。
通过观察llama的模型结构发现llama模型在onnx导出阶段进行了拆层处理,将一个模型拆分为多个部分进行实现
我进行模型加载的过程中报错gpu内粗分配失败,是因为没有做模型层的拆分吗?

好吧,拆完层,加载确实不会炸了,所以bm_alloc_gmem的最大分配上限是多少鸭