大模型训练中数据预处理优化

大模型训练中数据预处理优化踩坑记录

最近在做大模型训练时，发现数据预处理环节成了性能瓶颈，特此记录一些踩坑经验和优化方法。

问题背景

训练过程中发现GPU利用率低，主要原因是数据加载和预处理耗时过长。通过profile分析，发现80%的时间都花在了数据准备上。

常见坑点及解决方案

1. 数据读取效率低下 最初使用Python原生的文件读取方式：

for line in open('train.txt', 'r'):
    data.append(process_line(line))

优化后改用torch.utils.data.DataLoader配合自定义Dataset，同时开启多进程加载：

from torch.utils.data import DataLoader, Dataset

class MyDataset(Dataset):
    def __init__(self, file_path):
        self.data = load_data(file_path)
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return process_data(self.data[idx])

# 使用多进程加载
loader = DataLoader(MyDataset('train.txt'), batch_size=32, num_workers=4)

2. 内存使用不当 预处理时一次性加载所有数据到内存，导致OOM。解决方法是采用流式处理和分批处理：

# 错误做法
all_data = load_all_data()
processed_data = [preprocess(d) for d in all_data]

# 正确做法
def data_generator():
    for batch in get_batched_data():
        yield preprocess_batch(batch)

3. 编码效率问题 在处理tokenize时，使用了低效的循环：

# 优化前
for text in texts:
    tokens = tokenizer.encode(text)

# 优化后
batch_tokens = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')

复现步骤

1. 使用torch.utils.data.DataLoader并设置合适的num_workers
2. 预处理数据时使用tokenizer的批量处理功能
3. 通过torch.utils.data.IterableDataset实现流式数据加载

总结

数据预处理优化对大模型训练效率至关重要，建议在项目初期就考虑好数据管道的设计。