从段错误到稳定加载：LLaMA-Factory 在 ROCm 下的数据适配实战

最近在折腾 AMD GPU 上的大模型微调时，LLaMA-Factory 确实是个省心的框架，但在 ROCm 环境下处理大规模数据集时，还是踩了不少坑。最典型的就是训练刚开始，进程直接 Segmentation Fault 退出，没有任何报错堆栈，只留下一脸懵逼的开发者。经过几天的排查和代码调整，我发现问题的核心往往不在模型本身，而在于 DataLoader 的多进程机制与 ROCm 内存管理的冲突，以及数据集预处理方式的差异。今天就把这次排查过程和最终的解决方案整理出来，希望能帮到同样在 AMD 平台上奋斗的朋友。

定位“隐形杀手”：num_workers 引发的段错误

在使用 LLaMA-Factory 启动微调任务时，默认的配置文件通常会开启多进程数据加载（num_workers > 0），这在 NVIDIA CUDA 环境下通常运行良好。然而，在 ROCm 平台上，PyTorch 的多进程启动机制（基于 fork）有时会导致子进程继承父进程的 HIP 上下文状态不一致，进而引发底层的段错误。

起初我以为是驱动版本问题，重装了多次 ROCm 栈无果。后来尝试将 num_workers 设置为 0，训练竟然奇迹般地跑通了，虽然速度稍慢，但稳定性大幅提升。这证实了问题出在多进程派生上。

对于生产环境，完全关闭多进程显然不可取。经过反复测试，我发现将 num_workers 限制在物理核心数的一半，并显式设置 persistent_workers=False，能有效避免上下文污染。以下是我在 LLaMA-Factory 的 ds_z3_config.json 或启动命令中采用的调整策略：

# 推荐启动参数示例
CUDA_VISIBLE_DEVICES=0 python src/train.py \
    --model_name_or_path meta-llama/Llama-3-8B \
    --dataset my_custom_dataset \
    --num_workers 2 \
    --persistent_workers false \
    --prefetch_factor 2

注意，这里的 num_workers 不宜过大。在我的 MI250 测试机上，设置为 2 或 4 是甜点值，再高就会复现随机性的崩溃。此外，确保你的 PyTorch 版本是明确编译为 ROCm 支持的版本（如 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.0），混用 CUDA 版本的 wheel 包是导致此类诡异报错的常见原因。

大型数据集预处理：避免内存溢出的最佳实践

解决了进程崩溃，下一个拦路虎是内存溢出（OOM）。在微调长文本或超大语料时，如果直接在 DataLoader 中进行实时 Tokenization，主进程内存极易飙升。特别是在 AMD 显卡上，系统内存与显存的协同机制与 NVIDIA 略有不同，对内存峰值更敏感。

我的建议是**“预加工，轻加载”**。不要依赖训练时的动态处理，而是提前将数据集转换为二进制格式（如 .arrow 或打包好的 .bin 文件）。LLaMA-Factory 支持直接加载预处理后的文件，这样 DataLoader 只需要负责读取字节流，极大地降低了 CPU 内存压力。

具体操作流程如下：

1. 使用单独的脚本遍历原始 JSON/Text 数据。
2. 调用对应的 Tokenizer 进行编码，并添加必要的特殊 token。
3. 将结果保存为 Arrow 格式。

这种“空间换时间”的策略在 ROCm 环境下尤为有效，它不仅避免了训练过程中的内存抖动，还让 GPU 能更专注于计算而非等待数据预处理。

自定义 Dataset 类：确保 ROCm 兼容性的代码示例

如果官方支持的数据集格式无法满足需求，我们需要编写自定义的 Dataset 类。在 ROCm 环境下，关键在于确保 __getitem__ 方法中没有隐式的 CUDA 操作，并且返回的数据类型严格符合 PyTorch 的期望。

下面是一个经过验证的自定义 Dataset 模板，专门针对 AMD 环境优化，去除了可能导致上下文切换的冗余操作：

from torch.utils.data import Dataset
import torch

class RocmCompatibleDataset(Dataset):
    def __init__(self, data_path, tokenizer, max_length=2048):
        # 建议在此处加载预处理好的 arrow 文件，而非原始文本
        self.data = load_processed_data(data_path) 
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        item = self.data[index]
        # 确保输入_ids 和 attention_mask 在 CPU 上构建，由 Trainer 自动移至设备
        input_ids = torch.tensor(item['input_ids'], dtype=torch.long)
        attention_mask = torch.tensor(item['attention_mask'], dtype=torch.long)
        labels = torch.tensor(item['labels'], dtype=torch.long)
        
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": labels
        }

在这个实现中，特别注意不要在 __getitem__ 内部调用 .to('cuda') 或类似的设备转移操作，因为 DataLoader 的 worker 进程可能没有正确的 HIP 上下文。让训练循环（Trainer）统一处理设备转移是最安全的做法。同时，数据类型显式声明为 torch.long 可以避免因默认类型推断导致的精度不匹配问题。

结语

在 AMD ROCm 平台上运行 LLaMA-Factory 并非不可逾越的高山，关键在于理解其底层资源调度机制的差异。通过合理限制 num_workers、采用预加工数据策略以及编写规范的自定义 Dataset，我们完全可以构建出稳定高效的微调流水线。随着社区对 HIP 生态的贡献日益增多，像 SGLang 推理加速、TileLang 算子优化等工具也在不断补齐短板，AMD GPU 的性价比优势正逐渐转化为实实在在的生产力。

如果你也想亲手验证这些优化效果，或者需要更多算力来跑通自己的大模型实验，现在有个不错的机会。200 小时 GPU 算力已就位，快来领取：https://marketing.csdn.net/questions/Q2604140858304426315?utm_source=AIpaper

(注：上方为示意海报位置，实际活动请以链接页面展示为准)