Skip to content

Add synchronous fixed-budget Ouro rollouts for RL - #1

Draft
0z5a wants to merge 1 commit into
stack/ouro-rl-upstream-20260928from
feat/ouro-synchronous-rl
Draft

0z5a wants to merge 1 commit into
stack/ouro-rl-upstream-20260928from
feat/ouro-synchronous-rl

Conversation

@0z5a

@0z5a 0z5a commented Sep 28, 2026

Copy link
Copy Markdown
Owner

This adds a synchronous RL interface to the existing Ouro engine: one fixed depth for prefill and decode, selected-token full-vocabulary temperature-scaled logprobs, and full-weight publication while the engine is idle. It preserves the native model, scheduler, sampler and per-depth KV implementation. Publication advances the policy version only after copies finish; failed publication leaves generation disabled. Prefix reuse and asynchronous/speculative serving are excluded from this interface.

Base: ThinkFlowLab/vllm-rlt@ea680f7735eeaa357d2e09ba24355747496041dd. This fork draft is a companion to a VIME shared-parameter Megatron provider and global-budget GRPO recipe. It adds an in-process interface, not HTTP publication endpoints.

Validation on 2×H20:

Check Result
Existing engine/Ouro tests plus RL tests 55 passed
Tiny independent HF gradients, recompute, cached logprobs, publication Passed at K=2/3/4
Full pinned Ouro-1.4B FP32 cached selected-token logprobs vs dense training forward, K=4 Maximum absolute error 6.68e-6
Full BF16 cached selected-token logprobs vs SDPA training forward, K=2/3/4 Maximum absolute error 0.0478 / 0.1057 / 0.1046
Actual prefill block-tokens for an 8-token prompt, K=2/3/4 384 / 576 / 768
Actual decode block-tokens for seven cached steps, K=2/3/4 336 / 504 / 672

The BF16 paths are not bit-exact. Changing the HF causal-mask dispatch alone also changes BF16 logits; the matched SDPA training/HF paths are bit-exact and FP32 default-HF/dense training logits match. Real RL quality and time-to-reward improvement are not established by these numerical tests. The draft remains experimental pending the real short reward pilot.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
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.

1 participant