From d25500d363460d8b518f6b03643a7d3cb0448594 Mon Sep 17 00:00:00 2001 From: readlnh Date: Tue, 24 Mar 2026 01:45:04 +0100 Subject: [PATCH] [fix] gradient accumulation step alignment --- trainer/train_distillation.py | 2 +- trainer/train_dpo.py | 2 +- trainer/train_full_sft.py | 2 +- trainer/train_grpo.py | 2 +- trainer/train_lora.py | 2 +- trainer/train_ppo.py | 2 +- trainer/train_pretrain.py | 2 +- trainer/train_reason.py | 2 +- trainer/train_spo.py | 2 +- 9 files changed, 9 insertions(+), 9 deletions(-) diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index 5cc6269..e0d160f 100644 --- a/trainer/train_distillation.py +++ b/trainer/train_distillation.py @@ -91,7 +91,7 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st scaler.scale(loss).backward() - if (step + 1) % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) scaler.step(optimizer) diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index f1e567d..977953d 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -85,7 +85,7 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb= scaler.scale(loss).backward() - if (step + 1) % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) scaler.step(optimizer) diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index cc59cc7..16fadea 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -36,7 +36,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): scaler.scale(loss).backward() - if (step + 1) % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 5e63779..c7c122d 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -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.backward() - if (step + 1) % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0: if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) optimizer.step() diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 90bbf5a..fb1b1b1 100644 --- a/trainer/train_lora.py +++ b/trainer/train_lora.py @@ -37,7 +37,7 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None): scaler.scale(loss).backward() - if (step + 1) % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip) scaler.step(optimizer) diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index 88c11bc..4925651 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -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.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_(critic_model.parameters(), args.grad_clip) actor_optimizer.step() diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index a8ad97f..1b6347d 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -36,7 +36,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): scaler.scale(loss).backward() - if (step + 1) % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) diff --git a/trainer/train_reason.py b/trainer/train_reason.py index a9b2d55..c2bae0f 100644 --- a/trainer/train_reason.py +++ b/trainer/train_reason.py @@ -56,7 +56,7 @@ def train_epoch(epoch, loader, iters, tokenizer, lm_config, start_step=0, wandb= scaler.scale(loss).backward() - if (step + 1) % args.accumulation_steps == 0: + if step % args.accumulation_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) scaler.step(optimizer) diff --git a/trainer/train_spo.py b/trainer/train_spo.py index 219ddba..595bc50 100755 --- a/trainer/train_spo.py +++ b/trainer/train_spo.py @@ -191,7 +191,7 @@ def spo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_tokeni response_masks = completion_mask.float() # [B, R] 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: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) optimizer.step()