From 9cb07dd1ec6fee0928cd3c1c24b0c98dc670a506 Mon Sep 17 00:00:00 2001 From: "neil.ge" Date: Sat, 5 Sep 2026 19:16:51 +0800 Subject: [PATCH] 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 --- trainer/train_pretrain.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 83a0382..fee8905 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -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