Skip to content

Use one full-window running baseline for BCO gradient accumulation - #7267

Closed
gss10282025 wants to merge 1 commit into
huggingface:mainfrom
gss10282025:codex/fix-08-20260918
Closed

gss10282025 wants to merge 1 commit into
huggingface:mainfrom
gss10282025:codex/fix-08-20260918

Conversation

@gss10282025

Copy link
Copy Markdown

What does this PR do?

Collect detached rewards for the whole optimizer window, update RunningMoments once, and let every microbatch read the same updated mean. This keeps the unsplit BCO objective while allowing gradient accumulation.

The pending microbatches are collected in get_batch_samples, their rewards are computed without gradients, and running-statistic writes are suppressed during the loss calls. Evaluation stays read-only. Freezing the old mean and committing after backward would be a different objective, so that is not what this does.

At TRL 1.12.0 the one-vs-four-microbatch parameter-update difference drops from 0.06736848 relative L2 to zero; fixed-baseline and DPO controls agree across layouts; reproduced on a second host. The collection pass restores RNG state and registered model buffers. Full Trainer/GPU execution and multi-rank aggregation still need validation.

Before submitting

AI writing disclosure

  • No AI usage: the PR was written entirely by a human.
  • AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
  • AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

Anyone in the community is free to review the PR once the tests have passed.

Checks on this branch

Added tests/experimental/test_bco_window.py. Its isolated CPU execution passed using the production prefetch/loss methods and RunningMoments: full, equal and uneven windows agree in baseline and gradient; the prepass restores buffers and RNG; evaluation does not update the running state. This isolates the window logic and does not run a complete Trainer job. Ruff, Python syntax and git diff --check pass.

Fixes #7265. The separate removal proposal #7136 is still open; this draft addresses the behavior while BCO remains available.

Signed-off-by: sandel <326855337+gss10282025@users.noreply.github.com>

Assisted-by: OpenAI Codex
@qgallouedec

Copy link
Copy Markdown
Member

Closing in favour of #7136, which removes trl.experimental.bco.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] BCO uses different running baselines within one gradient-accumulation window

2 participants