Skip to content

feat: add async policy trainer primitives - #480

Open
VLa1111 wants to merge 1 commit into
inclusionAI:mainfrom
VLa1111:feat/add-async-policy-primitives
Open

feat: add async policy trainer primitives#480
VLa1111 wants to merge 1 commit into
inclusionAI:mainfrom
VLa1111:feat/add-async-policy-primitives

Conversation

@VLa1111

@VLa1111 VLa1111 commented Aug 14, 2026

Copy link
Copy Markdown

What does this PR do?

Adds the standalone building blocks for an experimental asynchronous policy trainer under areno.experimental.async_policy. The async trainer will overlaprollout/reward with training on single-node multi-GPU; these primitives — a bounded batch queue, batch metadata, and staleness control — are the first self-contained piece, landable and testable before any trainer integration. No algorithms are registered yet.

The change adds:

  • AsyncTrainBatch (batch.py): a slotted dataclass holding the materialized TrainSequences plus the rollout policy version and per-stage timing/size fields used for staleness control and metrics.
  • AsyncTrainBatchQueue (queue.py): a bounded FIFO queue with backpressure, per-operation put/get timeouts, close() to unblock waiters, and report_error() to propagate a worker failure to the next operation. It uses threading.Condition rather than queue.Queue because the stdlib queue has no close/unblock or error-propagation semantics.
  • staleness_delta / is_stale (staleness.py): helpers that compare a batch's rollout policy version against the current training policy version.
  • tests/test_async_policy_cpu.py: 19 CPU tests covering FIFO ordering, backpressure, timeouts, close, error propagation, and staleness.

A runnable example and an async-vs-sync evaluation report will accompany the async trainer itself; this PR is intentionally limited to the queue, batch, and staleness primitives.

Related issue

Fixes #487

  • feature: add experimental async policy trainer. This PR is the first of the series and lands the queue, batch, and staleness primitives.

Type of change

  • 🐛 Bug fix
  • ✨ New feature
  • 💥 Breaking change (public API / CLI behavior changes in a non-backward-compatible way)
  • 📝 Documentation update
  • ♻️ Refactoring
  • ⚡ Performance improvement
  • ✅ Test coverage improvement

How was it tested?

Local CPU validation on the final commit:

python -m pytest tests/test_async_policy_cpu.py -q
python -m pytest tests/test_algorithms_cpu.py tests/test_registry_discovery_cpu.py -q

Result: 19 passed and 5 passed respectively; existing algorithm/registry
discovery behavior is unaffected.

ruff check areno/experimental/async_policy tests/test_async_policy_cpu.py
ruff format --check areno/experimental/async_policy tests/test_async_policy_cpu.py

Result: all checks passed.

Hardware limitation: CPU-only validation (no GPU). This PR introduces no CUDA
execution path — the primitives are pure Python and gated behind
areno.experimental, so no GPU smoke test applies yet.

Checklist

  • The PR title summarizes the contribution.
  • Linked the related issue in the description (if any).
  • Existing tests pass (pytest tests/ -k cpu).
  • New behavior is covered by tests.
  • Described the test commands run and any hardware limitations.
  • Public API / CLI changes are additive and backward-compatible (see CONTRIBUTING.md).

Breaking change details

None — purely additive code under areno.experimental/; no areno/api/, CLI, or stable-surface changes.

Add the standalone building blocks for the experimental asynchronous
policy trainer under areno.experimental.async_policy.

- AsyncTrainBatch: batch metadata carrying the rollout policy version,
  timing, and size fields used for staleness control and metrics.
- AsyncTrainBatchQueue: a bounded queue with backpressure, put/get
  timeouts, close, and error propagation.
- staleness_delta/is_stale: policy-version staleness helpers.

All three are pure Python and covered by 19 CPU tests; algorithm
registration lands together with the async trainer itself.
@VLa1111
VLa1111 marked this pull request as draft August 14, 2026 02:49
@VLa1111
VLa1111 marked this pull request as ready for review August 14, 2026 02:51
@VLa1111

VLa1111 commented Aug 14, 2026

Copy link
Copy Markdown
Author

@xsuler Hi, I'd really appreciate a review on this PR when you have a chance.

Quick context so this PR doesn't look like a one-off: I'm applying to OSPP 2026 with AReno as my target project, and
my proposal is a native async policy trainer (async GRPO/GSPO) that overlaps rollout/reward with training on
single-node multi-GPU.

I'm landing that work as a sequence of small, independently-reviewable PRs so each one stays easy to audit. This first
PR is just the standalone building blocks — AsyncTrainBatch metadata, a bounded queue between rollout and training,
and staleness helpers. The trainer itself, a runnable example, and an async-vs-sync evaluation report will follow in
later PRs.

Let me know if this structure works for you, or if you'd prefer I bundle these primitives into the trainer PR instead.
Thanks!

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.

feature: add experimental async policy trainer

1 participant