分布式训练中数据预处理阶段性能瓶颈排查

分布式训练中数据预处理阶段性能瓶颈排查

在分布式大模型训练中，数据预处理阶段往往是性能瓶颈的关键环节。近期在部署LLaMA-7B模型时，发现单卡训练耗时20分钟，而分布式训练（8卡）却达到45分钟，明显超出预期。

问题定位

通过torch.profiler分析，发现在DataLoader的__iter__方法中存在大量CPU等待时间。具体表现为：

# 瓶颈代码段
for batch in dataloader:
    # 数据加载耗时约30s
    # 预处理耗时约15s
    pass

排查步骤

1. 数据读取优化：将本地文件读取改为torch.utils.data.IterableDataset，并使用num_workers=4并行加载
2. 预处理流水线：将文本tokenization与padding操作并行化
3. 内存优化：设置pin_memory=True，减少CPU到GPU的拷贝时间

复现代码

from torch.utils.data import DataLoader, Dataset

class CustomDataset(Dataset):
    def __init__(self, data_list):
        self.data = data_list
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        # 模拟预处理耗时
        text = self.data[idx]
        tokenized = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
        return tokenized

# 优化后的数据加载器
train_dataset = CustomDataset(data)
dataloader = DataLoader(
    train_dataset,
    batch_size=8,
    num_workers=4,
    pin_memory=True,
    persistent_workers=True
)

解决方案

最终通过调整num_workers=6、设置prefetch_factor=2，将预处理时间从30s降至12s，训练效率提升约30%。

优化后的配置：num_workers=6, pin_memory=True, prefetch_factor=2, persistent_workers=True。