Fix ordering of ppo epoch iteration

trlx/trainer/accelerate_base_trainer.py

@@ -542,21 +542,24 @@ def learn(self):  # noqa: C901
 
         # For each epoch
         for _ in range(self.config.train.epochs):
-            # For each batch
-            for mbs in MiniBatchIterator(self.train_dataloader, self.mb_size, self.num_mb):
-                # For each update per batch
-                for _ in range(self.n_updates_per_batch):
-                    # Note that whereas standard policy gradient methods perform one
-                    # gradient update per batch, PPO for example commonly performs
-                    # multiple gradient updates on the same batch of data.
-                    # https://arxiv.org/pdf/1707.06347.pdf
-                    forward_time = 0
-                    backward_time = 0
+            # For each ppo epoch
+            for _ in range(self.n_inner_epochs):
+                # Note that whereas standard policy gradient methods perform one
+                # gradient update per batch, PPO for example commonly performs
+                # multiple epochs of gradient updates on the same batch of data.
+                # https://arxiv.org/pdf/1707.06347.pdf
+
+                # We create a new dataloader (so new data ordering and shuffle) each inner epoch
+                train_dataloader = self.create_train_dataloader()
+                # For each batch
+                for minibatch in MiniBatchIterator(train_dataloader, self.mb_size, self.num_mb):
+                    forward_time = 0.0
+                    backward_time = 0.0
                     stats_accum = []
-                    for mb in mbs:
+                    for microbatch in minibatch:
                         with self._accumulate():
                             forward_time -= time()
-                            loss, stats = self.loss(mb)
+                            loss, stats = self.loss(microbatch)
                             forward_time += time()
                             backward_time -= time()
                             self.accelerator.backward(loss)
@@ -633,6 +636,12 @@ def learn(self):  # noqa: C901
             self.post_epoch_callback()
         tbar.close()
 
+    def create_train_dataloader(self, shuffle=True, accelerate_prepare=True):
+        dataloader = self.store.create_loader(self.config.train.batch_size, shuffle=shuffle)
+        if accelerate_prepare:
+            dataloader = self.accelerator.prepare_dataloader(dataloader)
+        return dataloader
+
     @abstractmethod
     def get_arch(self, config: TRLConfig):
         """Returns a specific wrapper of the decoder architecture"""

trlx/trainer/accelerate_ilql_trainer.py

@@ -158,18 +158,20 @@ def loss(self, batch: Union[ILQLBatch, ILQLSeq2SeqBatch]):
 
         return self.ilql.loss((logits, (qs, target_qs, vs)), batch)
 
+    def create_train_dataloader(self):
+        return self.accelerator.prepare(self.store.create_loader(self.config.train.batch_size))
+
     def prepare_learning(self):
-        train_dataloader = self.store.create_loader(self.config.train.batch_size)
+        self.train_dataloader = self.create_train_dataloader()
         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.accelerator.prepare(self.model, self.opt, eval_dataloader)
 
-        self.n_updates_per_batch = 1
+        self.n_inner_epochs = 1
         self.total_steps = self.config.train.epochs * len(self.train_dataloader)
         self.total_steps = min(self.total_steps, self.config.train.total_steps)
 

trlx/trainer/accelerate_ppo_trainer.py

@@ -220,16 +220,19 @@ def post_epoch_callback(self):
     def post_backward_callback(self):
         self.kl_ctl.update(self.mean_kl, n_steps=self.config.train.batch_size)
 
+    def create_train_dataloader(self):
+        return self.store.create_loader(self.config.train.batch_size)
+
     def prepare_learning(self):
         eval_dataloader = self.eval_pipeline.create_loader(self.config.method.chunk_size)
         self.eval_dataloader = self.accelerator.prepare_data_loader(eval_dataloader)
 
         self.make_experience(self.config.method.num_rollouts)
 
-        self.train_dataloader = self.store.create_loader(self.config.train.batch_size, shuffle=False)
+        self.train_dataloader = self.create_train_dataloader()
 
-        self.n_updates_per_batch = self.config.method.ppo_epochs
-        self.total_steps = self.config.train.epochs * self.n_updates_per_batch * len(self.train_dataloader)
+        self.n_inner_epochs = self.config.method.ppo_epochs
+        self.total_steps = self.config.train.epochs * self.n_inner_epochs * len(self.train_dataloader)
         self.total_steps = min(self.total_steps, self.config.train.total_steps)
 
     def add_prompt_pipeline(self, pipeline: PromptPipeline):

trlx/trainer/accelerate_sft_trainer.py

@@ -72,18 +72,20 @@ def loss(self, batch):
 
         return loss, stats
 
+    def create_train_dataloader(self):
+        return self.accelerator.prepare(self.store.create_loader(self.config.train.batch_size))
+
     def prepare_learning(self):
-        train_dataloader = self.store.create_loader(self.config.train.batch_size)
+        self.train_dataloader = self.create_train_dataloader()
         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.accelerator.prepare(self.model, self.opt, eval_dataloader)
 
-        self.n_updates_per_batch = 1
+        self.n_inner_epochs = 1
         self.total_steps = self.config.train.epochs * len(self.train_dataloader)
         self.total_steps = min(self.total_steps, self.config.train.total_steps)