gnh1201/Open-Llama
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

support multiple epochs

Browse Source
This commit is contained in:
LiangSong 2023-05-03 00:02:01 +08:00
parent f05e929aad
commit c2184c6dd1
1 changed files with 42 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
solver/trainer.py
View File

@ -2,7 +2,7 @@
Author: LiangSong(sl12160010@gmail.com)
Date: 2023-04-24 20:05:21
LastEditors: LiangSong(sl12160010@gmail.com)
LastEditTime: 2023-04-29 21:59:51
LastEditTime: 2023-05-02 23:55:37
FilePath: /Open-Llama/solver/trainer.py
Description:
@ -109,9 +109,12 @@ class Trainer:
logging.warn("No ckpt found in {}".format(self.work_dir))
if self.global_step > 0:
skip_steps = self.global_step * self.gradient_accumulation_steps
self.train_loader = self.accelerator.skip_first_batches(
logging.warn("Skiped {} steps.".format(skip_steps))
self.train_loader_skiped = self.accelerator.skip_first_batches(
self.train_loader, num_batches=skip_steps
)
else:
self.train_loader_skiped = self.train_loader
self.accelerator.wait_for_everyone()
def train_step(self, batch):
@ -129,7 +132,16 @@ class Trainer:
self.get_lr_scheduler()
self.prepare()
self.start_time = time.time()
for self.data_step, batch in enumerate(self.train_loader):
self.epoch = 0
self.data_step = 0
while True:
if self.data_step >= self.config["train"]["num_training_steps"]:
break
if self.epoch == 0:
train_loader = self.train_loader_skiped
else:
train_loader = self.train_loader
for batch in train_loader:
# end training
if self.data_step >= self.config["train"]["num_training_steps"]:
break
@ -158,6 +170,7 @@ class Trainer:
# save state
if self.data_step % self.save_interval == 0 and self.data_step > 0:
self.accelerator.save_state(self.work_dir)
self.data_step += 1
wandb.finish()
def log(self, losses):