前言
xAI开源了包含3000多亿参数的grok-1,想试试在本地跑。试了半天结果内存不够,结果以失败告终,结论是机器丐不了一点,想要跑起来内存必须要管够,显存应该也是需要的(xAI好像用的8*A100 80G NvLink)。如果内存的话应该可以跑起来?
本地硬件
OS: Ubuntu 20.04 5.15.0-101-generic
CPU: Intel i9 12900K
GPU: NVIDIA GeForce RTX 3090 24GB
Memory: 64GB+300GB SWAP
Storage: 1TB SSD+2TB HDD
NVIDIA Driver: 545.23.06
CUDA: V12.3.52 (cuDNN-linux-x86_64-8.9.7.29-cuda12)
代码和checkpoints
这些都可以在github上找到https://github.com/xai-org/grok-1,然后下载权重文件,大概296GB,我用的磁力链接和qBittorrent下的,然后权重文件夹ckpt-0
放到checkpoints
目录下。
修改代码
直接跑可能跑不了,需要在run.py
L60修改成自己的显卡数量(see issue)。
运行
bash输入pip install -r requirements.txt
,其中,jax的版本最好去掉,因为要和自己的机器环境一致(我重新安装了N卡驱动545.23.06和nvcc 8.9.7.29,jax版本为0.4.25+cuda12.cudnn89-cp310,参考JAX: 库安装和GPU使用,解决不能识别gpu问题
,【Jax报错】CUDA backend failed to initialize: Unable to load cuPTI )
环境配好之后,使用python run.py
运行,期间会在/dev/shm
里面拷贝那296G的文件作为共享内存啥的,最好先腾出来空间(/dev/shm 调整大小),而且IO如果差一些可能会像下面这样在Loading checkpoint at ./checkpints/ckpt-0
卡很长时间。
鼠鼠就在这里卡了几个小时,最后还是报内存不够。
可以看到项目对内存的需求还是很大的,64GB RAM + 300GB SWAP还是不够它耗的。此外,对显存的需求好像不是很大?因为我看显存占用一直是18GB左右,也可能是还没运行到吃显存的那一步。
后面估计显存也会是一个大问题,那么这个尝试就到此为止吧,,
GPU is all you need!!! So is money!!!
(实验室服务器内存大很多但是害怕搞崩了就不试了)