Skip to content

Add shared-parameter Ouro Megatron provider - #3

Closed
0z5a wants to merge 2 commits into
fix/vllm-030-rollout-endpointsfrom
feat/ouro-shared-provider
Closed

0z5a wants to merge 2 commits into
fix/vllm-030-rollout-endpointsfrom
feat/ouro-shared-provider

Conversation

@0z5a

@0z5a 0z5a commented Sep 28, 2026 •

Copy link
Copy Markdown
Owner

Incremental review: Files changed.

Direct dependency and ancestor review: #1. Baseline: cleaned #1 head 555a862a45bf243e4ff21258896875ca65133148. This fork PR shows only this layer.

Add a shared-parameter Ouro Megatron provider using the existing custom-provider hook. It preserves physical parameter layout, four sandwich norms, inter-loop normalization and shared-layer gradients; differentiable SDPA handles training. TP=PP=CP=1; unused gate weights are retained and frozen.

H20 correctness check Result
Tiny shared gradients / recomputation / cached scores at K=2/3/4 Logit error ≤1.2e-7
Packed-sequence isolation / physical parameter count Passed
Full BF16 logits vs matched HF implicit-causal SDPA Exact at K=2/3/4
Full K=4 FP32 logits vs default HF SDPA Error 0
Full K=4 FP32 cached scores vs dense Error 6.68e-6
Full BF16 cached scores vs dense Error 0.0478 / 0.1057 / 0.1046

An initial explicit-mask BF16 comparison failed; matched implicit-causal dispatch explains the logit result. BF16 execution is not generally bit-exact. Tiny two-rank training/save/restart passed. The recipe layer #4 completes six full-model updates and resumes a seventh. Reward convergence and long-context coverage remain outstanding; speed improvement is unmeasured.

This historical provider layer is consolidated with later recipe/conversion work in upstream #459. Test and validation files were removed from this diff.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
@0z5a

0z5a commented Oct 6, 2026

Copy link
Copy Markdown
Owner Author

Superseded by vllm-project#459, now merged into 0z5a-branch. This historical Ouro layer is included in the consolidated upstream integration; closing the duplicate PR.

@0z5a 0z5a closed this Oct 6, 2026
0z5a added a commit that referenced this pull request Nov 7, 2026
…project#483) (vllm-project#506)

