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

Compare commits

base: github:b3bd25fb19fcf2a25c6381a915c5b9d8bad20902
Branches Tags
github:master
github:website
github:v2
github:minimind-v1
...
compare: github:40efec2c24e3d55e3e0ef10881dd07f94e10df94
Branches Tags
github:master
github:website
github:v2
github:minimind-v1

3 Commits
b3bd25fb19 ... 40efec2c24

Author SHA1 Message Date
Neil Ge 40efec2c24
Merge 9cb07dd1ec into e41bf7cacf 2026-09-20 19:49:53 +08:00
neil.ge 9cb07dd1ec 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>
 2026-09-05 19:16:51 +08:00
neil.ge 859c8ee2a3 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>
 2026-09-05 18:37:04 +08:00
2 changed files with 2 additions and 2 deletions
Show Stats 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_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
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恢复状态 ==========
start_epoch, start_step = 0, 0

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_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
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恢复状态 ==========
start_epoch, start_step = 0, 0