obscura/minimind
1
0
Fork 0
mirror of https://github.com/jingyaogong/minimind.git synced 2026-09-05 11:09:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

[feat] remove empty_cache

Browse Source
This commit is contained in:
jingyaogong
2025-12-27 07:14:36 +08:00
parent 11b962da06
commit 7eae14f3ce
10 changed files with 11 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
trainer/train_distill_reason.py
+1 -2
View File
@@ -65,7 +65,6 @@ def train_epoch(epoch, loader, iters, tokenizer, lm_config, start_step=0, wandb=
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
torch.cuda.empty_cache()
if step % args.log_interval == 0 or step == iters - 1:
spend_time = time.time() - start_time
@@ -103,7 +102,7 @@ if __name__ == "__main__":
parser.add_argument("--learning_rate", type=float, default=1e-6, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=1, help="数据加载线程数")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
trainer/train_distillation.py
+1 -2
View File
@@ -96,7 +96,6 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
torch.cuda.empty_cache()
if step % args.log_interval == 0 or step == iters - 1:
spend_time = time.time() - start_time
@@ -141,7 +140,7 @@ if __name__ == "__main__":
parser.add_argument("--learning_rate", type=float, default=5e-6, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=1, help="数据加载线程数")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
trainer/train_dpo.py
+1 -2
View File
@@ -90,7 +90,6 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb=
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
torch.cuda.empty_cache()
if step % args.log_interval == 0 or step == iters - 1:
spend_time = time.time() - start_time
@@ -129,7 +128,7 @@ if __name__ == "__main__":
parser.add_argument("--learning_rate", type=float, default=4e-8, help="初始学习率（建议<=5e-8避免遗忘）")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=1, help="数据加载线程数")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
trainer/train_full_sft.py
+2 -3
View File
@@ -52,7 +52,6 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
scaler.update()
optimizer.zero_grad(set_to_none=True)
torch.cuda.empty_cache()
if step % args.log_interval == 0 or step == iters - 1:
spend_time = time.time() - start_time
@@ -91,11 +90,11 @@ if __name__ == "__main__":
parser.add_argument("--learning_rate", type=float, default=5e-7, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=1, help="数据加载线程数")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔")
parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
parser.add_argument('--max_seq_len', default=340, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
trainer/train_grpo.py
+1 -4
View File
@@ -149,7 +149,6 @@ def grpo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_token
optimizer.step()
scheduler.step()
optimizer.zero_grad()
torch.cuda.empty_cache()
if step % args.log_interval == 0 or step == iters:
policy_loss_val = loss.item()
@@ -183,8 +182,6 @@ def grpo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_token
del prompt_inputs, outputs, completion_ids, per_token_logps, ref_per_token_logps
del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask
torch.cuda.empty_cache()
gc.collect()
if __name__ == "__main__":
@@ -196,7 +193,7 @@ if __name__ == "__main__":
parser.add_argument("--learning_rate", type=float, default=8e-8, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=1, help="数据加载线程数")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--log_interval", type=int, default=1, help="日志打印间隔")
trainer/train_lora.py
+1 -2
View File
@@ -53,7 +53,6 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None):
scaler.update()
optimizer.zero_grad(set_to_none=True)
torch.cuda.empty_cache()
if step % args.log_interval == 0 or step == iters - 1:
spend_time = time.time() - start_time
@@ -85,7 +84,7 @@ if __name__ == "__main__":
parser.add_argument("--learning_rate", type=float, default=1e-4, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=1, help="数据加载线程数")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--log_interval", type=int, default=10, help="日志打印间隔")
trainer/train_ppo.py
+1 -3
View File
@@ -179,7 +179,6 @@ def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_sche
critic_scheduler.step()
actor_optimizer.zero_grad()
critic_optimizer.zero_grad()
torch.cuda.empty_cache()
if is_main_process():
response_ids = gen_out[:, enc.input_ids.shape[1]:]
@@ -237,7 +236,6 @@ def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_sche
del enc, gen_out, responses_text, rewards, full_mask, values_seq, values, advantages
del logits, labels, logp_tokens, final_mask, actor_logp, old_logits, old_logp, ref_logits, ref_logp
del kl, kl_ref, ratio, surr1, surr2, policy_loss, value_loss, loss
torch.cuda.empty_cache()
if __name__ == "__main__":
@@ -250,7 +248,7 @@ if __name__ == "__main__":
parser.add_argument("--critic_learning_rate", type=float, default=8e-8, help="Critic学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=1, help="数据加载线程数")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--log_interval", type=int, default=1, help="日志打印间隔")
trainer/train_pretrain.py
+2 -3
View File
@@ -52,7 +52,6 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
scaler.update()
optimizer.zero_grad(set_to_none=True)
torch.cuda.empty_cache()
if step % args.log_interval == 0 or step == iters - 1:
spend_time = time.time() - start_time
@@ -90,11 +89,11 @@ if __name__ == "__main__":
parser.add_argument("--learning_rate", type=float, default=5e-4, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=1, help="数据加载线程数")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔")
parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
parser.add_argument('--max_seq_len', default=340, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
trainer/train_spo.py
+1 -4
View File
@@ -192,7 +192,6 @@ def spo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_tokeni
optimizer.step()
scheduler.step()
optimizer.zero_grad()
torch.cuda.empty_cache()
if step % args.log_interval == 0 or step == iters:
policy_loss_val = loss.item()
@@ -231,8 +230,6 @@ def spo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_tokeni
del prompt_inputs, outputs, completion_ids, per_token_logps, ref_per_token_logps
del completions, rewards, advantages, completion_mask, baselines, response_masks
torch.cuda.empty_cache()
gc.collect()
if __name__ == "__main__":
@@ -244,7 +241,7 @@ if __name__ == "__main__":
parser.add_argument("--learning_rate", type=float, default=1e-7, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=1, help="数据加载线程数")
parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=4, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--log_interval", type=int, default=1, help="日志打印间隔")
trainer/trainer_utils.py
-1
View File
@@ -88,7 +88,6 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc
torch.save(resume_data, resume_tmp)
os.replace(resume_tmp, resume_path)
del state_dict, resume_data
gc.collect()
torch.cuda.empty_cache()
else: # 加载模式
if os.path.exists(resume_path):