feat: add minimal strict selected-logprob provider hook

docs/en/advanced/selected-logprob-provider.md

@@ -22,3 +22,19 @@ loaded or raises `SelectedLogprobProviderUnavailable`. `strict` rejects that
 case. Provider exceptions and invalid result shapes always fail the run. A
 strict provider result must include non-empty `backend_id` and `contract_id`
 and remain connected to autograd when logits require gradients.
+
+For the RL-Kernel WS2 provider, the validated launch contract is:
+
+```bash
+--tensor-model-parallel-size 2 \
+--context-parallel-size 2 \
+--rollout-top-p 1.0 \
+--selected-logprob-provider rl_engine.integrations.vime.logp.provider \
+--selected-logprob-provider-mode strict
+```
+
+The provider owns the TP vocabulary reduction only. Vime continues to own CP
+token-row layout, response extraction, and PPO/GRPO loss composition. The
+provider does not claim attention or FFN train/rollout consistency; those
+claims require runtime readback from both Megatron and vLLM and are reported by
+the RL-Kernel validation example.

scripts/run-qwen3-8B-rlkernel-tp2-cp2.sh

@@ -0,0 +1,167 @@
+#!/usr/bin/env bash
+# Qwen3-8B GRPO smoke/validation run with the RL-Kernel selected-logprob
+# provider.  Vime remains the launcher; RL-Kernel owns the provider and its
+# contract.  This script intentionally keeps the framework-side change small.
+
+set -euo pipefail
+
+VIME_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
+RL_KERNEL_ROOT="${RL_KERNEL_ROOT:-${VIME_ROOT}/../RL-Kernel}"
+
+if [[ ! -f "${RL_KERNEL_ROOT}/rl_engine/integrations/vime/logp.py" ]]; then
+  echo "RL_KERNEL_ROOT must point to an RL-Kernel checkout containing the Vime provider" >&2
+  exit 2
+fi
+
+export PYTHONUNBUFFERED=1
+export VIME_RL_KERNEL_STRICT="${VIME_RL_KERNEL_STRICT:-1}"
+MEGATRON_ROOT="${MEGATRON_ROOT:-/root/Megatron-LM}"
+export PYTHONPATH="${RL_KERNEL_ROOT}:${VIME_ROOT}:${MEGATRON_ROOT}:${PYTHONPATH:-}"
+
+# The provider is a vocab-parallel TP implementation.  CP owns token rows and
+# must not be used as a vocabulary reduction group.
+TP_SIZE="${TP_SIZE:-2}"
+CP_SIZE="${CP_SIZE:-2}"
+ACTOR_GPUS="${ACTOR_GPUS:-4}"
+ROLLOUT_GPUS="${ROLLOUT_GPUS:-4}"
+NUM_GPUS="${NUM_GPUS:-8}"
+ROLLOUT_GPUS_PER_ENGINE="${ROLLOUT_GPUS_PER_ENGINE:-2}"
+ROLLOUT_TOP_P="${ROLLOUT_TOP_P:-1.0}"
+COLOCATE="${COLOCATE:-0}"
+
+if [[ "${NUM_GPUS}" != "8" || "${ACTOR_GPUS}" != "4" || "${ROLLOUT_GPUS}" != "4" ]]; then
+  echo "This validation entry point requires an 8-GPU node with 4 actor GPUs and 4 rollout GPUs" >&2
+  exit 2
+fi
+if [[ "${COLOCATE}" != "0" && "${COLOCATE}" != "1" ]]; then
+  echo "COLOCATE must be 0 (default, disjoint train/rollout GPUs) or 1" >&2
+  exit 2
+fi
+
+if [[ "${TP_SIZE}" != "2" || "${CP_SIZE}" != "2" ]]; then
+  echo "This validation entry point is intentionally fixed to TP=2, CP=2" >&2
+  exit 2
+fi
+if [[ "${ROLLOUT_TOP_P}" != "1.0" ]]; then
+  echo "RL-Kernel strict selected-logprob validation requires ROLLOUT_TOP_P=1.0" >&2
+  exit 2
+fi
+
+source "${VIME_ROOT}/scripts/models/qwen3-8B.sh"
+
+MODEL_ROOT="${MODEL_ROOT:-/root/Qwen3-8B}"
+TORCH_DIST_ROOT="${TORCH_DIST_ROOT:-/root/Qwen3-8B_torch_dist}"
+VIME_CKPT="${VIME_CKPT:-/root/Qwen3-8B_vime_rlkernel_tp2_cp2}"
+PROMPT_DATA="${PROMPT_DATA:-/root/dapo-math-17k/dapo-math-17k.jsonl}"
+
+if ! command -v nvidia-smi >/dev/null 2>&1; then
+  echo "nvidia-smi is required; refusing to run the CUDA validation on an unknown device" >&2
+  exit 3
+fi
+GPU_NAMES="$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || true)"
+GPU_COUNT="$(printf '%s\n' "${GPU_NAMES}" | sed '/^$/d' | wc -l | tr -d ' ')"
+if [[ "${GPU_COUNT}" != "${NUM_GPUS}" ]]; then
+  echo "Expected ${NUM_GPUS} visible GPUs, found ${GPU_COUNT}" >&2
+  printf '%s\n' "${GPU_NAMES}" >&2
+  exit 3
+fi
+if [[ "${GPU_REQUIRE_H100:-1}" == "1" ]] && ! printf '%s\n' "${GPU_NAMES}" | grep -q 'H100'; then
+  echo "Expected H100 GPUs; refusing to run on a different GPU class" >&2
+  printf '%s\n' "${GPU_NAMES}" >&2
+  exit 3
+fi
+python3 - <<'PY'
+import torch
+
+if not torch.cuda.is_available() or torch.cuda.device_count() != 8:
+    raise SystemExit("PyTorch must expose 8 CUDA devices for this validation")
+PY
+for required_path in "${MODEL_ROOT}" "${TORCH_DIST_ROOT}" "${PROMPT_DATA}" "${MEGATRON_ROOT}"; do
+  if [[ ! -e "${required_path}" ]]; then
+    echo "Required runtime path does not exist: ${required_path}" >&2
+    exit 3
+  fi
+done
+python3 - <<'PY'
+from rl_engine.integrations.vime.logp import provider
+print(f"RL-Kernel provider import OK: {provider.__module__}.{provider.__name__}")
+PY
+
+CKPT_ARGS=(
+  --hf-checkpoint "${MODEL_ROOT}"
+  --ref-load "${TORCH_DIST_ROOT}"
+  --load "${VIME_CKPT}"
+  --save "${VIME_CKPT}"
+  --save-interval 100000
+)
+
+ROLLOUT_ARGS=(
+  --prompt-data "${PROMPT_DATA}"
+  --input-key prompt
+  --label-key label
+  --apply-chat-template
+  --rollout-shuffle
+  --rm-type deepscaler
+  --num-rollout "${NUM_ROLLOUT:-1}"
+  --rollout-batch-size "${ROLLOUT_BATCH_SIZE:-8}"
+  --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT:-2}"
+  --rollout-max-response-len "${MAX_RESPONSE_LEN:-1024}"
+  --rollout-temperature 1.0
+  --rollout-top-p "${ROLLOUT_TOP_P}"
+  --global-batch-size "${GLOBAL_BATCH_SIZE:-16}"
+  --balance-data
+)
+
+PARALLEL_ARGS=(
+  --tensor-model-parallel-size "${TP_SIZE}"
+  --context-parallel-size "${CP_SIZE}"
+  --pipeline-model-parallel-size 1
+  --sequence-parallel
+  --expert-model-parallel-size 1
+  --expert-tensor-parallel-size 1
+  --use-dynamic-batch-size
+  --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU:-2048}"
+)
+
+RL_KERNEL_ARGS=(
+  --selected-logprob-provider rl_engine.integrations.vime.logp.provider
+  --selected-logprob-provider-mode strict
+  --custom-megatron-init-path rl_engine.integrations.megatron_runtime.initialize_from_environment
+)
+
+MISC_ARGS=(
+  --attention-dropout 0.0
+  --hidden-dropout 0.0
+  --attention-softmax-in-fp32
+  --attention-backend flash
+  --no-gradient-accumulation-fusion
+  --rollout-num-gpus-per-engine "${ROLLOUT_GPUS_PER_ENGINE}"
+  --vllm-gpu-memory-utilization "${VLLM_GPU_MEMORY_UTILIZATION:-0.4}"
+)
+
+ray stop --force || true
+ray start --head --node-ip-address "${MASTER_ADDR:-127.0.0.1}" \
+  --num-gpus "${NUM_GPUS}" --disable-usage-stats \
+  --dashboard-host=0.0.0.0 --dashboard-port="${RAY_DASHBOARD_PORT:-8265}"
+
+TRAIN_LAYOUT_ARGS=()
+if [[ "${COLOCATE}" == "1" ]]; then
+  TRAIN_LAYOUT_ARGS+=(--colocate)
+else
+  TRAIN_LAYOUT_ARGS+=(--megatron-to-hf-mode bridge)
+fi
+
+ray job submit --address="http://127.0.0.1:${RAY_DASHBOARD_PORT:-8265}" \
+  --working-dir "${VIME_ROOT}" \
+  -- python3 train.py \
+  --train-backend megatron \
+  --actor-num-nodes 1 \
+  --actor-num-gpus-per-node "${ACTOR_GPUS}" \
+  --rollout-num-gpus "${ROLLOUT_GPUS}" \
+  "${TRAIN_LAYOUT_ARGS[@]}" \
+  "${MODEL_ARGS[@]}" \
+  "${CKPT_ARGS[@]}" \
+  "${ROLLOUT_ARGS[@]}" \
+  "${PARALLEL_ARGS[@]}" \
+  "${RL_KERNEL_ARGS[@]}" \
+  "${MISC_ARGS[@]}"