* [Bugfix][Hardware][AMD] Fix colocate weight sync on ROCm and add a ROCm 7.14 Dockerfile (vllm-project#418)

* [BugFix][ROCm] Widen rollout engine's visible devices so trainer IPC handles resolve on AMD GPUs

A HIP IPC handle carries the EXPORTING process's device ordinal, and
hipIpcOpenMemHandle resolves that ordinal against the IMPORTING process's
device list. CUDA instead honours the importer's current device. Colocate
weight sync therefore breaks on ROCm whenever a rollout engine is launched
with a mask narrower than the trainer's: the trainer sees the whole node and
exports at ordinal N, while the engine enumerates a single device where
ordinal N is out of range. Weight sync succeeds on GPU 0 and fails on every
other GPU with hipErrorInvalidValue.

_compute_server_args built the engine subprocess's _visible_devices from the
engine's own GPUs alone. Route it through _engine_visible_devices(), which
lists the engine's own GPUs first and the rest of the node's afterwards. The
engine still lands on its intended devices because its own are first; only
the mask's LENGTH matters for ordinal resolution, not any agreement between
the two numberings.

Gated on torch.version.hip: on CUDA the narrow mask is correct and widening
it would be a needless behaviour change.

Verified on 8x MI308X (gfx942), ROCm 7.14, torch 2.12.0+rocm7.14, vLLM
v0.29.0, colocate vLLM rollout + Megatron trainer. Single-variable A/B with
4 actor GPUs and --rollout-num-gpus-per-engine 1, so the engines are launched
narrower than the trainer: without this commit three of four engines fail the
first weight sync with a 500 from /update_weights and CUDA error: invalid
argument in rebuild_cuda_tensor -> UntypedStorage._new_shared_cuda; with it
the run exits 0 and all four engines sync. A framework-free two-process probe
isolates the same ordinal resolution in ~10s.

Co-authored-by: Claude
Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Treemann <rongzhang.zheng@amd.com>

* [Docker][ROCm] Add a ROCm 7.14 Dockerfile for AMD GPUs

docker/Dockerfile.rocm targets ROCm 7.0.2 / gfx950 and builds the whole stack
onto ubuntu:22.04 across nine stages. On a host that already has a prebuilt
ROCm 7.14 image most of that work is done, so this variant starts from
rocm/primus:v26.4 - which ships torch, triton, flash-attention,
TransformerEngine and aiter - and adds only vLLM, Megatron-LM and vime.

Single stage, no build context: every source is a pinned upstream git
coordinate with its commit asserted, and nothing is COPYed in. VIME_REF
accepts a branch or a full 40-hex sha; a sha is asserted, a branch has its
resolved sha printed.

Two properties are deliberate.

The base image's ROCm torch is never replaced. A pip constraints file is
generated from the installed versions right after the vLLM build, and every
later pip install runs under it. vime's requirements.txt asks for
vllm-router, transformers and ray[default] unpinned; without the constraints
a resolver can satisfy one of them with a PyPI torch or vllm wheel - those
are CUDA-only - and the breakage surfaces much later as a runtime failure on
the first GPU op. A pip install --dry-run guard aborts the build if the
resolver still plans to do it.

A verification layer fails the build rather than the first training run: ROCm
torch untouched, VLLM_TAG actually installed, vime/Megatron/Ray/
torch_memory_saver importing, and the vLLM API window vime needs present at
both ends. That last check is why the VLLM_TAG pin is a checked claim and not
a comment: vime needs NCCLTrainerInitInfo from
vllm.distributed.weight_transfer.nccl_engine, which v0.27.1 predates, and it
needs whichever cli_args location its current revision imports - the check
reads that from vime's own source so it stays meaningful as main advances.

Additive: docker/Dockerfile.rocm is untouched.

Co-authored-by: Claude
Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Treemann <rongzhang.zheng@amd.com>

* Remove the unnecessary max() operation

Signed-off-by: Rongzhang Zheng <rongzhang.zheng@amd.com>

* [Docker][ROCm] Fix --offload-train issues in rocm7.14 image

- Bump torch_memory_saver to 016938275e46e72e7b8d60e8d15046171b3e3c72 to align HIP VMM
  allocation sizes to allocation granularity
- Set ROC_ENABLE_LARGE_BAR=0 to avoid TMS VMM .item() SIGSEGV on MI300-class hosts
- Canonicalize duplicate libamd_comgr.so.3 under _rocm_sdk_* and fail the build if
  LD_PRELOADing the TMS hook still triggers duplicate LLVM registration on import torch

Signed-off-by: Treemann <rongzhang.zheng@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Signed-off-by: Treemann <rongzhang.zheng@amd.com>
Signed-off-by: Rongzhang Zheng <rongzhang.zheng@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* [Sync] Mirror Slime v0.4.0 training stack and release updates (vllm-project#441)

* sync: update from Slime through #2393

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix: complete vLLM sync translations

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* sync: complete Slime update through #2394

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* Sync Slime through score centering

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix: rebase vLLM patches onto current nightly

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix: remove stale async CI entry and obsolete metrics fallbacks

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* feat: transport sampling-mask logprobs for score centering

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix(vllm): enable scale-out endpoint for rollout TITO

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix(rollout): keep requests during sleep and enable external PD TITO

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix(ci): enable TITO on external OPD teacher

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* test: enable scale-out render endpoint in Geo3K e2e

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* feat: mirror Slime persistent rollout queue and distributed fully async

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* test: exercise straw transport in mirrored GPU suites

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* style: match repository formatter for straw sync

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* ci: install straw in CPU contract container

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* sync: mirror Slime #2427 straw checkpoint rollback

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* ci: install straw dependency for sync tests

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* test: pin vllm generation config for score centering

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix: bind straw coordinator in vllm rollout

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* docs: record full GPU gate procedure

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* sync: mirror Slime #2432 Straw queue and replay changes

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* ci: quote Straw requirement in Buildkite shell steps

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Mirror Slime PipelineRL and reference top-p fix

Port #2437 flush scheduling and #2434 reference log-prob replay fix. Replace request-bound TITO weight version with per-output propagation through vLLM #53199, and preserve exact top-p score-centering support.

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [CI] Add flush interval to weight updater test fixtures

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [CI] Accept versionless terminal abort chunk in PipelineRL probe

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Mirror concurrent rollout and HTTP regression tests

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Preserve buffered-token weight version on abort

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Avoid zero-KL assertion for top-p actor/ref mismatch

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Backport AuxOutput reset fix from vLLM #59060

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Stamp synthetic abort outputs with committed version

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Preserve coalesced rollout metadata and validate sampler probabilities

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Remove superseded per-file spill and mirror Straw references

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [CI] Inherit the available file descriptor limit for GPU jobs

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Backport terminal token-stream finish reasons

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Test] Cover terminal output with real prompt token metadata

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Mirror latest Slime and validate nightly Python/Rust adapters

Mirror Slime through 2f231865 with approved native overlays and Source test registration. Pin validated43b nightly, install tested tensor-dump/tracing Python and Rust patches, and parse native integer-ID top-k responses. Preserve original Straw assertions and saving policy. Keep profiler deferred and latest unchanged.

AI assistance: OpenAI Codex. Scoped validation171CPU tests, full precommit and11patch/98path installation proof pass; full remote CI and remaining candidate GPU validation pending.

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Mirror Slime external PD delta artifact assertion

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix: repair verified recovery and initial-policy CI contracts

Preserve training loss, numerical thresholds and checkpoint strategy. Add indexed precision dumps, offline-first manager setup, scoped reference load, same-distribution CI KL diagnostic, canonical optimizer masters and NIXL peer retirement.

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix: keep initial-model diagnostic optional for direct loss callers

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix: give the PD rollout debug dump a unique per-rollout path

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Mirror Slime v0.4.0 training stack and release updates

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Mirror latest Slime documentation and native engine contracts

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Restore upstream documentation release checks

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix(sync): mirror active DeepEP recipes and docs CI contracts

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* docs(sync): mirror Slime environment entry and preserve sampler parameters

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Preserve upstream export explanation and tokenization attribution

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Inherit Slime documentation deletions and navigation

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* Remove opt-in DeepEP auto switching

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [RL] Release native weight-transfer communicators and expose lifecycle APIs

Include fork vLLM vllm-project#119 on the fixed 43b4 baseline, including explicit communicator destruction and matching Python/Rust public controls. Keep the other sixteen-patch capabilities intact.

Assisted-by: OpenAI Codex
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Mirror Slime agent and collective updates through 58b08786

Source PRs: THUDM/slime#2453, #2450, #2458, #2457. Preserve approved Vime overlays and restore the omitted Source Straw error assertions. Native TITO, replay/SC, MTP and Buildkite adapters use the fixed 43b4 engine stack.

Validation: Slurm2366 301CPU/1skip,2376 Python3.11 agent step163pass/2optional skips,2377 Straw7+88pass,2378 nine precommit hooks. Native sandbox SSD prerequisites2373 pass; full new8GPU/model and final-image gates remain pending. No latest or release promotion.

Slime-commit: 58b0878679de339ec7844b45d884a87955a134be

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Build] Package native ARM Sunabako without changing x86 installation

