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

perf: enable fused AdamW in full SFT training 

Use PyTorch's fused AdamW implementation to reduce optimizer kernel
launch overhead and speed up the training step on CUDA devices.
Guard with device check so the CPU-only fallback path is unaffected.

Co-authored-by: TRAE CLI <traecli@bytedance.com>
pull/843/head
neil.ge 2026-09-05 18:37:04 +08:00
parent 6fc918beb6
commit 859c8ee2a3
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
trainer/train_full_sft.py
View File

@ -137,7 +137,7 @@ if __name__ == "__main__":
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, fused=('cuda' in args.device))
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========
start_epoch, start_step = 0, 0 start_epoch, start_step = 0, 0