From 13c3a2c093dfae42daeb6783304d8cafc3acf6f8 Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Mon, 1 Dec 2025 16:35:05 +0800 Subject: [PATCH] [fix] cuda memory --- trainer/train_pretrain_vlm.py | 3 ++- trainer/train_sft_vlm.py | 3 ++- trainer/trainer_utils.py | 5 ++++- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/trainer/train_pretrain_vlm.py b/trainer/train_pretrain_vlm.py index b4997b2..8a10f34 100644 --- a/trainer/train_pretrain_vlm.py +++ b/trainer/train_pretrain_vlm.py @@ -76,11 +76,12 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): clean_state_dict = { 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) vlm_checkpoint(vlm_config, weight=args.save_weight, model=model, optimizer=optimizer, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler) model.train() + del state_dict, clean_state_dict del X, Y, loss_mask, pixel_values, res, loss torch.cuda.empty_cache() diff --git a/trainer/train_sft_vlm.py b/trainer/train_sft_vlm.py index 180b4b6..1cec71e 100644 --- a/trainer/train_sft_vlm.py +++ b/trainer/train_sft_vlm.py @@ -76,11 +76,12 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): clean_state_dict = { 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) vlm_checkpoint(vlm_config, weight=args.save_weight, model=model, optimizer=optimizer, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler) model.train() + del state_dict, clean_state_dict del X, Y, loss_mask, pixel_values, res, loss torch.cuda.empty_cache() diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index 9d89475..bc813cf 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -87,7 +87,7 @@ def vlm_checkpoint(vlm_config, weight='pretrain_vlm', model=None, optimizer=None # 移除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' - 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) wandb_id = None @@ -119,6 +119,9 @@ def vlm_checkpoint(vlm_config, weight='pretrain_vlm', model=None, optimizer=None resume_tmp = resume_path + '.tmp' torch.save(resume_data, resume_tmp) os.replace(resume_tmp, resume_path) + del state_dict, clean_state_dict, resume_data + gc.collect() + torch.cuda.empty_cache() else: # 加载模式 if os.path.exists(resume_path): ckp_data = torch.load(resume_path, map_location='cpu')