Merge 9cb07dd1ec into 7a9137d2e9
commit
08e1bdb9d8
|
|
@ -137,7 +137,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
|
||||||
|
|
|
||||||
|
|
@ -136,7 +136,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
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue