Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

Suboptimal forward dispatch order in make_experience causes unnecessary serialization

Open
  • Is something wrong?
#1,358 1 comment 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

@RainieLLM is already working on this.

Since Sep 18, 2026.

  • #1359 by @RainieLLM — open
  • #1376 by @Asthenia0412 — open

Assessment

Difficulty
3/5
Estimated time
1-2 days
Newbie friendliness
28/100
Issue type
Refactor
Clarity
Clearly specified
Activity status
Stale
Tech stack
python
Domain
performance

Research direction

Read openrlhf/trainer/ppo_utils/experience_maker.py around RemoteExperienceMaker.make_experience() and the _dispatch_forward helper, tracing how args.train.colocate_all, colocate_actor_ref and colocate_critic_reward decide where ray.get blocks. Note that linked PR #1359 already targets this, so check it before starting. Done means the reward/actor batch is dispatched before the sync barriers and critic/reference dispatch overlaps, with _dispatch_forward's signature unchanged and all three flag combinations still correct.

Written by the indexing model from the issue text.

Description

Problem

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

# 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:

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:

# 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.

Dominant language
Python
Stars
10.1k
Forks
1k
Avg merge
5h 19m
Merged PRs (30d)
4

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from OpenRLHF/OpenRLHF

All issues in OpenRLHF/OpenRLHF

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.