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

add smart gradient accumulation

pull/204/head
powermano 2025-02-21 17:09:35 +08:00
parent e34d4e9371
commit 54ac9e4462
2 changed files with 10 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
train_full_sft.py
View File

@ -53,6 +53,10 @@ def train_epoch(epoch, wandb):
loss += res.aux_loss
loss = loss / args.accumulation_steps
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:

4
train_pretrain.py
View File

@ -53,6 +53,10 @@ def train_epoch(epoch, wandb):
loss += res.aux_loss
loss = loss / args.accumulation_steps
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: