在PySpark中运行PyTorch GPU模型时遭遇CUDA内存不足错误的排查求助

大家好，我最近在尝试用PySpark在GPU上跑PyTorch模型，结果一直碰到CUDA内存不足的错误，试了好几种调整方法都没解决，想请各位帮忙看看问题出在哪。

错误详情 #

触发错误的代码行是PyTorch的设备转换操作：

return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)

完整的错误日志如下：

torch.cuda.OutOfMemoryError: CUDA out of memory.
Tried to allocate 20.00 MiB (GPU 0; 14.76 GiB total capacity; 56.20 MiB already allocated; 18.75 MiB free; 58.00 MiB reserved in total by PyTorch)
If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.

测试用代码 #

为了排查问题，我写了一段用随机输入测试模型的代码，具体如下：

import torch
from pyspark.sql import SparkSession
from pyspark import SparkConf

appName = "PySpark Test"
conf = SparkConf().setAppName(appName)
conf.set("spark.executorEnv.LD_PRELOAD", "libnvblas.so")
conf.set("spark.executor.resource.gpu.amount", "1")
conf.set("spark.task.resource.gpu.amount", "1")  # 注：原代码里写的spark.spark.task应该是笔误，我已经改成spark.task了
conf.set("spark.executor.cores", "2")
conf.set("spark.executor.instances", "1")  # 注：原代码里的spark.executor.num是错误的配置项，正确应该是spark.executor.instances
conf.set("spark.rapids.sql.enabled", "true")

# 创建Spark会话
spark = SparkSession.builder.config(conf=conf).getOrCreate()

def load_model(device):
    # 加载模型并放到指定设备上，这里省略了具体加载逻辑
    return model

def run_dummy_model(i):
    assert torch.cuda.is_available()  # 这里能通过，说明GPU确实是可用的
    # 生成随机测试输入
    t = torch.rand(1, 3, 128, 128)
    device = torch.device("cuda")
    model = load_model(device)
    t.to(device)
    outputs = model(t)
    return outputs

# 生成10个测试任务
input = spark.sparkContext.parallelize([0] * 10)
result = input.map(run_dummy_model)
print(result.take(10))

运行环境 #

• 我是在GCP Dataproc的单GPU节点集群上运行这段代码的
• 代码里的assert torch.cuda.is_available()没有报错，说明GPU的基础配置应该是没问题的

已尝试的解决思路 #

之前我遇到过类似的CUDA内存不足问题，是因为多个PyTorch进程抢占同一块GPU导致的，所以这次我特意调整了Spark的executor配置，设置只启动1个executor，每个executor分配1块GPU，但还是出现了同样的错误。

现在实在找不到问题点了，想请教各位：

• 我的Spark配置还有哪里可能有问题？
• 除了多进程抢占GPU，还有哪些原因会导致这种看似内存充足但分配失败的情况？
• 错误提示里提到的max_split_size_mb参数应该怎么设置？

麻烦大家帮忙分析一下，非常感谢！

备注：内容来源于stack exchange，提问作者DzedCPT