[Bugfix] Fix runtime errors left by the 0z5a-branch sync merge

.buildkite/pipeline.yml

@@ -238,7 +238,6 @@ steps:
             tests/test_rollout_sample_hooks.py \
             tests/test_vllm_rollout.py \
             tests/test_mem_agent_nvlink_detect.py \
-            tests/test_server_idle_contract.py \
             tests/test_agent/test_sandbox_exec_and_wait.py \
             tests/test_stateless_adam.py \
             tests/test_training_recovery.py \

docs/en/advanced/draft-feature-collection.md

@@ -2,7 +2,7 @@
 
 `--draft-feature-mode collect-only` exports selected inputs to the actor's final, bias-free LM head during an existing log-probability forward. It does not train or publish a draft model. The default `off` path creates no export directory or hook.
 
-The first version supports TP=PP=CP=DP=1, no VPP, micro-batch=1 and the actor model. Disable `--keep-old-actor` and `--use-rollout-logprobs`, and use `--update-weights-interval 1` so the version counter cannot lag actor updates. A batch that reuses the training forward instead of doing a separate actor log-probability forward is rejected before exporting; collection never inserts a compensating forward. For synchronous GRPO, use more than one optimizer step per rollout so the ordinary actor path already computes log probabilities separately.
+The first version supports TP=PP=CP=DP=1, no VPP, micro-batch=1 and the actor model. Disable `--use-rollout-logprobs` and use `--update-weights-interval 1` so the version counter cannot lag actor updates. A batch that reuses the training forward instead of doing a separate actor log-probability forward is rejected before exporting; collection never inserts a compensating forward. For synchronous GRPO, use more than one optimizer step per rollout so the ordinary actor path already computes log probabilities separately.
 
 Add these options to a working training recipe:
 

docs/en/examples/workflows/index.rst

@@ -9,6 +9,7 @@ Training & systems workflows
    ../../examples/qwen3-4b-base-openhermes.md
    ../../_examples_synced/fully_async/README.md
    ../../_examples_synced/dspark/README.md
+   ../../_examples_synced/ouro/README.md
    ../../_examples_synced/mem_agent/README.md
    ../../_examples_synced/delta_weight_sync/README.md
    ../../_examples_synced/eval_multi_task/README.md

docs/en/rollout/index.rst

@@ -14,3 +14,4 @@ Place engines, schedule generation, and move data and weights.
    ../advanced/delta-weight-sync.md
    ../advanced/vllm-config.md
    ../advanced/speculative-decoding.md
+   ../advanced/draft-feature-collection.md

docs/zh/advanced/draft-feature-collection.md

@@ -2,7 +2,7 @@
 
 `--draft-feature-mode collect-only` 在已有 actor log-probability forward 中导出最终无 bias LM head 的选中输入，不训练或发布 draft 模型。默认 `off` 路径不会创建导出目录或 hook。
 
-首版支持 TP=PP=CP=DP=1、无 VPP、micro-batch=1 和 actor 模型。关闭 `--keep-old-actor`、`--use-rollout-logprobs`，并设置 `--update-weights-interval 1`，避免 version counter 落后于 actor 更新。如果真实 batch 复用训练 forward，没有独立 actor log-probability forward，会在导出之前报错；采集不会补做 forward。同步 GRPO 可配置每轮多个 optimizer step，让正常 actor 路径本来就需要独立计算 log probabilities。
+首版支持 TP=PP=CP=DP=1、无 VPP、micro-batch=1 和 actor 模型。关闭 `--use-rollout-logprobs`，并设置 `--update-weights-interval 1`，避免 version counter 落后于 actor 更新。如果真实 batch 复用训练 forward，没有独立 actor log-probability forward，会在导出之前报错；采集不会补做 forward。同步 GRPO 可配置每轮多个 optimizer step，让正常 actor 路径本来就需要独立计算 log probabilities。
 
 在已跑通的训练 recipe 中追加：
 

docs/zh/examples/workflows/index.rst

@@ -9,6 +9,7 @@
    ../../examples/qwen3-4b-base-openhermes.md
    ../../_examples_synced/fully_async/README.md
    ../../_examples_synced/dspark/README.md
+   ../../_examples_synced/ouro/README.md
    ../../_examples_synced/mem_agent/README.md
    ../../_examples_synced/delta_weight_sync/README.md
    ../../_examples_synced/eval_multi_task/README.md

docs/zh/rollout/index.rst

@@ -14,3 +14,4 @@ Rollout 与调度
    ../advanced/delta-weight-sync.md
    ../advanced/vllm-config.md
    ../advanced/speculative-decoding.md
+   ../advanced/draft-feature-collection.md

tests/test_draft_feature_collector.py

@@ -1,5 +1,8 @@
+import ast
 import json
 from argparse import Namespace
+from contextlib import nullcontext
+from pathlib import Path
 
 import pytest
 import torch
@@ -201,3 +204,58 @@ def forward(self, hidden):
     collector.finish()
     features = load_file(str(collector.root / "batch-0000.safetensors"))["features"]
     torch.testing.assert_close(features, hidden[[2, 3], 0])
