Baike.dev
Log in
Back to tool/Back to issues
#1358·OpenRLHF

Suboptimal forward dispatch order in make_experience causes unnecessary serialization

Author: aha-jansenCreated Sep 15, 2026Updated Sep 18, 2026

Problem

In RemoteExperienceMaker.make_experience() (openrlhf/trainer/ppo_utils/experience_maker.py), the current dispatch order is:

python
# Reward model
r_refs = self._dispatch_forward(
    self.reward_model_group,
    args.train.colocate_all,          # sync only if colocate_all
    ...
)

# Actor model
action_log_probs_ref = self._dispatch_forward(
    self.actor_model_group,
    args.train.colocate_all or args.train.colocate_actor_ref,  # sync if colocate_actor_ref!
    ...
)

# Critic model
if args.train.colocate_critic_reward and r_refs is not None:
    ray.get(r_refs)   # explicit barrier for reward
    ray.get(self.reward_model_group.async_run_method(method_name="empty_cache"))
value_ref = self._dispatch_forward(self.critic_model_group, ...)

# Reference model
base_action_log_probs_ref = self._dispatch_forward(
    self.initial_model_group,
    args.train.colocate_all or args.train.colocate_actor_ref,
    ...
)

where _dispatch_forward is:

python
def _dispatch_forward(self, group, sync_condition, **kwargs):
    ref = group.async_run_method_batch(method_name="forward", **kwargs)
    if sync_condition:
        ray.get(ref)                                          # blocks driver here
        ray.get(group.async_run_method(method_name="empty_cache"))
    return ref

This causes two distinct serialization issues:

Issue 1: colocate_actor_ref=True blocks Critic dispatch

When --train.colocate_actor_ref=True, _dispatch_forward calls ray.get(action_log_probs_ref) synchronously before returning. The driver is blocked until the Actor forward completes, so Critic dispatch is delayed even though Critic and Actor run on separate GPUs and could overlap.

Timeline with colocate_actor_ref=True:

  Reward GPU:  [=== reward forward ===]                         (async, not waited)
  Actor GPU:   [=== actor forward ===]                          (sync, driver waits)
  Critic GPU:                           [=== critic forward ===] ← delayed!
  Ref GPU:                                                       [=== ref forward ===] ← further delayed!

Issue 2: colocate_critic_reward=True blocks Reference dispatch

When --train.colocate_critic_reward=True (without colocate_all), the explicit ray.get(r_refs) before Critic also delays the Reference model dispatch, even though Reference runs on a separate GPU and could overlap with both Reward and Critic.

Timeline with colocate_critic_reward=True:

  Reward GPU:  [=== reward forward ===]
  Actor GPU:   [=== actor forward ===]  (async)
  Critic GPU:                            [=== critic forward ===]
  Ref GPU:                                                        [=== ref forward ===] ← delayed by reward barrier!

Proposed Fix

Separate the sync barriers from the dispatch calls. Dispatch all models first, then sync as required by colocation constraints:

python
# Batch 1: dispatch Reward + Actor immediately (always parallelizable)
r_refs = _dispatch_forward(reward_model_group, sync_condition=colocate_all, ...)
action_log_probs_ref = _dispatch_forward(actor_model_group, sync_condition=colocate_all, ...)

# Sync barriers (only when GPU memory must be freed before next dispatch)
if colocate_actor_ref and not colocate_all:
    ray.get(action_log_probs_ref)
    ray.get(actor_model_group.async_run_method("empty_cache"))

if colocate_critic_reward and not colocate_all and r_refs is not None:
    ray.get(r_refs)
    ray.get(reward_model_group.async_run_method("empty_cache"))

# Batch 2: dispatch Critic + Reference (parallelizable with each other)
value_ref = _dispatch_forward(critic_model_group, sync_condition=colocate_all, ...)
base_action_log_probs_ref = _dispatch_forward(reference_model_group, sync_condition=colocate_all, ...)

Expected timeline:

  Reward GPU:  [=== reward ===]
  Actor GPU:   [=== actor ===]    (barrier if colocate_actor_ref)
  Critic GPU:                      [=== critic ===]
  Ref GPU:                         [=== ref ===]    ← now overlaps with Critic!

Benefits

  • No changes to _dispatch_forward signature
  • Correctly handles colocate_actor_ref + colocate_critic_reward combined mode
  • Critic and Reference always overlap when running on separate GPUs

Affected Configurations

  • --train.colocate_actor_ref (without colocate_all)
  • --train.colocate_critic_reward (without colocate_all)
  • Both flags combined

Would appreciate confirmation on whether this analysis is correct before submitting a PR.

Source: OpenRLHF/OpenRLHF

View original on GitHubView discussion on GitHub