github
/
minimind-v
mirror of https://github.com/jingyaogong/minimind-v.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[fix] cuda memory

pull/107/head
jingyaogong 2025-12-01 16:35:05 +08:00
parent e03af46d12
commit 13c3a2c093
3 changed files with 8 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
trainer/train_pretrain_vlm.py
View File

@ -76,11 +76,12 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
clean_state_dict = { clean_state_dict = {
key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.') key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.')
} }
clean_state_dict = {k: v.half() for k, v in clean_state_dict.items()} # 半精度保存 clean_state_dict = {k: v.half().cpu() for k, v in clean_state_dict.items()} # 半精度保存并移到CPU
torch.save(clean_state_dict, ckp) torch.save(clean_state_dict, ckp)
vlm_checkpoint(vlm_config, weight=args.save_weight, model=model, optimizer=optimizer, vlm_checkpoint(vlm_config, weight=args.save_weight, model=model, optimizer=optimizer,
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler) epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
model.train() model.train()
del state_dict, clean_state_dict
del X, Y, loss_mask, pixel_values, res, loss del X, Y, loss_mask, pixel_values, res, loss
torch.cuda.empty_cache() torch.cuda.empty_cache()

3
trainer/train_sft_vlm.py
View File

@ -76,11 +76,12 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
clean_state_dict = { clean_state_dict = {
key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.') key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.')
} }
clean_state_dict = {k: v.half() for k, v in clean_state_dict.items()} # 半精度保存 clean_state_dict = {k: v.half().cpu() for k, v in clean_state_dict.items()} # 半精度保存并移到CPU
torch.save(clean_state_dict, ckp) torch.save(clean_state_dict, ckp)
vlm_checkpoint(vlm_config, weight=args.save_weight, model=model, optimizer=optimizer, vlm_checkpoint(vlm_config, weight=args.save_weight, model=model, optimizer=optimizer,
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler) epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
model.train() model.train()
del state_dict, clean_state_dict
del X, Y, loss_mask, pixel_values, res, loss del X, Y, loss_mask, pixel_values, res, loss
torch.cuda.empty_cache() torch.cuda.empty_cache()

5
trainer/trainer_utils.py
View File

@ -87,7 +87,7 @@ def vlm_checkpoint(vlm_config, weight='pretrain_vlm', model=None, optimizer=None
# 移除vision_encoder参数（不需要保存，因为是预训练的） # 移除vision_encoder参数（不需要保存，因为是预训练的）
clean_state_dict = {k: v for k, v in state_dict.items() if not k.startswith('vision_encoder.')} clean_state_dict = {k: v for k, v in state_dict.items() if not k.startswith('vision_encoder.')}
ckp_tmp = ckp_path + '.tmp' ckp_tmp = ckp_path + '.tmp'
torch.save({k: v.half() for k, v in clean_state_dict.items()}, ckp_tmp) torch.save({k: v.half().cpu() for k, v in clean_state_dict.items()}, ckp_tmp)
os.replace(ckp_tmp, ckp_path) os.replace(ckp_tmp, ckp_path)
wandb_id = None wandb_id = None
@ -119,6 +119,9 @@ def vlm_checkpoint(vlm_config, weight='pretrain_vlm', model=None, optimizer=None
resume_tmp = resume_path + '.tmp' resume_tmp = resume_path + '.tmp'
torch.save(resume_data, resume_tmp) torch.save(resume_data, resume_tmp)
os.replace(resume_tmp, resume_path) os.replace(resume_tmp, resume_path)
del state_dict, clean_state_dict, resume_data
gc.collect()
torch.cuda.empty_cache()
else: # 加载模式 else: # 加载模式
if os.path.exists(resume_path): if os.path.exists(resume_path):
ckp_data = torch.load(resume_path, map_location='cpu') ckp_data = torch.load(resume_path, map_location='cpu')