[runtime, deps] fix: initialize rollout state with the current VeOmni API

README.md

@@ -76,6 +76,11 @@ Extras: `gpu` (PyTorch + FlashAttention + kernels), `verl` and `veomni`
 `[tool.uv.sources]` block; to develop against a local checkout, set
 `editable = true` there (inline hints included).
 
+VeOmni integration requires its current `init_parallel_state_from_config` API;
+the legacy `init_parallel_state` API is not supported. Standalone rollout
+workers derive their non-EP topology from the initialized process group,
+even when they inherit a different `WORLD_SIZE` from the trainer.
+
 ## Components
 
 - [`vexact/batch_invariant_ops/`](vexact/batch_invariant_ops/README.md) — batch-invariant operators/kernels for true on-policy RL training.

pyproject.toml

@@ -36,18 +36,18 @@ verl = [
     "verl",
     # VeRL's v1 trainer imports TransferQueue unconditionally, but the
     # upstream package does not currently declare it as a dependency.
-    "TransferQueue==0.1.9",
+    "TransferQueue==0.1.10",
     "uvicorn",
     "fastapi",
     "cachetools",
 ]
 veomni = [
     "veomni",
-    # VeOmni's default install pins transformers==5.9.0 (via its
+    # VeOmni's default install pins transformers==5.16.1 (via its
     # `transformers-stable` dependency group). Mirror that pin here so vexact
     # users picking up the veomni extra resolve to the same version VeOmni
     # tests/develops against.
-    "transformers==5.9.0",
+    "transformers==5.16.1",
 ]
 
 [build-system]
@@ -133,10 +133,10 @@ known-third-party = [
 ]
 
 [tool.uv.sources]
-verl = { git = "https://github.com/verl-project/verl.git", rev = "e9618406de5bad40041d7612554e465ec2003ec1" }
+verl = { git = "https://github.com/verl-project/verl.git", rev = "8718ca30a3f002f93b7c4fd99b9b2506718681bc" }
 # To work on verl locally, point this at `{ path = "./verl", editable = true }`.
 # To work on VeOmni locally, point this at `{ path = "./VeOmni", editable = true }`.
-veomni = { git = "https://github.com/ByteDance-Seed/VeOmni.git", rev = "9651e80750a7d204d69423d71063c2987b03bf08" }
+veomni = { git = "https://github.com/ByteDance-Seed/VeOmni.git", rev = "aa8bb43c6ba457afde4b9bd35cceaf74823604cc" }
 torch = [
     { index = "pytorch", extra = "gpu" },
 ]
@@ -162,5 +162,5 @@ environments = ["sys_platform == 'linux'"]
 override-dependencies = [
     "opencv-python-headless<4.13.0",
     # Keep the vllm and veomni extras on the same Transformers version.
-    "transformers==5.9.0",
+    "transformers==5.16.1",
 ]

tests/test_veomni_parallel_state.py

