本地丐版运行xAI grok-1的尝试（失败版）

前言

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.pyL60修改成自己的显卡数量（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!!!

（实验室服务器内存大很多但是害怕搞崩了就不试了）