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)