From 859c8ee2a32ae968519bd49802e3ffb71f5089ff Mon Sep 17 00:00:00 2001 From: "neil.ge" Date: Sat, 5 Sep 2026 18:37:04 +0800 Subject: [PATCH] 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