问题

在DeepSpeed Offload时候需要编译并加载一个动态库，使得Adam优化器可以从GPU卸载到CPU。可以用以下方法测试编译是否成功：

1
2
import deepspeed
deepspeed.ops.op_builder.CPUAdamBuilder().load()

我刚开始的时候遇到了报错，提示找不到其中的一个动态链接库：

1
/home/xxx/miniconda3/envs/Visual-RFT/bin/../lib/gcc/x86_64 - conda - linux - gnu/11.2.0/../../../../x86_64 - conda - linux - gnu/bin/ld: cannot find - lcurand: No such file or directory

解决方案

如果要编译动态库，需要用到CUDA Toolkit，而且要是官网下载的（Conda/Pytorch安装的CUDA Toolkit是不全的）。

所以需要：

  1. 下载安装CUDA Toolkit
  2. 设置环境变量让DeepSpeed可以找到（因为我是非ROOT用户安装的，我装在了~/CUDA12.2下）

然后这样初始化即可：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def initial_deepspeed():
    cuda_lib_path = os.path.expanduser("~/CUDA12.2/lib64")
    nvjitlink_lib_path = os.path.expanduser(
        "~/miniconda3/envs/Visual-RFT/lib/python3.10/site-packages/nvidia/nvjitlink/lib"
    )
    conda_lib_path = os.path.expandvars("$CONDA_PREFIX/lib")
    cuda_bin_path = os.path.expanduser("~/CUDA12.2/bin")

    # 设置 LD_LIBRARY_PATH（追加到原值）
    os.environ["LD_LIBRARY_PATH"] = ":".join(filter(None, [
        cuda_lib_path,
        nvjitlink_lib_path,
        conda_lib_path,
        os.environ.get("LD_LIBRARY_PATH", ""),
    ]))

    # 设置 CUDA_PATH（直接覆盖）
    os.environ["CUDA_PATH"] = os.path.expanduser("~/CUDA12.2")

    # 设置 PATH（追加到原值）
    os.environ["PATH"] = ":".join(filter(None, [
        os.environ.get("PATH", ""),
        cuda_bin_path
    ]))

    import deepspeed

    deepspeed.ops.op_builder.CPUAdamBuilder().load()

设置好环境变量就可以初始化成功了，但要在import deepspeed之前就设置好环境变量。