在Colab配置CUDA C时遇ModuleNotFoundError: No module named 'nvcc_plugin'求助

原因分析 #

nvcc_plugin是旧版本Colab中用于编译CUDA代码的插件，当前Colab环境已更新，官方不再维护该插件，转而推荐使用更标准的编译方式构建CUDA扩展。

解决方案 #

方案1：用setuptools构建CUDA扩展（推荐） #

替换依赖nvcc_plugin的加载逻辑，改用torch.utils.cpp_extension完成编译：

1. 编写setup.py文件，示例代码：

from setuptools import setup
from torch.utils.cpp_extension import CUDAExtension, BuildExtension

setup(
    name='custom_cuda_ext',
    ext_modules=[
        CUDAExtension('custom_cuda_ext', [
            'src/main.cpp',
            'src/kernel.cu',
        ]),
    ],
    cmdclass={
        'build_ext': BuildExtension
    }
)

2. 在Colab中执行编译命令：

!python setup.py build_ext --inplace

3. 直接导入编译后的模块：

import custom_cuda_ext

方案2：直接用nvcc编译并通过ctypes加载 #

针对简单CUDA代码，可跳过插件直接编译：

1. 编译CUDA代码为共享库：

!nvcc -o cuda_kernel.so -shared -fPIC your_kernel.cu

2. 在Python中加载调用：

import ctypes
lib = ctypes.CDLL('./cuda_kernel.so')
# 调用库中函数，需根据CUDA函数签名定义参数类型

方案3：回退环境安装旧版插件（不推荐） #

如果必须使用nvcc_plugin，可尝试回退到兼容环境：

1. 卸载当前PyTorch相关包：

!pip uninstall -y torch torchvision torchaudio

2. 安装旧版PyTorch及依赖：

!pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 torchaudio==0.8.0 -f https://download.pytorch.org/whl/torch_stable.html

3. 安装nvcc_plugin：

!pip install git+https://github.com/andreinechaev/nvcc4jupyter.git

4. 加载插件：

%load_ext nvcc_plugin

注意事项 #

• 优先选择方案1或方案2，这是Colab当前官方支持的CUDA扩展开发方式，稳定性和兼容性更强。
• 方案3仅作为临时 workaround，旧版本环境存在功能缺失和安全风险。

内容的提问来源于stack exchange，提问作者Nabiha Faisal