[skyrl-train] Fix num_training_steps for workers being set incorrectly (#873)

`num_training_steps` was being set to the number of training batch steps
rather than the number of optimizer (mini-batch) steps, causing learning
rate decay to progress too quickly if using a non-constant learning rate
scheduler.


renames to `num_training_batches` for clarity, since each training batch
can contain several optimizer steps.

Closes #872

Author: erictang000

skyrl-train/skyrl_train/trainer.py

@@ -497,21 +497,29 @@ def build_models(self, PolicyWorker, CriticWorker, RefWorker):
             else:
                 critic_model = None
 
+        policy_steps_per_train_batch = (
+            cfg.trainer.train_batch_size // cfg.trainer.policy_mini_batch_size * cfg.trainer.update_epochs_per_batch
+        )
+        critic_steps_per_train_batch = 0
+        if cfg.trainer.critic.model.path:
+            critic_steps_per_train_batch = (
+                cfg.trainer.train_batch_size // cfg.trainer.critic_mini_batch_size * cfg.trainer.update_epochs_per_batch
+            )
         if not cfg.trainer.placement.colocate_all:
             refs = []
             if ref_model is not None:
                 refs.extend(ref_model.async_init_model(cfg.trainer.ref.model.path))
             refs.extend(
                 policy_model.async_init_model(
                     cfg.trainer.policy.model.path,
-                    num_training_steps=self.total_training_steps,
+                    num_training_steps=self.total_training_steps * policy_steps_per_train_batch,
                 )
             )
             if cfg.trainer.critic.model.path:
                 refs.extend(
                     critic_model.async_init_model(
                         cfg.trainer.critic.model.path,
-                        num_training_steps=self.total_training_steps,
+                        num_training_steps=self.total_training_steps * critic_steps_per_train_batch,
                     )
                 )
             ray.get(refs)
@@ -523,7 +531,7 @@ def build_models(self, PolicyWorker, CriticWorker, RefWorker):
             ray.get(
                 policy_model.async_init_model(
                     cfg.trainer.policy.model.path,
-                    num_training_steps=self.total_training_steps,
+                    num_training_steps=self.total_training_steps * policy_steps_per_train_batch,
                 )
             )
             ray.get(policy_model.async_run_ray_method("pass_through", "_set_pad_token_id", self.tokenizer.pad_token_id))
@@ -532,7 +540,7 @@ def build_models(self, PolicyWorker, CriticWorker, RefWorker):
                 ray.get(
                     critic_model.async_init_model(
                         cfg.trainer.critic.model.path,
-                        num_training_steps=self.total_training_steps,
+                        num_training_steps=self.total_training_steps * critic_steps_per_train_batch,
                     )
                 )
                 critic_model.offload_to_cpu()