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