From a157b66ace14d3423a7b8e5519311f6f0d34fc08 Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Thu, 27 Nov 2025 19:37:39 +0800 Subject: [PATCH] [feat] del cuda memory --- trainer/train_pretrain_vlm.py | 3 +++ trainer/train_sft_vlm.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/trainer/train_pretrain_vlm.py b/trainer/train_pretrain_vlm.py index 32b155d..b4997b2 100644 --- a/trainer/train_pretrain_vlm.py +++ b/trainer/train_pretrain_vlm.py @@ -82,6 +82,9 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler) model.train() + del X, Y, loss_mask, pixel_values, res, loss + torch.cuda.empty_cache() + if __name__ == "__main__": parser = argparse.ArgumentParser(description="MiniMind-V Pretrain") diff --git a/trainer/train_sft_vlm.py b/trainer/train_sft_vlm.py index 253436e..180b4b6 100644 --- a/trainer/train_sft_vlm.py +++ b/trainer/train_sft_vlm.py @@ -82,6 +82,9 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler) model.train() + del X, Y, loss_mask, pixel_values, res, loss + torch.cuda.empty_cache() + if __name__ == "__main__": parser = argparse.ArgumentParser(description="MiniMind-V SFT")