Make the per-step gc.collect() in RLEngine configurable

tests/rl/rl_cluster_test.py

@@ -341,6 +341,47 @@ def _create_test_rl_engine(
         actor=model, tokenizer=vocab, cluster_config=cluster_config
     )
 
+  @parameterized.named_parameters(
+      dict(testcase_name='enabled_by_default', gc_collect_after_weight_sync=True),
+      dict(testcase_name='disabled', gc_collect_after_weight_sync=False),
+  )
+  def test_sync_weights_gc_collect_after_weight_sync(
+      self, gc_collect_after_weight_sync
+  ):
+    mesh = Mesh(np.array(jax.devices()).reshape(-1, 1), ('fsdp', 'tp'))
+    cluster_config = rl_engine_lib.ClusterConfig(
+        role_to_mesh={
+            rl_engine_lib.Role.ACTOR: mesh,
+            rl_engine_lib.Role.REFERENCE: mesh,
+            rl_engine_lib.Role.ROLLOUT: mesh,
+        },
+        rollout_engine='vanilla',
+        offload_to_cpu=False,
+        gc_collect_after_weight_sync=gc_collect_after_weight_sync,
+        training_config=rl_engine_lib.RLTrainingConfig(
+            actor_optimizer=optax.sgd(1e-3),
+            eval_every_n_steps=1,
+            max_steps=10,
+            gradient_accumulation_steps=None,
+        ),
+        rollout_config=base_rollout.RolloutConfig(
+            max_tokens_to_generate=10,
+            max_prompt_length=256,
+            kv_cache_size=1024,
+            data_type=jnp.bfloat16,
+        ),
+    )
+    vocab = tc.MockVocab()
+    model = tc.ToyTransformer(
+        config=tc.ModelConfig(vocab_size=vocab.GetPieceSize()), rngs=nnx.Rngs(0)
+    )
+    rl_engine = rl_engine_lib.RLEngine(
+        actor=model, tokenizer=vocab, cluster_config=cluster_config
+    )
+    with mock.patch.object(rl_engine_lib.gc, 'collect') as mock_collect:
+      rl_engine.sync_weights()
+    self.assertEqual(mock_collect.called, gc_collect_after_weight_sync)
+
   def test_init_engine_invalid_engine_string(self):
     with self.assertRaisesRegex(
         ValueError, '`cluster_config.rollout_engine` should be one of'

tunix/common/configs.py

@@ -390,6 +390,11 @@ class ClusterConfig:
       Alternatively, if a subclass of `BaseRollout` is provided, it will be used
       as the rollout engine.
     offload_to_cpu: Whether to offload models to CPU at each step..
+    gc_collect_after_weight_sync: Whether to run `gc.collect()` after each
+      weight sync. This promptly releases the host-side references the sync
+      leaves behind, which matters when host memory is tight. On a colocated
+      setup a full host garbage collection on every step is a pure stall
+      (about one second per step observed), so it can be disabled.
     training_config: RL training config.
     rollout_config: Rollout config. It may be different for different modes,
       e.g. TRAIN vs EVAL.
@@ -407,6 +412,7 @@ class ClusterConfig:
   role_to_logical_axis_rule: dict[Role, flax.typing.LogicalRules] | None = None
   rollout_engine: str | type["base_rollout.BaseRollout"] = "vanilla"
   offload_to_cpu: bool = False
+  gc_collect_after_weight_sync: bool = True
 
   training_config: RLTrainingConfig
   rollout_config: dict[Mode, RolloutConfig] | RolloutConfig

tunix/rl/rl_cluster.py

@@ -1156,8 +1156,11 @@ def get_actor_per_token_logps(
       actor_per_token_logps = jnp.concatenate(outs, axis=0)
       if not anchor_on_device:
         del anchor_policy_state
-      gc.collect()
       if actor_trainer_state_on_device and self.cluster_config.offload_to_cpu:
+        # Release the host-side references the log-prob pass leaves behind
+        # before the model moves back; without offloading there is nothing
+        # to reclaim and a full collection is a pure stall.
+        gc.collect()
         self._put_model_on_memory_kind(
             self.actor_trainer.model, self._default_memory_kind
         )
@@ -1179,7 +1182,8 @@ def sync_weights(self):
       )
       src_filtered_params = nnx.state(self.actor_trainer.model, filter_types)
       self.rollout.update_params(src_filtered_params, filter_types)
-      gc.collect()
+      if self.cluster_config.gc_collect_after_weight_sync:
+        gc.collect()
       # The anchor policy state is snapshotted from actor_trainer.model.
       self._anchor_policy_state = rl_utils.put_params_on_memory_kind(
           nnx.state(self.actor_trainer.model), "pinned_host"