fix: fix mcore train_iters in grpo (#1383)

Signed-off-by: Yuki Huang <yukih@nvidia.com>

Author: yuki-97

nemo_rl/algorithms/grpo.py

@@ -467,7 +467,10 @@ def setup(
 
     if policy_config.get("megatron_cfg", {}).get("enabled", False):
         ## NOTE: this is equal to the total number of scheduler steps
-        total_train_iters = min(grpo_config["max_num_steps"], len(dataloader))
+        total_train_iters = min(
+            grpo_config["max_num_steps"],
+            grpo_config["max_num_epochs"] * len(dataloader),
+        )
         policy_config["megatron_cfg"]["train_iters"] = total_train_iters
 
     policy = Policy(