WSL2+Docker环境Spark-NLP GPU加速无差异问题求助

一、核心验证：GPU是否实际参与计算 #

虽然nvidia-smi显示模型已加载到GPU内存，但Spark-NLP可能未正确触发GPU加速，需先验证：

• 查看Spark任务日志，搜索GPU/CUDA关键词，确认是否有GPU设备初始化的记录
• 在代码中添加配置检查：

  from pyspark.sql import SparkSession
  spark = SparkSession.builder.getOrCreate()
  # 检查Spark-NLP GPU包是否加载
  print(spark.conf.get("spark.jars.packages"))
  # 检查GPU资源分配配置
  print(spark.conf.get("spark.executor.resource.gpu.amount"))
  print(spark.conf.get("spark.task.resource.gpu.amount"))
  

二、常见原因及修复方案 #

1. Spark未配置GPU资源分配 #

即使容器能访问GPU，Spark本身未分配GPU资源的话会默认用CPU：

• 启动SparkSession时必须指定GPU相关配置：

  spark = SparkSession.builder \
      .appName("SparkNLP-GPU") \
      .config("spark.driver.memory", "16G") \
      .config("spark.executor.memory", "16G") \
      .config("spark.executor.resource.gpu.amount", "1") \
      .config("spark.task.resource.gpu.amount", "1") \
      .config("spark.executor.resource.gpu.discoveryScript", "/usr/lib/spark/scripts/gpu/getGpusResources.sh") \
      .config("spark.jars.packages", "com.johnsnowlabs.nlp:spark-nlp_2.12:5.1.4,com.johnsnowlabs.nlp:spark-nlp-gpu_2.12:5.1.4") \
      .getOrCreate()
  

  注意：spark-nlp与spark-nlp-gpu版本必须完全一致，且需匹配Spark的Scala版本（如Spark 3.3.x对应Scala 2.12）

2. 模型不支持GPU加速 #

并非所有Spark-NLP预训练模型都能利用GPU：

• 仅标注GPU标识的Transformer模型（如BERT、RoBERTa系列）或深度学习类模型才会触发GPU计算
• 轻量规则类模型（如Tokenizer、StopWordsRemover）计算量小，GPU与CPU速度差异不明显

3. Docker容器GPU挂载不完整 #

启动容器时未正确传递GPU权限，导致Spark无法调用GPU：

• 启动容器必须添加--gpus all参数：

  docker run --gpus all -it your-spark-nlp-image bash
  

• 验证容器内CUDA环境：运行nvcc --version确认CUDA版本（Spark-NLP 5.x推荐CUDA 11.7+）

三、推荐版本组合（WSL2+Docker+GPU环境） #

• Spark: 3.3.3（Scala 2.12）
• Spark-NLP: 5.1.4（spark-nlp_2.12与spark-nlp-gpu_2.12均为该版本）
• CUDA: 11.7/12.0
• PySpark: 3.3.3（与Spark版本严格一致）
• Python: 3.8~3.10

四、最终验证 #

• 运行大批次Transformer模型推理任务（如处理10万+条文本），对比CPU/GPU运行时间
• 实时查看nvidia-smi的Utilization列，确认推理过程中GPU利用率是否上升（若始终为0%，则GPU未被实际调用）

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