Preserve Slime SDK0.1.1 contracts and the original x86 build target. Build a verified native ARM wheel with matching seccomp/syscall/ABI and driver paths before requirements resolution.

Source-Slime: 58b0878679de339ec7844b45d884a87955a134be
Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* Mirror rollout weight-group lifecycle and preserve draft state across sleep

Restore Source connect/disconnect layout and Agent MTP translation. Reuse vLLM parameter retention for draft-private weights instead of adding a second backup switch.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* sync: align draft retention patch with reviewed native selector

Mirror the simplified native proposal, including an empty fully-shared draft. No trainer, saving, or inference-path policy change.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* [Sync] Mirror router health-check ownership with native gateway patch

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix(sync): mirror source documentation and MTP verification budgets

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix(rust): mirror multimodal render-to-generate support

Keep the pinned vLLM base and existing Rust overlays; apply the independent multimodal bridge after them. Preserve explicit UUIDs, native validation and metadata-only EC support. Remove the GB10 launcher resurrected after the approved platform deletion.

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix(epd): forward producer transfer metadata and preserve media identity

Use the existing rollout endpoint metadata instead of inert SGLang-only engine keys. Preserve ordinary/streaming placement, complete media parts and caller UUIDs. Do not alias language_only to the unrelated modality-limit flag.

