From 914d8ecb24bd3893aff58cdca65cb5ce6f398972 Mon Sep 17 00:00:00 2001 From: Tianyu Gu Date: Thu, 10 Sep 2026 23:30:25 +0000 Subject: [PATCH] Make the gc.collect() after weight sync configurable RLEngine ran a full host gc.collect() on every training step: after each weight sync and after the actor log-prob pass. Those collections exist to promptly release host-side references when models are offloaded to CPU or host memory is tight; on a colocated setup with no offloading a full collection is a pure stall of about one second per step. - The gc.collect() after the actor log-prob pass now runs only inside the offload_to_cpu branch, where the references it reclaims actually exist, so it needs no knob. - The one after weight sync is behind a new ClusterConfig field, gc_collect_after_weight_sync (default True, preserving today's behavior). --- tests/rl/rl_cluster_test.py | 41 +++++++++++++++++++++++++++++++++++++ tunix/common/configs.py | 9 ++++++++ tunix/rl/rl_cluster.py | 8 ++++++-- 3 files changed, 56 insertions(+), 2 deletions(-) diff --git a/tests/rl/rl_cluster_test.py b/tests/rl/rl_cluster_test.py index 361e0c0c01..5bf78771b3 100644 --- a/tests/rl/rl_cluster_test.py +++ b/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' diff --git a/tunix/common/configs.py b/tunix/common/configs.py index 6aa0ce0671..9640d23abf 100644 --- a/tunix/common/configs.py +++ b/tunix/common/configs.py @@ -390,6 +390,14 @@ 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. A weight sync creates another copy of the weights on HBM, + and that copy is only freed once Python's garbage collector releases + it; if the collector does not run, the copies accumulate and the run + can OOM over time. Collecting right after the sync frees the copy + promptly. A full collection costs about one second per step on a + colocated setup, so it can be disabled for performance, provided the + memory headroom is there. training_config: RL training config. rollout_config: Rollout config. It may be different for different modes, e.g. TRAIN vs EVAL. @@ -407,6 +415,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 diff --git a/tunix/rl/rl_cluster.py b/tunix/rl/rl_cluster.py index ee35b516d7..2b01c874a9 100644 --- a/tunix/rl/rl_cluster.py +++ b/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"