[问题/Issue] 章节5：LLM模型预训练过程中内存不断增长 #59
Description
1. 遇到问题的章节 / Affected Chapter
Chapter6
2. 具体问题描述 / Problem Description
模型预训练过程中System Memory Utilization不断增长，直至100%进程被杀死，怀疑是dataset加载写法有问题
部分原因也是因为我将num_workers设置为了32，但减小只能缓解不能根本上解决问题。
实验配置H200*2+256G内存
deepseek给出的解析可以部分参考一下，部分回答存在歧义：
从你提供的代码和描述来看，系统内存利用率不断增长直至崩溃的问题可能与内存泄漏有关。以下是可能的原因和解决方案：
主要问题分析
数据集加载方式：
你一次性将所有数据加载到内存中（self.data = f.readlines()），这对于大文件会占用大量内存
每次迭代时都会创建新的numpy数组和torch张量，这些可能没有及时释放
DataLoader配置：
num_workers=32 设置过高，每个worker都会复制数据集，可能导致内存爆炸
pin_memory=True 会锁定内存页，可能增加内存压力
梯度累积：
虽然梯度累积本身不会直接导致内存泄漏，但在长时间训练中可能加剧内存问题
解决方案
优化数据集加载：
class PretrainDataset(Dataset):
def init(self, data_path, tokenizer, max_length=512):
super().init()
self.data_path = data_path
self.tokenizer = tokenizer
self.max_length = max_length
self.padding = 0
# 只存储行数，不加载全部内容
with open(data_path, 'r', encoding='utf-8') as f:
self.length = sum(1 for _ in f)
def __len__(self):
return self.length
def __getitem__(self, index: int):
# 按需读取单行
with open(self.data_path, 'r', encoding='utf-8') as f:
for i, line in enumerate(f):
if i == index:
sample = json.loads(line)
break
# 其余代码保持不变...
调整DataLoader参数：
train_loader = DataLoader(
train_ds,
batch_size=args.batch_size,
pin_memory=False, # 禁用pin_memory
drop_last=False,
shuffle=True,
num_workers=4, # 减少worker数量
persistent_workers=True # 保持worker进程存活
)
添加内存监控：
在训练循环中添加内存监控代码：
import psutil
def train_epoch(epoch):
# ...
for step, (X, Y, loss_mask) in enumerate(train_loader):
# 打印内存使用情况
mem = psutil.virtual_memory()
Logger(f"Memory usage: {mem.percent}% (used: {mem.used/1024/1024:.2f}MB)")
# 其余代码...
定期清理缓存：
import gc
def train_epoch(epoch):
# ...
if step % 100 == 0:
torch.cuda.empty_cache()
gc.collect()
使用生成器式数据集：
考虑使用IterableDataset替代Dataset，特别是对于非常大的数据集：
from torch.utils.data import IterableDataset
class PretrainIterableDataset(IterableDataset):
def init(self, data_path, tokenizer, max_length=512):
super().init()
self.data_path = data_path
self.tokenizer = tokenizer
self.max_length = max_length
self.padding = 0
def __iter__(self):
with open(self.data_path, 'r', encoding='utf-8') as f:
for line in f:
sample = json.loads(line)
# 处理逻辑与之前相同...
yield X, Y, loss_mask
其他建议
减小验证频率：
减少保存检查点的频率（如每5000步改为每10000步）
混合精度训练：
确保正确使用混合精度训练，减少显存占用
分批处理大文件：
如果数据文件非常大，考虑先将其分割成多个小文件
监控GPU内存：
添加GPU内存监控：torch.cuda.memory_allocated()/1024/1024
通过这些修改，应该能显著降低内存使用量并防止内存泄漏。如果问题仍然存在，建议使用内存分析工具（如memory_profiler）来定位具体的内存泄漏点。
3. 问题重现材料 / Reproduction Materials
见描述
确认事项 / Verification
- 此问题未在过往Issue中被报告过 / This issue hasn't been reported before