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:c6a855929d5d4d1bd4a92a012fd4bb1cadf39766
Branches Tags
github:master
github:website
github:v2
github:minimind-v1
...
compare: github:cc82f38b75246ccb893dcd86b9617801996b4d17
Branches Tags
github:master
github:website
github:v2
github:minimind-v1

3 Commits
c6a855929d ... cc82f38b75

Author SHA1 Message Date
Neil Ge cc82f38b75
Merge 9cb07dd1ec into a3c7b01cc0 2026-09-10 21:11:19 -04: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_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

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