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