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

6
train_full_sft.py
View File

@ -53,7 +53,11 @@ def train_epoch(epoch, wandb):
loss += res.aux_loss loss += res.aux_loss
loss = loss / args.accumulation_steps 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: if (step + 1) % args.accumulation_steps == 0:
scaler.unscale_(optimizer) scaler.unscale_(optimizer)

6
train_pretrain.py
View File

@ -53,7 +53,11 @@ def train_epoch(epoch, wandb):
loss += res.aux_loss loss += res.aux_loss
loss = loss / args.accumulation_steps 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: if (step + 1) % args.accumulation_steps == 0:
scaler.unscale_(optimizer) scaler.unscale_(optimizer)