Intel Extension for PyTorch在Intel GPU上运行大语言模型时的段错误问题分析

问题背景

在使用Intel Extension for PyTorch（IPEX）配合accelerate库运行33B参数量的WizardCoder模型时，开发者遇到了段错误(Segmentation fault)问题。该问题发生在模型生成文本的过程中，特别是在使用Intel Arc A770显卡和CPU混合设备映射的情况下。

环境配置

开发者使用的环境配置如下：

• 硬件：Intel Arc A770 16GB显卡 + Intel Core i5-6500 CPU
• 软件：
  • PyTorch 2.1.0a0
  • Intel Extension for PyTorch 2.1.10+xpu
  • accelerate 0.26.1
  • transformers 4.37.1
• 操作系统：Fedora Linux 39

问题现象

开发者最初尝试使用以下方式加载和运行模型：

1. 使用init_empty_weights初始化空模型
2. 通过infer_auto_device_map自动分配设备映射
3. 加载预训练模型并创建文本生成管道

在生成文本时，程序会突然崩溃并产生段错误。核心转储显示错误发生在parallel_cat操作中，这表明问题可能与张量拼接操作有关。

问题排查

开发者进行了多方面排查：

1. 内存使用检查：通过xpu-smi工具确认GPU内存使用约为8GB（qint8精度）或4GB（bfloat16精度），排除了OOM（内存不足）的可能性。
2. 简化模型分配：即使只将2个模型层分配到GPU上，问题仍然存在。
3. 修改加载方式：直接使用pipeline的device_map="auto"参数加载模型，问题得到解决。

技术分析

从技术角度来看，这个问题可能涉及以下几个方面：

1. 设备间数据传输：当模型部分在GPU、部分在CPU时，需要进行频繁的设备间数据传输。Intel GPU的并行处理机制可能对这种跨设备操作支持不够完善。

2. 张量拼接操作：核心转储显示错误发生在parallel_cat操作中，这可能是由于不同设备上的张量拼接时出现了同步问题。

3. 内存管理：虽然总内存使用量不高，但可能存在内存碎片化或特定操作的内存分配问题。

解决方案

开发者最终找到了有效的解决方案：

1. 避免手动设备映射，直接使用pipeline的自动设备分配功能。
2. 在创建pipeline时指定完整的模型参数，包括数据类型、内存限制和缓冲区卸载选项。

修正后的代码如下：

pipe = pipeline("text-generation", 
               model=model_id, 
               model_kwargs={
                   "torch_dtype": torch.bfloat16,
                   "device_map": "auto",
                   "max_memory": {0: "8GB", "cpu": "128GB"},
                   "offload_buffers": True
               })

经验总结

1. 对于大型语言模型，推荐使用pipeline内置的设备分配机制，而不是手动管理。
2. 在Intel GPU上运行模型时，确保使用最新版本的驱动和软件栈。
3. 监控GPU内存使用情况时，不仅要关注总量，还要注意内存分配模式。
4. 跨设备操作（GPU-CPU）可能引入额外的复杂性，应尽量减少这类操作。

这个问题展示了在使用新兴硬件（如Intel独立显卡）进行深度学习推理时可能遇到的挑战，也体现了PyTorch生态系统在异构计算方面的持续改进空间。