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