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 pretrain training 

Same fused AdamW optimization as the full SFT script, using the
existing device_type variable so the CPU-only path is unaffected.

Co-authored-by: TRAE CLI <traecli@bytedance.com>
pull/843/head
neil.ge 2026-09-05 19:16:51 +08:00
parent 859c8ee2a3
commit 9cb07dd1ec
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_pretrain.py
View File

@ -136,7 +136,7 @@ if __name__ == "__main__":
train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_ds = PretrainDataset(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=(device_type == 'cuda'))
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========
start_epoch, start_step = 0, 0 start_epoch, start_step = 0, 0