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

docs/en/examples/qwen3-4B.md

@@ -292,8 +292,8 @@ In this case, 2 GPUs will be allocated for training, and 6 GPUs will be allocate
    --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256)
    ```
 
-### Asynchronous Training
+### Fully-Asynchronous Rollout
 
-When you separate training and inference, you may notice that the training and inference GPUs are always waiting for each other. To prevent these resources from being idle, we can enable asynchronous training. This can be done by changing `train.py` to `train_async.py` in the startup script. By doing this, vime will generate data for the next rollout while training on the current one.
-
-The only difference between `train.py` and `train_async.py` lies in the synchronization logic of the training loop. We achieve this by using Ray's asynchronous features (`.remote`, `ray.get`).
+For long-tail rollout workloads, use `train.py` together with
+`--rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async`.
+The background rollout worker keeps generation warm while the normal training loop consumes completed groups.

docs/en/examples/qwen3-4b-base-openhermes.md

@@ -81,5 +81,3 @@ You can compare [run-qwen3-4B-base-sft.sh](https://github.com/vllm-project/vime/
     As for `--calculate-per-token-loss`, this is because `vime` defaults to calculating the per-sample mean for GRPO. In general SFT training, the average is taken over all unmasked tokens in a batch, so it is recommended to configure this.
 
     Finally, `--disable-compute-advantages-and-returns` indicates that there is no need to pre-calculate log probabilities during the SFT process, and `--debug-train-only` means that `vllm` does not need to be initialized.
-
-3.  Used `train_async.py` instead of `train.py`. This is to leverage the asynchronous training process to implement data prefetching.

docs/zh/examples/qwen3-4B.md

@@ -292,10 +292,10 @@ ray job submit ... \
    --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256)
    ```
 
-### 异步训练
+### 全异步 Rollout
 
-当进行训推分离时，你会发现训练和推理的 GPU 总是相互等待着，为了避免这种资源空闲，我们可以开启异步训练。开启的方式即为将启动脚本中的 `train.py` 改变为 `train_async.py`。这样 vime 就会在进行当前 rollout 的训练时进行下一个 rollout 的数据生成了。
+对于长尾 rollout 工作负载，使用 `train.py`，并设置
+`--rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async`。
+后台 rollout worker 会持续生成数据，普通训练循环则消费已经完成的 group。
 
 ⚠️  异步训练时，vLLM 的性能统计日志与训练日志可能混在一起。可通过 `--vllm-disable-log-stats` 关闭性能统计，并用 `--vllm-uvicorn-log-level warning` 降低服务日志量。
-
-`train.py` 和 `train_async.py` 的差别只在于 train loop 的同步逻辑，我们通过 ray 的异步（`.remote`, `ray.get`）实现了这点。

docs/zh/examples/qwen3-4b-base-openhermes.md

@@ -81,5 +81,3 @@ bash scripts/run-qwen3-4B-base-sft.sh
    至于 `--calculate-per-token-loss`，这是因为 vime 默认是以 GRPO 的 per sample mean 进行计算的，而一般 sft 训练都是按一个 batch 的所有不被 mask 的 token 取平均，所以建议配置上。
 
    最后 `--disable-compute-advantages-and-returns` 表示 sft 的过程中不需要预先计算 log prob，`--debug-train-only` 表示不需要初始化 vllm。
-
-3. 使用了 `train_async.py` 而不是 `train.py`。这是为了利用异步训练的流程，来实现数据 prefetch。

examples/fully_async/README.md

@@ -41,13 +41,12 @@ fully-async rollout 0: done in ...s, queue_left=...
 
 ## How To Plug Your Own Generate Into This
 
-Two pieces flip the standard pipeline into fully-async:
+Fully-async uses the standard `python3 train.py` entrypoint. Select the
+rollout implementation with:
 
-1. Use the async training driver: `python3 train_async.py` (not `train.py`).
-2. Set the rollout function path:
-   ```
-   --rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async
-   ```
+```
+--rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async
+```
 
 For custom per-sample logic, use vime's standard plug-in points — they
 work unchanged under fully-async:

examples/fully_async/run-qwen2.5-0.5B-fully_async.sh

