From 54ac9e446207377841eb28d4f9dbd6129972025e Mon Sep 17 00:00:00 2001
From: powermano <894773140@qq.com>
Date: Fri, 21 Feb 2025 17:09:35 +0800
Subject: [PATCH] add smart gradient accumulation

---
 train_full_sft.py | 6 +++++-
 train_pretrain.py | 6 +++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/train_full_sft.py b/train_full_sft.py
index 5185861..836d741 100644
--- a/train_full_sft.py
+++ b/train_full_sft.py
@@ -53,7 +53,11 @@ def train_epoch(epoch, wandb):
             loss += res.aux_loss
             loss = loss / args.accumulation_steps
 
-        scaler.scale(loss).backward()
+        if (step + 1) % args.accumulation_steps != 0:
+            with model.no_sync():  # Disable DDP gradient synchronization
+                scaler.scale(loss).backward()
+        else:
+            scaler.scale(loss).backward()
 
         if (step + 1) % args.accumulation_steps == 0:
             scaler.unscale_(optimizer)
diff --git a/train_pretrain.py b/train_pretrain.py
index a4465b6..29feb79 100644
--- a/train_pretrain.py
+++ b/train_pretrain.py
@@ -53,7 +53,11 @@ def train_epoch(epoch, wandb):
             loss += res.aux_loss
             loss = loss / args.accumulation_steps
 
-        scaler.scale(loss).backward()
+        if (step + 1) % args.accumulation_steps != 0:
+            with model.no_sync():  # Disable DDP gradient synchronization
+                scaler.scale(loss).backward()
+        else:
+            scaler.scale(loss).backward()
 
         if (step + 1) % args.accumulation_steps == 0:
             scaler.unscale_(optimizer)