@@ -0,0 +1,112 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Check current VeOmni initialization in standalone spawned rollout workers."""
+
+import os
+import sys
+from datetime import timedelta
+from types import ModuleType
+from unittest.mock import Mock
+
+import pytest
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from vexact.utils.veomni_parallel_state import init_veomni_parallel_state
+
+
+def _install_veomni():
+    class AcceleratorConfig:
+        def __init__(self):
+            self.world_size = self.dp_size = self.dp_shard_size = int(os.environ.get("WORLD_SIZE", "1"))
+            self.dp_replicate_size = self.pp_size = self.tp_size = self.cp_size = self.ulysses_size = 1
+            self.ep_size = 1
+
+    veomni = ModuleType("veomni")
+    arguments = ModuleType("veomni.arguments")
+    arguments.AcceleratorConfig = AcceleratorConfig
+    distributed = ModuleType("veomni.distributed")
+    parallel_state = ModuleType("veomni.distributed.parallel_state")
+    distributed.parallel_state = parallel_state
+    veomni.arguments = arguments
+    veomni.distributed = distributed
+    return {
+        "veomni": veomni,
+        "veomni.arguments": arguments,
+        "veomni.distributed": distributed,
+        "veomni.distributed.parallel_state": parallel_state,
+    }
+
+
+@pytest.mark.parametrize("inherited_world_size", ["1", "16"])
+def test_uses_rollout_group_topology(monkeypatch, inherited_world_size):
+    modules = _install_veomni()
+    initializer = Mock()
+    modules["veomni.distributed.parallel_state"].init_parallel_state_from_config = initializer
+    for name, value in modules.items():
+        monkeypatch.setitem(sys.modules, name, value)
+    monkeypatch.setattr(dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(dist, "get_world_size", lambda: 4)
+    monkeypatch.setenv("WORLD_SIZE", inherited_world_size)
+
+    assert init_veomni_parallel_state() is initializer.return_value
+    initializer.assert_called_once()
+    accelerator = initializer.call_args.args[0]
+    assert accelerator.world_size == accelerator.dp_size == accelerator.dp_shard_size == 4
+    assert accelerator.dp_replicate_size == accelerator.pp_size == accelerator.tp_size == 1
+    assert accelerator.cp_size == accelerator.ulysses_size == accelerator.ep_size == 1
+    assert initializer.call_args.kwargs == {"name": None}
+    assert os.environ["WORLD_SIZE"] == inherited_world_size
+
+
+def test_uninitialized_group_does_not_import_veomni(monkeypatch):
+    monkeypatch.setattr(dist, "is_initialized", lambda: False)
+    monkeypatch.setitem(sys.modules, "veomni", None)
+    assert init_veomni_parallel_state() is None
+
+
+def _spawned_worker(rank, world_size, init_method):
+    assert "verl" not in sys.modules
+    assert os.environ["WORLD_SIZE"] == "16"
+    dist.init_process_group(
+        "gloo", init_method=init_method, rank=rank, world_size=world_size, timeout=timedelta(seconds=30)
+    )
+    try:
+        from unittest.mock import patch
+
+        from veomni.distributed import parallel_state
+
+        # Exercise a real VeOmni mesh on CPU even on the GPU CI runner.
+        with patch.object(parallel_state, "get_device_type", return_value="cpu"):
+            first = init_veomni_parallel_state()
+            assert first is init_veomni_parallel_state()
+        assert first is parallel_state.get_parallel_state()
+        assert first.dp_size == first.dp_shard_size == first.world_size == world_size
+        assert first.dp_replicate_size == first.pp_size == first.tp_size == 1
+        assert first.cp_size == first.ulysses_size == 1
+        assert not first.ep_enabled
+        assert first.device_mesh.device_type == "cpu"
+        assert first.dp_group.size() == world_size
+        assert "verl" not in sys.modules
+        assert os.environ["WORLD_SIZE"] == "16"
+        dist.barrier()
+    finally:
+        dist.destroy_process_group()
+
+
+def test_spawned_rollout_group_uses_current_api(monkeypatch, tmp_path):
+    pytest.importorskip("veomni")
+    monkeypatch.setenv("WORLD_SIZE", "16")
+    mp.spawn(_spawned_worker, args=(4, f"file://{tmp_path / 'rendezvous'}"), nprocs=4, join=True)

tests/verl_integrations/test_async_server_model_version.py

@@ -20,7 +20,8 @@
 
 pytest.importorskip("verl")
 
-from vexact.integrations.verl.async_server import VeXactServer  # noqa: E402
+from verl.checkpoint_engine.base import CheckpointEngineManager  # noqa: E402
+from vexact.integrations.verl.async_server import VeXactReplica, VeXactServer  # noqa: E402
 from vexact.integrations.verl.rollout import ServerAdapter  # noqa: E402
 
 
@@ -77,6 +78,31 @@ async def test_unversioned_weight_update_clears_previous_model_version():
     assert server.global_steps is None
 
 
+@pytest.mark.asyncio
+async def test_abort_accepts_current_verl_checkpoint_keyword():
+    server = VeXactServer.__new__(VeXactServer)
+    remote_abort = AsyncMock(side_effect=server.abort_all_requests)
+    replica = VeXactReplica.__new__(VeXactReplica)
+    replica.servers = [SimpleNamespace(abort_all_requests=SimpleNamespace(remote=remote_abort))]
+
+    manager = CheckpointEngineManager.__new__(CheckpointEngineManager)
+    manager.replicas = [replica]
+    await manager.abort_replicas(reject_request=False)
+
+    remote_abort.assert_awaited_once_with(reject_request=False)
+
+
+@pytest.mark.asyncio
+async def test_abort_does_not_silently_ignore_request_rejection():
+    server = VeXactServer.__new__(VeXactServer)
+    remote_abort = AsyncMock(side_effect=server.abort_all_requests)
+    replica = VeXactReplica.__new__(VeXactReplica)
+    replica.servers = [SimpleNamespace(abort_all_requests=SimpleNamespace(remote=remote_abort))]
+
+    with pytest.raises(NotImplementedError, match="rejecting incoming generation requests"):
+        await replica.abort_all_requests(reject_request=True)
+
+
 class _RemoteMethod:
     def __init__(self, name, calls):
         self.name = name