@@ -121,7 +121,7 @@ ROLLOUT_GPUS=${ROLLOUT_GPUS:-$((NUM_GPUS - ACTOR_GPUS))}
 
 ray job submit --address="http://127.0.0.1:8265" \
    --runtime-env-json="${RUNTIME_ENV_JSON}" \
-   -- python3 train_async.py \
+   -- python3 train.py \
    --actor-num-nodes 1 \
    --actor-num-gpus-per-node "${ACTOR_GPUS}" \
    --rollout-num-gpus "${ROLLOUT_GPUS}" \

examples/fully_async/run-qwen3.5-9B-fully_async.sh

@@ -125,7 +125,7 @@ ROLLOUT_GPUS=${ROLLOUT_GPUS:-$((NUM_GPUS - ACTOR_GPUS))}
 
 ray job submit --address="http://127.0.0.1:8265" \
    --runtime-env-json="${RUNTIME_ENV_JSON}" \
-   -- python3 train_async.py \
+   -- python3 train.py \
    --actor-num-nodes 1 \
    --actor-num-gpus-per-node "${ACTOR_GPUS}" \
    --rollout-num-gpus "${ROLLOUT_GPUS}" \

scripts/run-qwen3-235B-A22B-sft.sh

@@ -138,7 +138,7 @@ RUNTIME_ENV_JSON="{
 
 ray job submit --address="http://127.0.0.1:8265" \
    --runtime-env-json="${RUNTIME_ENV_JSON}" \
-   -- python3 train_async.py \
+   -- python3 train.py \
    --actor-num-nodes 4 \
    --actor-num-gpus-per-node 8 \
    ${MODEL_ARGS[@]} \

scripts/run-qwen3-4B-base-sft.sh

@@ -114,7 +114,7 @@ RUNTIME_ENV_JSON="{
 
 ray job submit --address="http://127.0.0.1:8265" \
    --runtime-env-json="${RUNTIME_ENV_JSON}" \
-   -- python3 train_async.py \
+   -- python3 train.py \
    --actor-num-nodes 1 \
    --actor-num-gpus-per-node 8 \
    ${MODEL_ARGS[@]} \

scripts/run-qwen3-8B-async-amd.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Async (non-colocated) GRPO for Qwen3-8B on ROCm (gfx950 / MI350): train_async.py,
+# Fully-async non-colocated GRPO for Qwen3-8B on ROCm (gfx950 / MI350),
 # actor and rollout on disjoint GPUs, RCCL-broadcast weight sync (no colocate IPC).
 
 # for rerun the task
