[bugfix] initialize rollout manager first to calculate num_rollout (#473)

Author: zhuzilin

slime/ray/placement_group.py

@@ -182,4 +182,7 @@ def create_rollout_manager(args, pg, wandb_run_id):
         args.num_rollout = num_rollout_per_epoch * args.num_epoch
     assert args.num_rollout > 0
 
+    if args.offload:
+        ray.get(rollout_manager.offload.remote())
+
     return rollout_manager, num_rollout_per_epoch

train.py

@@ -16,16 +16,16 @@ def train(args):
 
         _TensorboardAdapter(args)
 
-    # create the actor and critic models
-    actor_model, critic_model = create_training_models(args, pgs, wandb_run_id=wandb_run_id)
-
     # create the rollout manager, with sglang engines inside.
+    # need to initialize rollout manager first to calculate num_rollout
     rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"], wandb_run_id=wandb_run_id)
 
+    # create the actor and critic models
+    actor_model, critic_model = create_training_models(args, pgs, wandb_run_id=wandb_run_id)
+
     actor_model.set_rollout_manager(rollout_manager)
 
     if args.offload:
-        ray.get(rollout_manager.offload.remote())
         ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]))
 
     # always update weight first so that sglang has the loaded weights from training.

train_async.py

@@ -11,12 +11,13 @@ def train(args):
     pgs = create_placement_groups(args)
     wandb_run_id = init_wandb_primary(args)
 
-    # create the actor and critic models
-    actor_model, critic_model = create_training_models(args, pgs, wandb_run_id=wandb_run_id)
-
     # create the rollout manager, with sglang engines inside.
+    # need to initialize rollout manager first to calculate num_rollout
     rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"], wandb_run_id=wandb_run_id)
 
+    # create the actor and critic models
+    actor_model, critic_model = create_training_models(args, pgs, wandb_run_id=wandb_run_id)
+
     actor_model.set_rollout_manager(rollout_manager)
 
     # always update weight first so that sglang has the loaded weights from training.