Spark+PyTorch多GPU环境CUDA内存不足及模型复用优化咨询

使用配备448GB内存与4块GPU的计算设备，运行以下调用Mistral-7B模型的Spark+pandas_udf代码：

import torch
import pandas as pd
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import StringType
from pyspark.sql import SparkSession
from transformers import pipeline

# Initialize the SparkSession
spark = SparkSession.builder \
    .appName("MistralUDFExample") \
    .getOrCreate()

# Initialize the pipeline outside the UDF
pipe = pipeline("text-generation", model="HuggingFaceH4/mistral-7b-sft-beta", torch_dtype=torch.bfloat16, device_map="auto")

@pandas_udf(StringType())
def mistral_udf(texts: pd.Series) -> pd.Series:
    results = []
    for text in texts:
        messages = [
            {
                "role": "system",
                "content": "You are a helpful assistant which only task is to analyze the text of emails sent to a customer service mailbox.  From this text you must extract a category.  No category labels will be provided. Email:",
            },
            {"role": "user", "content": text},
        ]
        prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        outputs = pipe(prompt, max_new_tokens=8, do_sample=True, temperature=0, top_k=50, top_p=0.95)
        generated_text = outputs[0]["generated_text"]
        results.append(generated_text)
    return pd.Series(results)

# Example usage
input_df = spark.createDataFrame(pd.DataFrame({'texts': ['Terrible experience shopping online?', 'When I called, the human on the other side was vert friendly']}))
output_df = input_df.withColumn('generated_text', mistral_udf(input_df['texts']))
output_df.show(truncate=False)

代码首次执行正常，但二次执行时触发以下错误：

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 112.00 MiB (GPU 1; 15.77 GiB total capacity; 3.16 GiB already allocated; 47.12 MiB free; 3.22 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

需求：优化代码以复用已加载的模型，解决CUDA内存不足问题。

1. 模型复用：使用Spark广播变量 #

当前代码在Driver端初始化模型，但Spark的pandas_udf运行在Executor节点，Driver端模型不会自动传递给Executor，二次执行时Executor可能重复加载模型导致内存泄漏。用广播变量将模型分发到每个Executor，确保每个Executor仅加载一次模型。

修改后代码：

import torch
import pandas as pd
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import StringType
from pyspark.sql import SparkSession
from transformers import pipeline
import os

# 设置PyTorch内存配置，减少碎片化
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

def init_model():
    return pipeline(
        "text-generation",
        model="HuggingFaceH4/mistral-7b-sft-beta",
        torch_dtype=torch.bfloat16,
        device_map="auto",
        model_kwargs={"use_cache": True},  # 复用计算结果，降低内存占用
        pipeline_kwargs={"return_full_text": False}  # 仅返回生成内容，减少数据传输
    )

# 初始化SparkSession并配置GPU资源
spark = SparkSession.builder \
    .appName("MistralUDFExample") \
    .config("spark.executor.resource.gpu.amount", "1")  # 每个Executor分配1块GPU
    .config("spark.task.resource.gpu.amount", "1")  # 每个Task分配1块GPU
    .getOrCreate()

# 广播模型到所有Executor
broadcast_pipe = spark.sparkContext.broadcast(init_model())

@pandas_udf(StringType())
def mistral_udf(texts: pd.Series) -> pd.Series:
    pipe = broadcast_pipe.value
    # 批量构造对话模板
    messages_list = [
        [
            {"role": "system", "content": "You are a helpful assistant which only task is to analyze the text of emails sent to a customer service mailbox.  From this text you must extract a category.  No category labels will be provided. Email:"},
            {"role": "user", "content": text}
        ] for text in texts
    ]
    # 批量生成prompt
    prompts = [pipe.tokenizer.apply_chat_template(msg, tokenize=False, add_generation_prompt=True) for msg in messages_list]
    # 批量推理，减少内存波动
    outputs = pipe(
        prompts,
        max_new_tokens=8,
        do_sample=False,  # temperature=0时无需采样，降低内存开销
        temperature=0,
        top_k=50,
        top_p=0.95,
        pad_token_id=pipe.tokenizer.eos_token_id  # 指定pad token，避免额外内存消耗
    )
    # 提取结果
    results = [output["generated_text"] for output in outputs]
    # 清理CUDA缓存
    torch.cuda.empty_cache()
    return pd.Series(results)

# 示例调用
input_df = spark.createDataFrame(pd.DataFrame({'texts': ['Terrible experience shopping online?', 'When I called, the human on the other side was very friendly']}))
output_df = input_df.withColumn('generated_text', mistral_udf(input_df['texts']))
output_df.show(truncate=False)

# 作业结束后释放资源
broadcast_pipe.unpersist()
torch.cuda.empty_cache()
spark.stop()

2. 内存优化关键措施 #

• 关闭冗余采样逻辑：temperature=0时设置do_sample=False，避免不必要的内存开销
• 批量处理文本：替换逐行循环为批量推理，减少内存碎片化，提升GPU利用率
• 配置Spark GPU资源：明确每个Executor/Task的GPU分配，避免多任务抢占同一GPU内存
• 优化PyTorch内存：设置PYTORCH_CUDA_ALLOC_CONF参数减少内存碎片化
• 手动清理缓存：在UDF结束后调用torch.cuda.empty_cache()释放未使用的显存
• 复用模型计算缓存：初始化模型时启用use_cache=True，复用之前的计算结果

3. 避免二次执行内存泄漏 #

通过广播变量确保每个Executor仅加载一次模型，后续任务复用该实例，避免重复加载导致的内存累积；作业结束后主动释放广播变量和GPU缓存，清除残留内存占用。

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