tests/test_logprob_response_spans.py

@@ -5,7 +5,13 @@
 import torch
 
 from megatron.core import mpu
-from vime.backends.megatron_utils.loss import _build_topp_keep_mask, get_rollout_top_p_logprob_kwargs
+from vime.backends.megatron_utils.loss import (
+    _build_topp_keep_mask,
+    _maybe_capture_log_probs,
+    drain_captured_log_probs,
+    enable_log_prob_capture,
+    get_rollout_top_p_logprob_kwargs,
+)
 
 
 NUM_GPUS = 0
@@ -97,5 +103,20 @@ def test_top_p_mask_aligns_with_cp1_response_rows(monkeypatch):
     assert masked_rows == {2: [13], 3: [14], 5: [21], 6: [22], 7: [23]}
 
 
+@pytest.mark.unit
+def test_logprob_capture_uses_partition_keys_and_detaches_values():
+    enable_log_prob_capture()
+    first = torch.tensor([1.0, 2.0], requires_grad=True)
+    second = torch.tensor([3.0], requires_grad=True)
+
+    _maybe_capture_log_probs({"partition": [7, 3]}, [first, second])
+    captured = drain_captured_log_probs()
+
+    assert set(captured) == {3, 7}
+    torch.testing.assert_close(captured[7], first)
+    torch.testing.assert_close(captured[3], second)
+    assert not captured[7].requires_grad
+
+
 if __name__ == "__main__":
     raise SystemExit(pytest.main([__file__]))

