Bug in accelerate_sft_trainer.py: Incorrect calculation of total_steps

[rockmagma02]

🐛 Describe the bug

I found a bug in your accelerate_sft_trainer.py file that I would like to report. Specifically, I noticed an issue with the prepare_learning() function. In the current implementation, self.total_steps is calculated by multiplying the number of epochs (self.config.train.epochs) with the length of the train_dataloader. However, since train_dataloader is assigned before accelerator preparation, it does not reflect the actual number of training steps taken if multiple GPUs are used.

As a result, the total_steps value ends up being larger than the true total number of training steps, leading to training ending prematurely.

To fix this issue, I recommend modifying the prepare_learning() function to calculate self.total_steps using the self.train_dataloader variable instead, which correctly reflects the number of training steps after accelerator preparation. Here is the suggested modification:

    def prepare_learning(self):
        train_dataloader = self.store.create_loader(self.config.train.batch_size)
        eval_dataloader = self.eval_pipeline.create_loader(self.config.train.batch_size)

        (
            self.model,
            self.opt,
            self.train_dataloader,
            self.eval_dataloader,
        ) = self.accelerator.prepare(self.model, self.opt, train_dataloader, eval_dataloader)

        self.n_updates_per_batch = 1
        self.total_steps = self.config.train.epochs * len(self.train_dataloader)
        self.total_steps = min(self.total_steps, self.config.train.total_steps)

I hope this helps! Let me know if you have any questions or if there's anything else I can assist you with.

Which trlX version are you using?

newest

Additional system and package information

No response

[jon-tow]

Thanks for reporting! Fixed with @reciprocated's #432 patch.