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

[fix] gradient accumulation step alignment

pull/825/head
readlnh 2026-03-24 01:45:04 +01:00
parent f798445826
commit 9fe4343b7b
9 changed files with 9 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
trainer/train_distillation.py
View File

@ -91,7 +91,7 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st
scaler.scale(loss).backward() scaler.scale(loss).backward()
if (step + 1) % args.accumulation_steps == 0: if step % args.accumulation_steps == 0:
scaler.unscale_(optimizer) scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
scaler.step(optimizer) scaler.step(optimizer)

2
trainer/train_dpo.py
View File

@ -85,7 +85,7 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb=
scaler.scale(loss).backward() scaler.scale(loss).backward()
if (step + 1) % args.accumulation_steps == 0: if step % args.accumulation_steps == 0:
scaler.unscale_(optimizer) scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
scaler.step(optimizer) scaler.step(optimizer)

2
trainer/train_full_sft.py
View File

@ -36,7 +36,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
scaler.scale(loss).backward() scaler.scale(loss).backward()
if (step + 1) % args.accumulation_steps == 0: if step % args.accumulation_steps == 0:
scaler.unscale_(optimizer) scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)

2
trainer/train_grpo.py
View File

@ -148,7 +148,7 @@ def grpo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_token
loss = (policy_loss + aux_loss) / args.accumulation_steps # scalar loss = (policy_loss + aux_loss) / args.accumulation_steps # scalar
loss.backward() loss.backward()
if (step + 1) % args.accumulation_steps == 0: if step % args.accumulation_steps == 0:
if args.grad_clip > 0: if args.grad_clip > 0:
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
optimizer.step() optimizer.step()

2
trainer/train_lora.py
View File

@ -37,7 +37,7 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None):
scaler.scale(loss).backward() scaler.scale(loss).backward()
if (step + 1) % args.accumulation_steps == 0: if step % args.accumulation_steps == 0:
scaler.unscale_(optimizer) scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip) torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip)
scaler.step(optimizer) scaler.step(optimizer)

2
trainer/train_ppo.py
View File

@ -174,7 +174,7 @@ def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_sche
loss = (policy_loss + args.vf_coef * value_loss + args.kl_coef * kl_ref + aux_loss) / args.accumulation_steps # scalar loss = (policy_loss + args.vf_coef * value_loss + args.kl_coef * kl_ref + aux_loss) / args.accumulation_steps # scalar
loss.backward() loss.backward()
if (step + 1) % args.accumulation_steps == 0: if step % args.accumulation_steps == 0:
clip_grad_norm_(actor_model.parameters(), args.grad_clip) clip_grad_norm_(actor_model.parameters(), args.grad_clip)
clip_grad_norm_(critic_model.parameters(), args.grad_clip) clip_grad_norm_(critic_model.parameters(), args.grad_clip)
actor_optimizer.step() actor_optimizer.step()

2
trainer/train_pretrain.py
View File

@ -36,7 +36,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
scaler.scale(loss).backward() scaler.scale(loss).backward()
if (step + 1) % args.accumulation_steps == 0: if step % args.accumulation_steps == 0:
scaler.unscale_(optimizer) scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)

2
trainer/train_reason.py
View File

@ -56,7 +56,7 @@ def train_epoch(epoch, loader, iters, tokenizer, lm_config, start_step=0, wandb=
scaler.scale(loss).backward() scaler.scale(loss).backward()
if (step + 1) % args.accumulation_steps == 0: if step % args.accumulation_steps == 0:
scaler.unscale_(optimizer) scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
scaler.step(optimizer) scaler.step(optimizer)

2
trainer/train_spo.py
View File

@ -191,7 +191,7 @@ def spo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_tokeni
response_masks = completion_mask.float() # [B, R] response_masks = completion_mask.float() # [B, R]
rho = value_tracker.update(rewards, per_token_logps.detach(), response_masks) rho = value_tracker.update(rewards, per_token_logps.detach(), response_masks)
if (step + 1) % args.accumulation_steps == 0: if step % args.accumulation_steps == 0:
if args.grad_clip > 0: if args.grad_clip > 0:
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
optimizer.step() optimizer.step()