@@ -48,6 +48,7 @@ CKPT_ARGS=(
 )
 
 ROLLOUT_ARGS=(
+   --rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async
    --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl
    --input-key prompt
    --label-key label
@@ -131,7 +132,7 @@ RUNTIME_ENV_JSON="{
 
 ray job submit --address="http://127.0.0.1:8265" \
    --runtime-env-json="${RUNTIME_ENV_JSON}" \
-   -- python3 train_async.py \
+   -- python3 train.py \
    --actor-num-nodes 1 \
    --actor-num-gpus-per-node ${ACTOR_GPUS} \
    --rollout-num-gpus ${ROLLOUT_GPUS} \

scripts/run-qwen3.5-35B-A3B-sft.sh

@@ -150,7 +150,7 @@ RUNTIME_ENV_JSON="{
 
 ray job submit --address="http://127.0.0.1:8265" \
    --runtime-env-json="${RUNTIME_ENV_JSON}" \
-   -- python3 train_async.py \
+   -- python3 train.py \
    --actor-num-nodes 1 \
    --actor-num-gpus-per-node 8 \
    ${MODEL_ARGS[@]} \

tests/test_accelerator.py

@@ -55,18 +55,31 @@ def reset_accelerator_selection(monkeypatch):
     selected = accelerator._ACCELERATOR
     patch_imported = accelerator._MUSA_PATCH_IMPORTED
     bootstrap_checked = accelerator._MUSA_BOOTSTRAP_CHECKED
-    for name in ("VIME_ACCELERATOR", "MUSA_VISIBLE_DEVICES", "MUSA_PATCH_PATH", "CUDA_VISIBLE_DEVICES"):
+    supa_imported = accelerator._SUPA_RUNTIME_IMPORTED
+    supa_bootstrap_checked = accelerator._SUPA_BOOTSTRAP_CHECKED
+    for name in (
+        "VIME_ACCELERATOR",
+        "MUSA_VISIBLE_DEVICES",
+        "MUSA_PATCH_PATH",
+        "CUDA_VISIBLE_DEVICES",
+        "SUPA_VISIBLE_DEVICES",
+        "BIREN_HOME",
+    ):
         monkeypatch.delenv(name, raising=False)
     accelerator._REGISTRY.clear()
     accelerator.reset_accelerator()
     accelerator._MUSA_PATCH_IMPORTED = False
     accelerator._MUSA_BOOTSTRAP_CHECKED = False
+    accelerator._SUPA_RUNTIME_IMPORTED = False
+    accelerator._SUPA_BOOTSTRAP_CHECKED = False
     yield
     accelerator._REGISTRY.clear()
     accelerator._REGISTRY.update(registry)
     accelerator._ACCELERATOR = selected
     accelerator._MUSA_PATCH_IMPORTED = patch_imported
     accelerator._MUSA_BOOTSTRAP_CHECKED = bootstrap_checked
+    accelerator._SUPA_RUNTIME_IMPORTED = supa_imported
+    accelerator._SUPA_BOOTSTRAP_CHECKED = supa_bootstrap_checked
 
 
 @pytest.mark.unit
@@ -108,6 +121,7 @@ def import_musa_patch():
 @pytest.mark.unit
 def test_cpu_only_initialization_does_not_require_an_accelerator(monkeypatch):
     monkeypatch.setattr(accelerator, "is_musa_available", lambda: False)
+    monkeypatch.setattr(accelerator, "is_supa_available", lambda: False)
     monkeypatch.setattr(accelerator, "_cuda_available", lambda: False)
 
     assert accelerator.initialize_accelerator() is None
@@ -146,6 +160,7 @@ class RegisteredAccelerator(FakeAccelerator):
         name = "registered"
 
     monkeypatch.setattr(accelerator, "is_musa_available", lambda: False)
+    monkeypatch.setattr(accelerator, "is_supa_available", lambda: False)
     monkeypatch.setattr(accelerator, "_cuda_available", lambda: False)
     accelerator.register_accelerator("registered", RegisteredAccelerator, lambda: True, priority=300)
 
@@ -199,5 +214,70 @@ def test_musa_availability_handles_missing_torch_namespace(monkeypatch):
     assert accelerator.is_musa_available() is False
 
 
+@pytest.mark.unit
+def test_cuda_selection_does_not_bootstrap_supa(monkeypatch):
+    monkeypatch.setenv("VIME_ACCELERATOR", "cuda")
+    monkeypatch.setenv("SUPA_VISIBLE_DEVICES", "0")
+    monkeypatch.setattr(accelerator, "_cuda_available", lambda: True)
+    monkeypatch.setattr(accelerator.CUDAAccelerator, "is_available", lambda self: True)
+    monkeypatch.setattr(
+        accelerator,
+        "_import_torch_supa",
+        lambda: pytest.fail("CUDA selection must not import torch_supa"),
+    )
+
+    assert accelerator.get_accelerator().name == "cuda"
+    assert accelerator.process_group_backend() == "nccl"
+    assert accelerator.visible_devices_env_key() == "CUDA_VISIBLE_DEVICES"
+
+
+@pytest.mark.unit
+def test_selected_supa_bootstraps_runtime_once(monkeypatch):
+    imports = []
+    fake_supa = SimpleNamespace(is_available=lambda: True)
+
+    def import_torch_supa():
+        imports.append("torch_supa")
+        monkeypatch.setattr(accelerator.torch, "supa", fake_supa, raising=False)
+        return True
+
+    monkeypatch.setenv("SUPA_VISIBLE_DEVICES", "0")
+    monkeypatch.setattr(accelerator, "_import_torch_supa", import_torch_supa)
+
+    assert imports == []
+    assert accelerator.initialize_accelerator().name == "supa"
+    assert accelerator.initialize_accelerator().name == "supa"
+    assert imports == ["torch_supa"]
+
+
+@pytest.mark.unit
+def test_supa_backend_maps_devices_and_process_groups(monkeypatch):
+    monkeypatch.setattr(accelerator.SUPAAccelerator, "is_available", lambda self: True)
+    monkeypatch.setenv("SUPA_VISIBLE_DEVICES", "2,5")
+    accelerator.set_accelerator(accelerator.SUPAAccelerator())
+
+    assert accelerator.device_type() == "supa"
+    assert accelerator.visible_devices_env_key() == "SUPA_VISIBLE_DEVICES"
+    assert accelerator.resolve_visible_device_id("5") == 1
+    assert accelerator.process_group_backend() == "bccl"
+    assert accelerator.weight_update_backend() == "cpu:gloo,supa:bccl"
+    assert accelerator.process_group_backend("gloo") == "gloo"
+    assert accelerator.distributed_device_id() is None
+
+
+@pytest.mark.unit
+def test_supa_backend_is_recognized_as_accelerator_backend():
+    assert accelerator.is_accelerator_backend("bccl") is True
+    assert accelerator.is_accelerator_backend("cpu:gloo,supa:bccl") is True
+    assert accelerator.is_accelerator_backend("gloo") is False
+
+
+@pytest.mark.unit
+def test_supa_availability_handles_missing_torch_namespace(monkeypatch):
+    monkeypatch.delattr(accelerator.torch, "supa", raising=False)
+
+    assert accelerator.is_supa_available() is False
+
+
 if __name__ == "__main__":
     raise SystemExit(pytest.main([__file__]))

tests/test_cuda_stack_offload.py

@@ -0,0 +1,47 @@
+import ctypes
+from types import SimpleNamespace
+
+import pytest
+
+from vime.utils import memory_utils
+
+NUM_GPUS = 0
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("cuda,hip,initialized", [("12.0", "6.0", True), ("12.0", None, False), (None, None, True)])
+def test_stack_reset_does_not_initialize_unsupported_cuda(monkeypatch, cuda, hip, initialized):
+    monkeypatch.setattr(memory_utils.torch.version, "cuda", cuda)
+    monkeypatch.setattr(memory_utils.torch.version, "hip", hip)
+    monkeypatch.setattr(memory_utils.torch.cuda, "is_initialized", lambda: initialized)
+    monkeypatch.setattr(memory_utils, "_cuda_stack_api", lambda: pytest.fail("must not load CUDA"))
+
+    memory_utils.reset_cuda_stack_size()
+
+
+@pytest.mark.unit
+def test_stack_reset_restores_cuda_default(monkeypatch):
+    calls = []
+
+    class Driver:
+        def cuCtxGetLimit(self, output, limit):
+            ctypes.cast(output, ctypes.POINTER(ctypes.c_size_t)).contents.value = 32768
+            calls.append(("get", limit))
+            return 0
+
+        def cuCtxSetLimit(self, limit, value):
+            calls.append(("set", limit, value))
+            return 0
+
+    monkeypatch.setattr(memory_utils.torch, "version", SimpleNamespace(cuda="12.0", hip=None))
+    monkeypatch.setattr(memory_utils.torch.cuda, "is_initialized", lambda: True)
+    monkeypatch.setattr(memory_utils.torch.cuda, "synchronize", lambda: calls.append(("sync",)))
+    monkeypatch.setattr(memory_utils, "_cuda_stack_api", lambda: Driver())
+
+    memory_utils.reset_cuda_stack_size()
+
+    assert calls == [("sync",), ("get", 0), ("set", 0, 1024)]
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__]))

