search

THUDM / SwissArmyTransformer

SwissArmyTransformer is a flexible and powerful library to develop your own Transformer variants.
https://THUDM.github.io/SwissArmyTransformer
Apache License 2.0
951 stars 90 forks source link

单机多卡训练时内存占用过高 #158

Closed zodiacg closed 7 months ago

zodiacg commented 7 months ago

使用8*3090+约400G内存精调CogVLM时候（MP_SIZE=4），发现使用8张显卡精调内存会在加载参数阶段爆掉，使用4张就能够正常加载进行精调。

稍微调试了一下发现‎BaseModel.from_pretrained方法在修改了model_parallel_size时，似乎每个进程都会在cpu加载完整模型然后再分割参数。请问这个是必要的吗？有没有办法使用其它方式减少这部分的主内存占用？

1049451037 commented 7 months ago

并不会在每个进程都加载cpu，只有0号卡会加载cpu，所以4卡和8卡应该是一样的内存。

zodiacg commented 7 months ago

找到问题了，NUM_GPUS_PER_WORKER一直写的1……导致rank计算有问题