Validation: frozen700-path CPU2502 passes92 cases; real two-host ECCPU/NIXL2503 passes both frontends with four actual transfers and baseline-exact tokens/UUID metadata-only. Explicit connector/body assembly, not default deployment or all-model certification.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix(epd): mirror encoder deployment with native EC transfer

Keep Slime deployment sequencing and consumer rendering. Use native EC transport and DP endpoint ownership, preserve decoder input processing while omitting encoder weights, and supply resolved engine config and embedding support to Rust. Reuse existing upstream authentication fix; retain test assertions.

Validation: complete 23-patch stack on43b;83CLI/468server and all510Rust paths; installed experimentalfefcf78runtime/705sourcepaths/92CPU; no-overlay original3GPU E-P-D2552 passes Python/Rust. One-hour2554 is running, not claimed complete. Native review worktrees remain uncommitted. AI assistance used.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

---------

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>

* [CI][ROCm] Add ROCm CI: Dockerfile, Buildkite pipeline, and test variants (vllm-project#474)

* ci(rocm): ROCm CI test suite + review fixes, rebased onto main

Rebased onto upstream main (which now carries the Slime v0.4.0 sync vllm-project#441:
rebased vLLM patches, enable_scale_out, test_accelerator changes). Drops the
now-redundant docker/patch/vllm_engine content this branch previously copied
from vllm-project#432 — main and vllm-project#432 own those. Keeps only the ROCm CI additions and the
PR #3 review fixes:

- .buildkite/pipeline-rocm.yaml: RAY_EXPERIMENTAL_NOSET_* + NCCL env, the
  HIP->CUDA visibility wrapper (normalizes CUDA_VISIBLE_DEVICES to the locked
  HIP set in-command, keeping tests/ci/gpu_lock_exec.py untouched), and the
  vllm-config/quick-start test steps.
- tests/*-rocm.py: 8 standalone ROCm test variants (self-contained U.is_rocm
  shim; experimental/unwired PPO-critic note; corrected mixed-offload
  level-1-sleep restore description; dead code removed).
- docs amd_tutorial: `sleep infinity` so detached container stays up for exec.

Co-Authored-By: Claude <noreply@anthropic.com>

* ci(rocm): keep green vllm-config suite; bundle validated ROCm image

Retest on the rebuilt image (vllm/vime-rocm:pr-441) cleared startup and the
health gate for all suites; the 4 vllm-config GRPO tests pass to a train step.
Drop the 4 that fail for reasons outside this CI work:
  - quick_start_glm4_9B: GLM-specific step-0 KL exceeds the strict <1e-8 assert
    (ROCm rollout/train logprob divergence; the 0.5B GRPO tests pass the same check)
  - qwen3_4B_ppo{,_train_critic_only,_disaggregate}: SIGSEGV in critic model
    creation (known experimental ROCm PPO-critic path)

Bundle the validated Dockerfile.rocm (layers on vllm/vllm-openai-rocm:nightly,
credit @indianspeedster vllm-project#432) with the fixes the retest surfaced:
  - pin pyarrow==25.0.1: the base's 26.0.0 is a numpy-2 build that rejects the
    numpy 1.26.4 pin at import, breaking Ray serialization + HF datasets
  - apply vllm-generation-health.patch: provides /health_generate, which the
    post-vllm-project#441 rollout health monitor polls (slim 3-patch set had dropped it)
  - drop unused megatron-energon / multi-storage-client
  - CMD sleep infinity so `docker run -d` stays up for exec

Co-Authored-By: Claude <noreply@anthropic.com>

* docker(rocm): isolate ROCm-rebased vLLM patches from the shared CUDA set

d21bad2 overwrote docker/patch/latest/{vllm,vllm-pull_weights,
vllm-inflight-queue-diagnostics}.patch with ROCm-rebased versions, but the CUDA
docker/Dockerfile applies those same files with strict `git apply` — the rebase
would break the CUDA image build.

The ROCm nightly base genuinely needs the rebased patches: main's CUDA versions
apply without rejects but silently misplace the weight_version hunk under --fuzz
(lands on the wrong protocol class), failing the import-check assert.

Resolve by isolation:
  - docker/patch/latest/ reverted to main's versions (CUDA untouched)
  - ROCm-rebased patches moved to docker/amd_patch/latest/ (ROCm-only dir, already
    home to the ROCm megatron patches); Dockerfile.rocm points there
  - vllm-generation-health.patch stays shared (applies clean on both bases)

Rebuilt vllm/vime-rocm:pr-441 from docker/amd_patch/ — import check passes
(weight_version correctly placed). Also trims Dockerfile/pipeline comment blocks
to <=2 lines.

Co-Authored-By: Claude <noreply@anthropic.com>

---------

Co-authored-by: Claude <noreply@anthropic.com>

* [Bugfix][Rollout] Drop the leftover drain check from abort()

The sync merge took main's server_control.py, whose abort_servers_until_idle()
already waits until every engine reports idle, but kept 0z5a-branch's
report-only drain check at the end of abort(). That block reads
server_abort, verify_server_drain and _ABORT_DRAIN_VERIFY_DEADLINE_S, none
of which exist after the merge, so the first abort() of every rollout raised
NameError: name 'server_abort' is not defined.

abort() is now the same as on main. tests/test_server_idle_contract.py only
covered the removed probe API (verify_server_drain, query_engine_inflight,
abort_inflight_requests) and failed at import; main's tests for
abort_servers_until_idle cover the replacement.
test_abort_deletes_inflight_without_pause_resume patched the removed
verify_server_drain; without that patch it is the same as on main.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

* [Bugfix][Training] Stop reading the removed keep_old_actor argument

main removed --keep-old-actor in the Slime v0.4.0 sync (vllm-project#441), and the sync
merge dropped the argument, but two checks added on 0z5a-branch by the
external-draft collector (vllm-project#454) still read it:

- actor.py, the condition for reusing log-probs in the loss, raised
  AttributeError: 'Namespace' object has no attribute 'keep_old_actor' on
  the first train step whenever the terms before it hold, as in a plain
  GRPO run (one micro-batch group, policy loss, kl_coef 0);
- arguments.py, the collect-only validation, raised the same error for every
  collect-only configuration that passed the earlier checks.

Both terms are dropped. The option no longer exists, so they were always
false; the log-prob reuse condition is now the same as on main. The
collect-only test case that set keep_old_actor=True is removed with it, and
the draft-feature-collection docs (en, zh) no longer ask users to disable the
flag.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

* [Bugfix][Training] Restore the draft feature collector in the actor log-prob forward

The sync merge took main's compute_log_prob(), which has no
draft_feature_collector parameter, and dropped the lines in train_actor()
that create the collector, but kept the call that passes it. Every run that
needs a separate actor log-prob forward (can_reuse_log_probs_in_loss is
false, e.g. kl_coef != 0, more than one micro-batch group, a critic or GSPO)
raised NameError: name 'draft_feature_collector' is not defined, and
--draft-feature-mode collect-only could not collect anything.

Both pieces are restored as they were on 0z5a-branch (vllm-project#454), on top of
main's code: train_actor() creates the collector in collect-only mode, and
compute_log_prob() passes it to forward_only(), which still accepts it.
This also uses the unwrap_model import again.

tests/test_draft_feature_collector.py gains two CPU checks on the real
methods (isolated with ast, as tests/test_reference_top_p_replay.py does):
compute_log_prob() hands the collector to forward_only(), and train_actor()
assigns every collector it passes. Both fail on the merge commit. The fake
forward_only() in test_reference_top_p_replay.py accepts the new keyword.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

* [Bugfix][Rollout] Restore the inference log-prob validation from vllm-project#454

The sync merge kept the first half of 0z5a-branch's
_inference_generate_tokens_and_logprobs() (the docstring and the token_ids
check) but took main's second half, which zips token_ids with whatever
logprobs.content holds. A response with fewer log-prob entries than tokens
was silently truncated, and a missing, non-numeric or non-finite log-prob
(NaN, inf, True, 10**1000) went straight into training. The try/except in
generate() that adds the request id to these errors was dropped too, which
left `import math` unused (ruff F401) and fails the pre-commit CI step.

The parser tail and the wrapper are restored as they were in vllm-project#454, on top of
main's generate(). vllm-project#454 also changed main's
test_inference_generate_tokens_and_logprobs_aligns_partial_content to expect
the length-mismatch error; the merge took main's version of that test, so it
is restored to 0z5a-branch's.

Fixes 11 tests in tests/test_vllm_rollout.py that fail on the merge commit:
test_inference_logprob_rejects_invalid_numbers (6) and
test_bad_metadata_does_not_mutate_training_fields (5).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

* [CI] Stop running the removed server idle contract test

fdc372c removed tests/test_server_idle_contract.py together with the
leftover drain check it tested, but the synchronized upstream CPU step in
.buildkite/pipeline.yml still ran it. pytest exits non-zero on a missing
file, so that step would fail as soon as pre-commit passes. main's pipeline
does not list the file either.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

* [Test] Update two merged tests to the merged rollout and argument APIs

Both tests come from one side of the sync merge and now run against the
other side's code. Each fails on the merge commit.

- test_colocate_offload_resolution (from vllm-project#454, 18 cases) checked the
  namespace it passed to vime_validate_args(). Since vllm-project#441 on main that
  function validates a deep copy and returns it, so the test now reads the
  returned args, as the other tests in the file do.
- test_generate_requests_sampler_topk and
  test_generate_requests_complete_top_p_probabilities (from main) expected
  generate() to post without headers. vllm-project#454 tags every generate request with
  an x-request-id header (the id that the metadata errors restored in the
  previous commit report), so they now expect exactly that header.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

* [Doc] List the draft-feature and Ouro pages in main's docs navigation

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>

* [Bugfix][Training] Stop reading the removed update_weights_interval argument

main removed --update-weights-interval in vllm-project#441 and now publishes weights
after every actor update (train.py calls update_weights() each rollout).
The collect-only validation added by vllm-project#454 still read
args.update_weights_interval, so --draft-feature-mode collect-only failed
in vime_validate_args() with AttributeError: 'Namespace' object has no
attribute 'update_weights_interval' before anything ran.

The check is dropped: publication after every actor update is no longer
optional. The CPU tests did not see this because collect_args() set the
field by hand; it no longer does, so
test_collect_only_supported_arguments now covers the real namespace (it
fails on the previous head). The test case for interval 2 and the flag's
mention in the draft-feature-collection docs (en, zh) go with it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

* [Bugfix][Training] Give collect-only training data the identity the collector reads

main moved the conversion of rollout samples into training data from
vime/ray/rollout.py to vime/data/batch_builder.py. vllm-project#454 had added two things
to the old code for --draft-feature-mode collect-only, and the sync merge
kept neither: the group_indices and weight_versions columns, and passing
them through the per-rank DP split. DraftFeatureCollector._sequences() reads
both, so the first actor log-prob forward of a collect-only run failed with
KeyError: 'group_indices' (seen on a GH200 after rollout 0).

Both are restored in BatchBuilder as vllm-project#454 had them: convert() adds the two
columns in collect-only mode only, and split_by_dp() keeps them per sample.
The object-store split only copies keys that are present, so other modes are
unchanged.

tests/test_draft_feature_metadata.py checks that a collect-only conversion
carries both columns and that the object-store split hands each rank the
values of its own samples (fails on the previous head), and that conversion
without collection adds neither.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

---------

Signed-off-by: Treemann <rongzhang.zheng@amd.com>
Signed-off-by: Rongzhang Zheng <rongzhang.zheng@amd.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Co-authored-by: Rongzhang Zheng <rongzhang.zheng@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Joy Song <146360951+pancake0003@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>
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