tests/test_fully_async_rollout.py

@@ -43,6 +43,8 @@
 import pytest
 
 import vime.rollout.fully_async_rollout as fa
+from vime.rollout.filter_hub.base_types import DynamicFilterOutput
+from vime.utils.staleness import compute_staleness_metrics, sample_staleness
 from vime.utils.types import Sample
 
 
@@ -79,7 +81,7 @@ def _make_group(index: int) -> list[Sample]:
 
 def _make_worker(monkeypatch, data_buffer=None, concurrency=4) -> fa.AsyncRolloutWorker:
     monkeypatch.setattr(fa, "GenerateState", _FakeGenerateState)
-    args = SimpleNamespace(rollout_global_dataset=True, rollout_batch_size=4)
+    args = SimpleNamespace(rollout_global_dataset=True, rollout_batch_size=4, n_samples_per_prompt=1)
     return fa.AsyncRolloutWorker(args, data_buffer or _FakeDataBuffer([]), concurrency=concurrency)
 
 
@@ -112,6 +114,66 @@ def test_get_completed_groups_limit(monkeypatch):
     assert worker.get_completed_groups(limit=3) == []
 
 
+@pytest.mark.unit
+def test_concurrency_is_scaled_by_samples_per_prompt(monkeypatch):
+    monkeypatch.setattr(fa, "GenerateState", _FakeGenerateState)
+    args = SimpleNamespace(n_samples_per_prompt=4)
+    assert fa.AsyncRolloutWorker(args, _FakeDataBuffer([]), concurrency=10).concurrency == 2
+    assert fa.AsyncRolloutWorker(args, _FakeDataBuffer([]), concurrency=2).concurrency == 1
+
+
+@pytest.mark.unit
+def test_dynamic_filter_drops_groups_and_refills(monkeypatch):
+    worker = _make_worker(monkeypatch)
+    for gid in range(8):
+        group = _make_group(gid)
+        group[0].reward = float(gid % 2)
+        worker.output_queue.put((gid, group))
+    monkeypatch.setattr(fa, "_get_global_worker", lambda args, data_buffer: worker)
+
+    def keep_odd(args, group):
+        return DynamicFilterOutput(keep=bool(group[0].reward), reason="even")
+
+    monkeypatch.setattr(fa, "load_function", lambda path: keep_odd)
+    args = SimpleNamespace(
+        rollout_global_dataset=True,
+        rollout_batch_size=3,
+        dynamic_sampling_filter_path="test.keep_odd",
+    )
+    result = asyncio.run(fa._generate_rollout_async(args, rollout_id=7, data_buffer=None))
+
+    assert [group[0].index for group in result.samples] == [1, 3, 5]
+    dropped = result.metrics["_dropped_samples"]
+    assert [group[0].index for group in dropped] == [0, 2, 4]
+    assert all(sample.remove_sample for group in dropped for sample in group)
+    assert result.metrics["rollout/dynamic_filter/drop_even"] == 3
+    assert [gid for gid, _ in worker.get_completed_groups()] == [6, 7]
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "versions,expected",
+    [(["10"], 1), (["8"], 3), (["3", "10"], 8), ([], None), (["bad"], None), (["11"], None), (["²"], None)],
+)
+def test_sample_staleness_uses_oldest_valid_version(versions, expected):
+    assert sample_staleness(SimpleNamespace(weight_versions=versions), 10) == expected
+
+
+@pytest.mark.unit
+def test_staleness_metrics_use_serving_snapshot():
+    samples = [
+        SimpleNamespace(weight_versions=["2"]),
+        SimpleNamespace(weight_versions=["3"]),
+        SimpleNamespace(weight_versions=[]),
+    ]
+    assert compute_staleness_metrics(samples, 10) == {
+        "staleness/unknown_count": 1,
+        "staleness/mean": 8.5,
+        "staleness/max": 9,
+    }
+    assert compute_staleness_metrics(samples, None) == {}
+
+
 @pytest.mark.unit
 def test_done_callback_never_blocks_event_loop_thread(monkeypatch):
     """The callback runs on the loop thread; blocking there freezes every

tests/test_megatron_argument_validation.py

@@ -251,7 +251,6 @@ def make_vime_validate_args(**overrides):
         rollout_max_prompt_len=None,
         train_backend="megatron",
         release_train=False,
-        keep_old_actor=False,
         only_train_params_name_list=None,
         freeze_params_name_list=None,
         update_weight_transport="nccl",

tests/test_qwen2.5_0.5B_fully_async_short.py

@@ -133,7 +133,6 @@ def execute():
         train_args=train_args,
         num_gpus_per_node=NUM_GPUS,
         megatron_model_type=MODEL_TYPE,
-        train_script="train_async.py",
     )