大模型数据处理流水线性能优化

在大模型训练中，数据处理流水线的性能直接影响训练效率。本文分享一套可复现的优化方案。

1. 数据读取优化

使用 torch.utils.data.DataLoader 的 num_workers 参数并行读取数据：

from torch.utils.data import DataLoader

# 建议设置为CPU核心数或稍大
dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=8,
    pin_memory=True,
    prefetch_factor=2
)

2. 数据预处理流水线优化

使用 torchdata 或自定义 Dataset 类：

import torch
from torch.utils.data import Dataset

class OptimizedDataset(Dataset):
    def __init__(self, data_paths):
        self.data_paths = data_paths
        
    def __getitem__(self, idx):
        # 预加载并缓存常用数据
        data = load_and_preprocess(self.data_paths[idx])
        return data

3. 内存管理优化

通过 torch.cuda.empty_cache() 和批量处理避免内存溢出：

import torch
for batch in dataloader:
    # 处理批次数据
    outputs = model(batch)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    
    # 定期清理缓存
    if step % 100 == 0:
        torch.cuda.empty_cache()

4. 性能监控

使用 torch.profiler 分析瓶颈：

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
    record_shapes=True
) as prof:
    outputs = model(inputs)

优化后可将数据处理时间降低30-50%。