tests/test_megatron_argument_validation.py

@@ -336,6 +336,8 @@ def add_cli_args(parser, **_kwargs):
     module.get_vime_extra_args_provider()(parser)
     args = parser.parse_args(
         [
+            "--rollout-batch-size",
+            "1",
             "--selected-logprob-provider",
             "rl_engine.integrations.vime.logp.provider",
             "--selected-logprob-provider-mode",
@@ -347,5 +349,17 @@ def add_cli_args(parser, **_kwargs):
     assert args.selected_logprob_provider_mode == "strict"
 
 
+@pytest.mark.unit
+def test_strict_selected_logprob_provider_rejects_top_p_replay(monkeypatch):
+    module = load_arguments_module(monkeypatch)
+    args = argparse.Namespace(
+        selected_logprob_provider="rl_engine.integrations.vime.logp.provider",
+        selected_logprob_provider_mode="strict",
+        rollout_top_p=0.9,
+    )
+    with pytest.raises(ValueError, match="rollout-top-p 1.0"):
+        module._validate_selected_logprob_provider_args(args)
+
+
 if __name__ == "__main__":
     raise SystemExit(pytest.main([__file__]))

tests/test_selected_logprob_provider.py

@@ -35,6 +35,27 @@ def _native(*args, **kwargs):
     return logits[:, :1], entropy
 
 
+def _structural_request(**overrides) -> SelectedLogprobRequest:
+    values = dict(
+        logits=torch.randn(3, 5),
+        target_ids=torch.tensor([1, 2, 3]),
+        tensor_parallel_group=None,
+        context_parallel=ContextParallelLayout(world_size=1, rank=0, layout="single"),
+        with_entropy=False,
+        with_entropy_grad=False,
+        chunk_size=64,
+        hidden=torch.randn(3, 4),
+        lm_head_weight=torch.randn(5, 4),
+        lm_head_bias=torch.randn(5),
+        vocab_start_index=0,
+        global_vocab_size=5,
+        real_vocab_size=4,
+        temperature=torch.ones(3),
+    )
+    values.update(overrides)
+    return SelectedLogprobRequest(**values)
+
+
 def _install_provider(monkeypatch, provider):
     module_name = "selected_logprob_provider_fixture"
     module = types.ModuleType(module_name)
@@ -81,6 +102,50 @@ def provider(actual_request):
     torch.testing.assert_close(entropy, request.logits.sum(dim=-1))
 
 
+def test_structural_request_accepts_aligned_hidden_and_lm_head():
+    request = _structural_request()
+
+    assert request.hidden is not None and request.hidden.shape == (3, 4)
+    assert request.lm_head_weight is not None and request.lm_head_weight.shape == (5, 4)
+
+
+@pytest.mark.parametrize(
+    ("overrides", "match"),
+    [
+        ({"lm_head_weight": None}, "lm_head_weight"),
+        ({"hidden": torch.randn(2, 4)}, "hidden rows"),
+        ({"lm_head_weight": torch.randn(5, 6)}, "hidden width"),
+        ({"temperature": torch.tensor([1.0, 0.0, 1.0])}, "temperature must be positive"),
+    ],
+)
+def test_structural_request_rejects_incomplete_or_misaligned_inputs(overrides, match):
+    with pytest.raises(ValueError, match=match):
+        _structural_request(**overrides)
+
+
+def test_provider_may_return_a_structural_result_from_an_external_package(monkeypatch):
+    request = _request()
+
+    def provider(actual_request):
+        return SimpleNamespace(
+            selected_logprobs=actual_request.logits[:, :1],
+            entropy=None,
+            backend_id="external.structural",
+            contract_id="external.structural.v1",
+            provenance={"tp_reduction": "provider_owned"},
+        )
+
+    path = _install_provider(monkeypatch, provider)
+    actual, entropy = compute_selected_logprobs(
+        args=SimpleNamespace(selected_logprob_provider=path, selected_logprob_provider_mode="strict"),
+        request=request,
+        native=_native,
+    )
+
+    assert entropy is None
+    torch.testing.assert_close(actual, request.logits[:, :1])
+
+
 def test_auto_mode_only_falls_back_for_explicit_unavailability(monkeypatch):
     request = _request()
     calls = {"native": 0}

tests/test_train_dump_utils.py

@@ -0,0 +1,23 @@
+from types import SimpleNamespace
+
+import torch
+
+from vime.utils.train_dump_utils import save_debug_train_data
+
+
+def test_debug_dump_adds_rank_and_replaces_atomically(tmp_path, monkeypatch):
+    monkeypatch.setattr(torch.distributed, "get_rank", lambda: 2)
+    args = SimpleNamespace(save_debug_train_data=str(tmp_path / "{rollout_id}.pt"))
+
+    save_debug_train_data(
+        args,
+        rollout_id=5,
+        rollout_data={"log_probs": [torch.tensor([1.0])]},
+    )
+
+    output = tmp_path / "5.rank2.pt"
+    payload = torch.load(output, weights_only=True)
+    assert payload["rollout_id"] == 5
+    assert payload["rank"] == 2
+    torch.testing.assert_close(payload["rollout_data"]["log_probs"][0], torch.tensor([1.0]))
+    assert list(tmp_path.glob(".*.tmp")) == []