+
+
+def _actor_methods():
+    source = Path(__file__).resolve().parents[1] / "vime/backends/megatron_utils/actor.py"
+    actor = next(
+        node
+        for node in ast.parse(source.read_text()).body
+        if isinstance(node, ast.ClassDef) and node.name == "MegatronTrainRayActor"
+    )
+    return {node.name: node for node in actor.body if isinstance(node, ast.FunctionDef)}
+
+
+def test_compute_log_prob_passes_the_collector_to_forward_only():
+    seen = []
+
+    def fake_forward_only(*_args, draft_feature_collector=None, **_kwargs):
+        seen.append(draft_feature_collector)
+        return {}
+
+    # Importing the whole actor needs Ray and GPU Megatron; run the real method alone.
+    isolated = ast.Module(body=[_actor_methods()["compute_log_prob"]], type_ignores=[])
+    namespace = {
+        "DataIterator": object,
+        "torch": torch,
+        "timer": lambda _name: nullcontext(),
+        "forward_only": fake_forward_only,
+        "get_log_probs_and_entropy": None,
+    }
+    exec(compile(ast.fix_missing_locations(isolated), "actor.py", "exec"), namespace)
+    actor = Namespace(args=None, model=None)
+    collector = object()
+    namespace["compute_log_prob"](actor, [], [1], draft_feature_collector=collector)
+    namespace["compute_log_prob"](actor, [], [1], store_prefix="ref_")
+    assert seen == [collector, None]
+
+
+def test_train_actor_creates_the_collector_it_passes():
+    train_actor = _actor_methods()["train_actor"]
+    passed = [
+        (node.lineno, keyword.value.id)
+        for node in ast.walk(train_actor)
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "compute_log_prob"
+        for keyword in node.keywords
+        if keyword.arg == "draft_feature_collector" and isinstance(keyword.value, ast.Name)
+    ]
+    assigned = {
+        target.id: node.lineno
+        for node in ast.walk(train_actor)
+        if isinstance(node, ast.Assign)
+        for target in node.targets
+        if isinstance(target, ast.Name)
+    }
+    assert passed, "the actor log-prob forward no longer receives the draft feature collector"
+    for lineno, name in passed:
+        assert name in assigned and assigned[name] < lineno, f"{name} is passed before it is assigned"

tests/test_megatron_argument_validation.py

@@ -748,7 +748,7 @@ def test_colocate_offload_resolution(monkeypatch, colocate, offload_train, relea
         update_weight_transport="disk",
         update_weight_disk_dir="/tmp/weights",
     )
-    module.vime_validate_args(args)
+    args, _ = module.vime_validate_args(args)
     expected_train = False if colocate and release_train else (colocate if offload_train is None else offload_train)
     expected_rollout = True if colocate and release_train else (colocate if offload_rollout is None else False)
     assert args.offload_train is expected_train
@@ -859,7 +859,6 @@ def collect_args(**overrides):
         ({"draft_feature_max_tokens": 0}, "budgets must be positive"),
         ({"draft_feature_max_batches": -1}, "budgets must be positive"),
         ({"draft_feature_max_bytes": 0}, "budgets must be positive"),
-        ({"keep_old_actor": True}, "actor log-prob forward"),
         ({"use_rollout_logprobs": True}, "actor log-prob forward"),
         ({"compute_advantages_and_returns": False}, "actor log-prob forward"),
         ({"update_weights_interval": 2}, "publication after every actor update"),

tests/test_reference_top_p_replay.py

@@ -33,7 +33,15 @@ def test_reference_log_probs_use_full_vocab_while_actor_replays_top_p(monkeypatc
     monkeypatch.setattr(mpu, "get_tensor_model_parallel_rank", lambda: 0, raising=False)
 
     def fake_forward_only(
-        callback, args, model, data_iterator, _num_microbatches, *, store_prefix, use_rollout_top_p_replay
+        callback,
+        args,
+        model,
+        data_iterator,
+        _num_microbatches,
+        *,
+        store_prefix,
+        use_rollout_top_p_replay,
+        draft_feature_collector=None,
     ):
         kwargs = {}
         if use_rollout_top_p_replay:

tests/test_score_centering_transport.py

@@ -222,7 +222,8 @@ async def post(url, payload, **kwargs):
     sample = asyncio.run(rollout.generate(a, Sample(prompt="test"), {"max_new_tokens": 8}))
     payload, kwargs = captured[0]
     assert payload["sampling_params"]["logprobs"] == 4
-    assert kwargs == {"headers": None}
+    assert list(kwargs) == ["headers"]
+    assert list(kwargs["headers"]) == ["x-request-id"]
     assert sample.rollout_topk_token_ids.tolist() == [[3, 1, 4]]
 
 
@@ -471,7 +472,8 @@ def test_generate_requests_complete_top_p_probabilities(monkeypatch):
 
     async def post(url, payload, **kwargs):
         assert payload["sampling_params"]["logprobs"] == 1
-        assert kwargs == {"headers": None}
+        assert list(kwargs) == ["headers"]
+        assert list(kwargs["headers"]) == ["x-request-id"]
         return {
             "choices": [
                 {