Ouro GRPO with the shared rollout engine

Ouro GRPO with the shared rollout engine#

This recipe uses VIME’s existing Megatron optimizer, GRPO loss, gradient reduction and distributed checkpoints. Ouro keeps one physical copy of every decoder layer, all four sandwich norms, the inter-loop norm and the exit gate. The gate remains frozen during this externally budgeted recipe and participates in conversion and publication.

The companion is engine draft #78, the training-contract branch of 0z5a/vllm-rlt, based on ThinkFlowLab/vllm-rlt ecb1f8b505b7e831815b40aec3b4598619cca23a; tested companion commit f6cfefe5d7df8e31884b9396701245b934590bd6. The adapter calls public LLM.generate, start_weight_update, update_weights, finish_weight_update and get_weight_version. It uses the same contract for local and disaggregated prefill/decode (PD) engines. The previous private RLEngine prototype is superseded.

Launch#

Make VIME, the companion and an existing Megatron source checkout importable. The tested model is ByteDance/Ouro-1.4B at 574fa66cb8bf5abdc979642d01cf2b79b16bfab1. JSONL rows contain prompt, label and metadata.problem_id.

export PYTHONPATH="$PWD:/path/to/vllm-rlt:$MCORE"
GPUS_PER_NODE=1 bash examples/ouro/run.sh /models/Ouro-1.4B train.jsonl /runs/ouro 5 \
  --ouro-depths 2 3 4 --ouro-cuda-graphs --ouro-kv-blocks 128 \
  --rollout-max-response-len 16 --seq-length 512 --max-position-embeddings 2048 \
  --ouro-export-hf /runs/ouro/export

The launcher runs one trainer directly when GPUS_PER_NODE=1; multiple local trainer ranks use torchrun. --debug-train-only bypasses VIME’s server/router setup while this recipe still generates real online samples. TP=PP=CP=1, BF16, sequence parallelism and gradient accumulation fusion disabled are required. GRPO uses positive-temperature, full-vocabulary sampling without a reference, critic or KL loss. The small launch defaults are execution checks rather than a convergence protocol.

  • --ouro-depths 2 3 4 cycles a globally agreed decode budget by update ID.

  • --ouro-cuda-graphs enables recurrent graph replay across full publications.

  • --ouro-exit-threshold 0 exercises native early exit within the budget.

  • --ouro-async-scheduling uses the shared delayed-exit asynchronous path.

  • --ouro-speculative-tokens 3 --ouro-depths 4 uses the existing full-depth target policy with a K=2 draft. The upstream engine requires synchronous speculation.

  • --ouro-prefill-devices 0 --ouro-decode-devices 1 --ouro-max-num-seqs 1 selects real spawned PD workers. This pilot supports one trainer rank, which may share GPU 0 with P. PD and self-speculation cannot be combined in the upstream engine. Fixed batch shapes are required for the reported BF16 bitwise resume comparisons; dynamic BF16 batching can produce different scores and optimizer hashes.

  • --ouro-export-hf DIR exports through VIME’s ordinary HF saver.

  • --ouro-reward-function vime_plugins.ouro.reward.boxed_answer enables exact integer boxed-answer scoring; the default is Deepscaler.

  • --ouro-eval-data FILE --ouro-eval-prompts N --ouro-eval-interval M evaluates K=2/3/4 on separate prompts, or K=4 for speculation.

For a fresh-process resume, keep the same schedule, data, sampling and PD pool flags, increase the update count and add:

--ouro-resume --load /runs/ouro/checkpoint --override-opt-param-scheduler

The override is needed when extending the requested constant-LR horizon. Model, Adam, scheduler, RNG and the deterministic data/budget cursor are restored. ouro-plan.json rejects incompatible execution settings. Completed saves remove older iteration directories within that experiment’s save directory.

On the tested A100 container, default UCX CUDA IPC did not deliver GPU writes in an independent two-process probe. Per-command UCX_TLS=tcp,cuda_copy passed the probe and PD checks with the installed NIXL; this is host-staged transport. The recipe disables forced worker termination and closes successful PD workers through the normal stop/acknowledgment protocol.

Probability and conversion semantics#

Every prompt uses full-depth prefill. Subsequent decode inputs use each returned exit_depths value, including mixed early exits and speculative target depths. The differentiable provider replays those physical depths and retains the last exited KV at deeper planes. VIME’s existing per-sample forward kwargs carry the trace, including sequence padding and activation recomputation.

Outputs include selected-token processed logprobs, effective sampling parameters and the committed policy version. VIME’s existing loss applies the same rollout temperature. Each GRPO group validates its prompt ID, K, version and configuration hash. Failed or incomplete publication blocks generation until a newer full update completes. Publication changes tensors in place and invalidates cached KV.

Ouro is registered in the general HF-to-Megatron and Megatron-to-HF converters. Physical names are preserved after removing MCore/DDP wrapper prefixes; recurrent loops do not duplicate parameter tensors. The same HfWeightIteratorDirect is used for initial and per-update publication. Full HF exports include all 269 physical model tensors.

Metrics separate rollout, training, publication and whole-step time. Decode work counts physical layer/token evaluations, including speculative verification rows; trainer counts include padding and recomputation. These are work counters, not measured backward FLOPs. PD GPU-hours count distinct devices once.