pull/843/merge
Neil Ge 2026-09-29 16:15:42 +08:00 committed by GitHub
commit b3bd25fb19
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 2 additions and 2 deletions

View File

@ -128,7 +128,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

View File

@ -127,7 +127,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