数据预处理阶段的资源调度策略

在大模型训练过程中，数据预处理阶段的资源调度往往被忽视，但却是影响整体训练效率的关键环节。最近在处理一个100GB的文本数据集时，我踩了一个大坑。

问题重现

最初我直接使用pandas读取整个数据集：

import pandas as pd
# 错误做法
raw_data = pd.read_csv('large_dataset.csv')

结果导致内存溢出，系统直接崩溃。经过排查发现，单个文件就占用了超过12GB内存。

正确的资源调度策略

分块读取 + 内存优化：

import pandas as pd
from tqdm import tqdm

def process_large_dataset(filename, chunk_size=10000):
    results = []
    total_rows = sum(1 for _ in open(filename)) - 1  # 减去header行
    
    # 分块处理数据
    for chunk in pd.read_csv(filename, chunksize=chunk_size):
        # 每个chunk进行预处理
        processed_chunk = chunk.dropna()  # 示例处理
        results.append(processed_chunk)
        
        # 显示进度
        pbar = tqdm(total=total_rows, desc="处理进度")
        pbar.update(len(chunk))
    
    return pd.concat(results, ignore_index=True)

内存监控与释放：

import psutil
import gc

def monitor_memory():
    process = psutil.Process()
    print(f"当前内存使用: {process.memory_info().rss / 1024 / 1024:.2f} MB")

# 处理完及时释放内存
processed_data = process_large_dataset('large_dataset.csv')
gc.collect()  # 强制垃圾回收