Repository navigation
[Training][Rollout] Integrate Ouro shared layers and policy updates - #459
Conversation
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces Ouro shared-layer training documentation, validation reports, and associated utility scripts for model weight conversion, checkpoint comparison, and cleanup. The feedback on the newly added validation scripts focuses on improving portability and compatibility, specifically by avoiding hardcoded absolute paths and virtual environment python paths, and ensuring Path objects are explicitly converted to strings when passed to PyTorch distributed checkpoint APIs.
| import shutil | ||
| from pathlib import Path | ||
|
|
||
| root = Path("/workspace/0z5a/work/ouro-vime-contract-20261001") |
There was a problem hiding this comment.
The path root is hardcoded to an absolute path /workspace/0z5a/work/ouro-vime-contract-20261001. This makes the script non-portable and prone to failure in other environments. Consider dynamically resolving the repository root relative to the script's location.
| root = Path("/workspace/0z5a/work/ouro-vime-contract-20261001") | |
| root = Path(__file__).resolve().parents[3] |
|
|
||
|
|
||
| def hashes(folder): | ||
| metadata = dcp.FileSystemReader(folder).read_metadata() |
There was a problem hiding this comment.
dcp.FileSystemReader expects a string path. Passing a Path object directly may cause compatibility issues depending on the PyTorch version. It is safer to convert the Path object to a string.
| metadata = dcp.FileSystemReader(folder).read_metadata() | |
| metadata = dcp.FileSystemReader(str(folder)).read_metadata() |
| for name, value in metadata.state_dict_metadata.items() | ||
| } | ||
| planner = DefaultLoadPlanner(flatten_state_dict=False) | ||
| dcp.load(state, checkpoint_id=folder, planner=planner) |
There was a problem hiding this comment.
The checkpoint_id argument expects a string path. Passing a Path object directly may cause compatibility issues depending on the PyTorch version. It is safer to convert the Path object to a string.
| dcp.load(state, checkpoint_id=folder, planner=planner) | |
| dcp.load(state, checkpoint_id=str(folder), planner=planner) |
| @@ -0,0 +1,16 @@ | |||
| #!/bin/bash | |||
| set -euo pipefail | |||
| root=/workspace/0z5a/work/ouro-vime-contract-20261001 | |||
There was a problem hiding this comment.
The root path is hardcoded to /workspace/0z5a/work/ouro-vime-contract-20261001. This makes the script non-portable. Consider dynamically resolving the repository root relative to the script's location.
| root=/workspace/0z5a/work/ouro-vime-contract-20261001 | |
| root="$(cd "$(dirname "${BASH_SOURCE[0]}")"/../../.. && pwd)" |
| /venv/main/bin/python compare_checkpoint.py train-full-pd-eager/checkpoint/iter_0000002 \ | ||
| train-full-pd-graph/checkpoint/iter_0000002 full-pd-checkpoint.json >full-pd-compare.log 2>&1 | ||
| /venv/main/bin/python check_hf_export.py train-full-pd-graph/checkpoint/iter_0000002 \ | ||
| train-full-pd-graph/export full-pd-hf.json >full-pd-hf.log 2>&1 | ||
| /venv/main/bin/python cleanup_a100.py full |
There was a problem hiding this comment.
The virtualenv python path /venv/main/bin/python is hardcoded. This will fail on systems where the virtualenv is located elsewhere or named differently. Consider using the active environment's python directly.
| /venv/main/bin/python compare_checkpoint.py train-full-pd-eager/checkpoint/iter_0000002 \ | |
| train-full-pd-graph/checkpoint/iter_0000002 full-pd-checkpoint.json >full-pd-compare.log 2>&1 | |
| /venv/main/bin/python check_hf_export.py train-full-pd-graph/checkpoint/iter_0000002 \ | |
| train-full-pd-graph/export full-pd-hf.json >full-pd-hf.log 2>&1 | |
| /venv/main/bin/python cleanup_a100.py full | |
| python compare_checkpoint.py train-full-pd-eager/checkpoint/iter_0000002 \ | |
| train-full-pd-graph/checkpoint/iter_0000002 full-pd-checkpoint.json >full-pd-compare.log 2>&1 | |
| python check_hf_export.py train-full-pd-graph/checkpoint/iter_0000002 \ | |
| train-full-pd-graph/export full-pd-hf.json >full-pd-hf.log 2>&1 | |
| python cleanup_a100.py full |
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
main reorganised the docs into topic index pages and builds them with -W in CI. The sync merge took main's index pages, so two 0z5a-branch pages are no longer in any toctree, and Sphinx reports "document isn't included in any toctree" for each of them in en and zh. Under -W that fails the docs lab browser step. - advanced/draft-feature-collection.md (vllm-project#454) is listed in rollout/index.rst right after speculative-decoding.md, where it sat in 0z5a-branch's Advanced Features toctree. - The synced examples/ouro README (vllm-project#459) is listed in examples/workflows/index.rst next to the DSpark example. 0z5a-branch never listed it, but its pipeline did not build the docs with -W. `docs/build.sh {en,zh} -W --keep-going` now passes for both languages. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Incremental review: Files changed.
Base:
0z5a-branchat791f101e20db377b28e15f12a00863baaea2605f.Purpose
Integrate Ouro's physical shared layers with VIME's existing Megatron GRPO optimizer, loss, distributed checkpoint and HF conversion paths. Rollouts use the companion engine's public generation and transactional full-weight publication contract for local and real prefill/decode engines.
The provider replays full-depth prompt prefill and each decode token's actual execution depth, including the last-exited KV state at deeper planes. Physical parameters retain their checkpoint names across MCore/DDP wrappers; recurrent loops do not duplicate parameter storage.
examples/ouro/run.shprovides a multi-iteration launch and fresh-process resume recipe.Target:
0z5a-branchat791f101e20db377b28e15f12a00863baaea2605f. The submitted diff contains the implementation and user recipe: 17 files, 870 additions and 28 deletions. Added tests, validation harnesses, raw logs and coordination artifacts have been removed. Their results are summarized here.Coordination and dependency
Companion ThinkFlowLab/vllm-rlt#78 remains a draft. Following Minhao's PR 1 design, VIME uses full four-loop prompt prefill, actual decode-depth replay, selected-token processed scores and returned effective seeds. The reduced-prefill prototype is superseded.
The reference currently exposes
log_probsand per-requestlogprobs_mode. The proposed upstream API useslogprobsand engine-levelraw_logprobs/processed_logprobs, with weight publication in PR 2. Adapting to and testing Minhao's exact implementation remain pending availability of that branch.Validation
After porting to the target branch, the mixed-depth probability/recomputed-gradient CPU check and physical-parameter HF round trip pass on the unchanged Thor runtime. Changed Python files pass Ruff, Black and isort checks. Five full-model K=2/3/4/early/delayed cases exactly replay tokens and selected scores from returned entropy seeds.
Actual-depth replay reduces early-exit actor/inference score MAE from 3.622689 to 0.011142; fixed K=4 still has MAE 0.026607. Training/inference score equality is not asserted for the full model.
Speed comparisons
Full Ouro-1.4B warmed rollout medians, five alternating-order pairs, two prompts, 32 response tokens, temperature 0.9 and fixed seeds; initialization and capture excluded:
Thor component medians use matched updates 1–3 at K=3/4/2, four prompts × four samples and 16-token responses. A100 medians include first-use capture and fresh resume under existing GPU contention. Whole-training acceleration and reward convergence are not demonstrated.
The A100 runtime's default UCX GPU IPC probe failed; passing PD runs use command-scoped
UCX_TLS=tcp,cuda_copywith host-staged transport. BF16 bitwise comparisons require matched one-slot batches; differing dynamic batch shapes fail strict parity. TP/PP/CP=1 is required by this recipe.Launch, parameter conversion and resume instructions. Completed task-owned model/checkpoint/export files are cleaned after checks. No process termination, environment updates or lcpu NFS changes were performed.
Historical layer review: provider #3 → budget recipe #4 → public contract/conversion #5. This upstream PR is the consolidated diff against
0z5a-branch.