From 4ea72626618ba4ed06fd9606075b703ccf1b875c Mon Sep 17 00:00:00 2001 From: Codex Date: Fri, 11 Sep 2026 01:56:21 +0800 Subject: [PATCH 1/8] perf(rocm): restore VIME TP4 decode throughput --- rl_engine/integrations/linear_logp.py | 6 +- .../ops/rocm/loss/vocab_parallel_logp.py | 251 +++++++++++++----- .../kernels/ops/triton/matmul/mfma_gemm.py | 30 ++- tests/test_rocm_mfma_gemm.py | 13 + 4 files changed, 231 insertions(+), 69 deletions(-) diff --git a/rl_engine/integrations/linear_logp.py b/rl_engine/integrations/linear_logp.py index 667f989f..60887442 100644 --- a/rl_engine/integrations/linear_logp.py +++ b/rl_engine/integrations/linear_logp.py @@ -20,7 +20,6 @@ ) from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import DEFAULT_NUM_VOCAB_TILES - _ALIGNMENT_DIAGNOSTIC_LOCK = threading.Lock() _ALIGNMENT_DIAGNOSTIC_CALLS = 0 @@ -450,8 +449,7 @@ def _rocm_contract( tp_rank=rank, tp_world_size=world, vocab_shard_bounds=tuple( - (index * local_vocab, (index + 1) * local_vocab) - for index in range(world) + (index * local_vocab, (index + 1) * local_vocab) for index in range(world) ), real_vocab_size=real_vocab_size, padded_vocab_size=global_vocab_size, @@ -509,6 +507,8 @@ def _rocm_from_local_logits( tp_group=tp_group, num_vocab_tiles=DEFAULT_NUM_VOCAB_TILES, deterministic=True, + targets_validated=True, + cache_preflight=target == "rollout", ) _record_alignment_diagnostics( target=target, diff --git a/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py index 7f58862f..e4d9b962 100644 --- a/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py +++ b/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py @@ -21,6 +21,7 @@ from __future__ import annotations +import hashlib from collections import OrderedDict from typing import Any @@ -30,11 +31,8 @@ from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( DEFAULT_NUM_VOCAB_TILES, VocabParallelLogprobOp, - _gather_tile_stats, _merge_tile_partials, - _preflight_cross_rank_agreement, _tile_size, - _validate_active_targets, _validate_invocation, ) @@ -48,6 +46,11 @@ _METADATA_CACHE_LIMIT = 32 _ACTIVE_MASK_CACHE: OrderedDict[tuple[Any, ...], torch.Tensor] = OrderedDict() _SHARD_START_CACHE: OrderedDict[tuple[Any, ...], torch.Tensor] = OrderedDict() +_PREFLIGHT_FINGERPRINT_CACHE: OrderedDict[tuple[Any, ...], torch.Tensor] = OrderedDict() +_VERIFIED_PREFLIGHT_CACHE: OrderedDict[tuple[Any, ...], None] = OrderedDict() +_LOGP_GATHER_CACHE: OrderedDict[tuple[Any, ...], tuple[torch.Tensor, list[torch.Tensor]]] = ( + OrderedDict() +) def _device_key(device: torch.device) -> tuple[str, int | None]: @@ -78,15 +81,11 @@ def _cached_active_mask( return cached, all_active -def _cached_shard_starts( - bounds: tuple[tuple[int, int], ...], device: torch.device -) -> torch.Tensor: +def _cached_shard_starts(bounds: tuple[tuple[int, int], ...], device: torch.device) -> torch.Tensor: key = (_device_key(device), bounds) cached = _SHARD_START_CACHE.get(key) if cached is None: - cached = torch.tensor( - [start for start, _ in bounds], dtype=torch.long, device=device - ) + cached = torch.tensor([start for start, _ in bounds], dtype=torch.long, device=device) _SHARD_START_CACHE[key] = cached if len(_SHARD_START_CACHE) > _METADATA_CACHE_LIMIT: _SHARD_START_CACHE.popitem(last=False) @@ -95,13 +94,128 @@ def _cached_shard_starts( return cached -def _gather_target_logit_cached( +def _cached_logp_gather_buffers( + *, + rows: int, + max_tiles: int, + world_size: int, + device: torch.device, +) -> tuple[torch.Tensor, list[torch.Tensor]]: + key = ( + _device_key(device), + int(world_size), + int(rows), + int(max_tiles), + ) + cached = _LOGP_GATHER_CACHE.get(key) + if cached is None: + width = 2 * max_tiles + 1 + local = torch.empty((rows, width), dtype=torch.float32, device=device) + gathered = [torch.empty_like(local) for _ in range(world_size)] + cached = (local, gathered) + _LOGP_GATHER_CACHE[key] = cached + if len(_LOGP_GATHER_CACHE) > _METADATA_CACHE_LIMIT: + _LOGP_GATHER_CACHE.popitem(last=False) + else: + _LOGP_GATHER_CACHE.move_to_end(key) + return cached + + +def _cached_preflight_fingerprint( + contract: LogprobContract, + num_vocab_tiles: int, + device: torch.device, +) -> tuple[torch.Tensor, bytes]: + payload = ( + f"{contract.cross_rank_fingerprint()}:{BACKEND_ID}:" f"{int(num_vocab_tiles)}:deterministic" + ).encode("utf-8") + digest = hashlib.sha256(payload).digest() + key = (_device_key(device), digest) + cached = _PREFLIGHT_FINGERPRINT_CACHE.get(key) + if cached is None: + cached = torch.tensor(tuple(digest), dtype=torch.uint8, device=device) + _PREFLIGHT_FINGERPRINT_CACHE[key] = cached + if len(_PREFLIGHT_FINGERPRINT_CACHE) > _METADATA_CACHE_LIMIT: + _PREFLIGHT_FINGERPRINT_CACHE.popitem(last=False) + else: + _PREFLIGHT_FINGERPRINT_CACHE.move_to_end(key) + return cached, digest + + +def _preflight_cross_rank_agreement_device( + contract: LogprobContract, + tp_group: Any, + num_vocab_tiles: int, + device: torch.device, + *, + cache_result: bool = False, +) -> None: + """Check the immutable contract without object collectives or host readback.""" + + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + raise LogprobContractError("vocab-parallel logprob requires initialized torch.distributed") + world = torch.distributed.get_world_size(group=tp_group) + local, digest = _cached_preflight_fingerprint( + contract, + num_vocab_tiles, + device, + ) + cache_key = ( + id(tp_group), + _device_key(device), + digest, + ) + if cache_result and cache_key in _VERIFIED_PREFLIGHT_CACHE: + _VERIFIED_PREFLIGHT_CACHE.move_to_end(cache_key) + return + gathered = torch.empty( + (world, local.numel()), + dtype=local.dtype, + device=device, + ) + torch.distributed.all_gather_into_tensor(gathered, local, group=tp_group) + agreement = torch.all(gathered == local.unsqueeze(0)) + if agreement.is_cuda: + torch._assert_async( + agreement, + "cross-rank preflight failed: TP ranks disagree on the logprob contract", + ) + elif not bool(agreement): + raise LogprobContractError( + "cross-rank preflight failed: TP ranks disagree on the logprob contract" + ) + if cache_result: + _VERIFIED_PREFLIGHT_CACHE[cache_key] = None + if len(_VERIFIED_PREFLIGHT_CACHE) > _METADATA_CACHE_LIMIT: + _VERIFIED_PREFLIGHT_CACHE.popitem(last=False) + + +def _validate_active_targets_device( + target_1d: torch.Tensor, + active_mask: torch.Tensor, + real_vocab_size: int, +) -> None: + valid = torch.all(~active_mask | ((target_1d >= 0) & (target_1d < real_vocab_size))) + if valid.is_cuda: + torch._assert_async( + valid, + f"active target_ids must lie in the real vocabulary [0, {real_vocab_size})", + ) + elif not bool(valid): + raise LogprobContractError( + f"active target_ids must lie in the real vocabulary [0, {real_vocab_size})" + ) + + +def _gather_logp_partials( + local_m: torch.Tensor, + local_s: torch.Tensor, z_masked: torch.Tensor, safe_target: torch.Tensor, contract: LogprobContract, tp_group: Any, -) -> torch.Tensor: - """ROCm copy of the exact owner gather with cached immutable metadata.""" +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Gather tile statistics and the selected logit with one TP collective.""" sharding = contract.sharding n = z_masked.shape[0] @@ -117,24 +231,40 @@ def _gather_target_logit_cached( ).contiguous() if sharding.tp_world_size == 1: - stacked = local_contrib.unsqueeze(0) - else: - if ( - not torch.distributed.is_available() - or not torch.distributed.is_initialized() - ): - raise LogprobContractError( - "vocab-parallel logprob requires initialized torch.distributed" - ) - gathered = [ - torch.empty_like(local_contrib) for _ in range(sharding.tp_world_size) - ] - torch.distributed.all_gather(gathered, local_contrib, group=tp_group) - stacked = torch.stack(gathered, dim=0) + return local_m.contiguous(), local_s.contiguous(), local_contrib.float() + + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + raise LogprobContractError("vocab-parallel logprob requires initialized torch.distributed") + tile_size = sharding.local_vocab_size // local_m.shape[1] + tile_counts = [(end - start) // tile_size for start, end in sharding.vocab_shard_bounds] + max_tiles = max(tile_counts) + packed, gathered = _cached_logp_gather_buffers( + rows=n, + max_tiles=max_tiles, + world_size=sharding.tp_world_size, + device=z_masked.device, + ) + local_tiles = local_m.shape[1] + packed[:, :local_tiles].copy_(local_m) + packed[:, max_tiles : max_tiles + local_tiles].copy_(local_s) + packed[:, -1].copy_(local_contrib) + # vLLM's ROCm full-graph path captures list-form all_gather, while + # all_gather_into_tensor remains an eager _allgather_base_ call per token. + torch.distributed.all_gather(gathered, packed, group=tp_group) + + m_parts = [gathered[rank][:, : tile_counts[rank]] for rank in range(sharding.tp_world_size)] + s_parts = [ + gathered[rank][:, max_tiles : max_tiles + tile_counts[rank]] + for rank in range(sharding.tp_world_size) + ] + m_all = torch.cat(m_parts, dim=1).contiguous() + s_all = torch.cat(s_parts, dim=1).contiguous() starts = _cached_shard_starts(sharding.vocab_shard_bounds, safe_target.device) owner = torch.bucketize(safe_target, starts, right=True) - 1 - return stacked[owner, rows] + stacked = torch.stack(gathered, dim=0) + target_logit = stacked[owner, rows, -1] + return m_all, s_all, target_logit def _native_backward_available() -> bool: @@ -176,9 +306,7 @@ class _RocmVocabParallelLogprobFunction(torch.autograd.Function): """ROCm tile statistics and backward with the shared WS2 merge contract.""" @staticmethod - def forward( - ctx, local_logits, target_1d, active_mask, contract, tp_group, tile, all_active - ): + def forward(ctx, local_logits, target_1d, active_mask, contract, tp_group, tile, all_active): sharding = contract.sharding shard = local_logits.contiguous() local_tiles = sharding.local_vocab_size // tile @@ -190,24 +318,19 @@ def forward( sharding.real_vocab_size, local_tiles, ) - tile_counts = [ - (end - start) // tile for start, end in sharding.vocab_shard_bounds - ] - m_all, s_all = _gather_tile_stats( - local_m.contiguous(), - local_s.contiguous(), - contract, - tp_group, - tile_counts, - ) safe_target = ( target_1d if all_active else torch.where(active_mask, target_1d, torch.zeros_like(target_1d)) ) - target_logit = _gather_target_logit_cached( - shard, safe_target, contract, tp_group - ).float() + m_all, s_all, target_logit = _gather_logp_partials( + local_m.contiguous(), + local_s.contiguous(), + shard, + safe_target, + contract, + tp_group, + ) lse = _merge_tile_partials(m_all, s_all) selected_logp = ( target_logit - lse @@ -245,15 +368,9 @@ def backward(ctx, grad_logp, grad_lse): ).contiguous() else: coef_logp = lse.new_zeros((rows,)) - target_local = torch.full( - (rows,), -1, dtype=torch.long, device=shard.device - ) + target_local = torch.full((rows,), -1, dtype=torch.long, device=shard.device) has_lse_grad = grad_lse is not None - coef_lse = ( - grad_lse.float().contiguous() - if has_lse_grad - else lse.new_zeros((rows,)) - ) + coef_lse = grad_lse.float().contiguous() if has_lse_grad else lse.new_zeros((rows,)) grad = _HipKernels.backward( shard, lse.contiguous(), @@ -275,33 +392,41 @@ def _apply_with_kernels( tp_group: Any, num_vocab_tiles: int, validate: bool, + targets_validated: bool, + cache_preflight: bool, ) -> tuple[torch.Tensor, torch.Tensor]: if not isinstance(contract, LogprobContract): raise LogprobContractError("contract must be a LogprobContract") tile = _tile_size(contract, num_vocab_tiles) _validate_invocation(local_logits, target_ids, contract, tp_group) - target_1d = target_ids.reshape(-1).to( - device=local_logits.device, dtype=torch.long - ) + target_1d = target_ids.reshape(-1).to(device=local_logits.device, dtype=torch.long) active_mask, all_active = _cached_active_mask(contract, local_logits.device) + if validate and not targets_validated: + _validate_active_targets_device(target_1d, active_mask, contract.sharding.real_vocab_size) if validate: - _validate_active_targets( - target_1d, active_mask, contract.sharding.real_vocab_size - ) if contract.sharding.tp_world_size > 1: - _preflight_cross_rank_agreement( - contract, tp_group, num_vocab_tiles, True + _preflight_cross_rank_agreement_device( + contract, + tp_group, + num_vocab_tiles, + local_logits.device, + cache_result=cache_preflight, ) selected_logp, lse = _RocmVocabParallelLogprobFunction.apply( local_logits, target_1d, active_mask, contract, tp_group, tile, all_active ) - if validate and bool((~torch.isfinite(lse) & active_mask).any().item()): - raise LogprobContractError( + if validate: + finite = torch.all(torch.isfinite(lse) | ~active_mask) + message = ( "non-finite logsumexp on an active row: logits over the real " "vocabulary must be finite for every active token" ) + if finite.is_cuda: + torch._assert_async(finite, message) + elif not bool(finite): + raise LogprobContractError(message) return selected_logp, lse @@ -322,6 +447,8 @@ def apply( num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES, validate: bool = True, deterministic: bool = True, + targets_validated: bool = False, + cache_preflight: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: if not deterministic: return super().apply( @@ -345,6 +472,8 @@ def apply( tp_group=tp_group, num_vocab_tiles=num_vocab_tiles, validate=validate, + targets_validated=targets_validated, + cache_preflight=cache_preflight, ) diff --git a/rl_engine/kernels/ops/triton/matmul/mfma_gemm.py b/rl_engine/kernels/ops/triton/matmul/mfma_gemm.py index e9bcd3fe..7c4a1847 100644 --- a/rl_engine/kernels/ops/triton/matmul/mfma_gemm.py +++ b/rl_engine/kernels/ops/triton/matmul/mfma_gemm.py @@ -59,6 +59,8 @@ class MfmaGemmConfig: _DECODE_CONFIG = MfmaGemmConfig(16, 32, 2, waves_per_eu=0, num_stages=2, group_m=1) +_QWEN_QKV_GATE_DECODE_CONFIG = MfmaGemmConfig(32, 64, 4, waves_per_eu=0, num_stages=2, group_m=1) +_QWEN_LM_HEAD_DECODE_CONFIG = MfmaGemmConfig(16, 128, 4, waves_per_eu=2, num_stages=2, group_m=1) _SMALL_CONFIG = MfmaGemmConfig(64, 128, 4, waves_per_eu=2, num_stages=2, group_m=8) _LARGE_CONFIG = MfmaGemmConfig(128, 128, 4, waves_per_eu=2, num_stages=2, group_m=8) @@ -66,8 +68,14 @@ class MfmaGemmConfig: def select_config(m_size: int, n_size: int, k_size: int) -> MfmaGemmConfig: """Pick a performance configuration. Never affects the result bits.""" - del n_size, k_size if m_size <= SPLIT_SCHEDULE_MAX_ROWS: + # Qwen3-8B TP4 decode is bandwidth-bound and benefits from more + # N-parallel programs on its widest projections. Keep the one-row + # QKV case on the lower-overhead default. + if k_size == 4096 and (n_size == 6144 or (n_size == 1536 and m_size > 1)): + return _QWEN_QKV_GATE_DECODE_CONFIG + if k_size == 4096 and n_size >= 32768: + return _QWEN_LM_HEAD_DECODE_CONFIG return _DECODE_CONFIG if m_size <= 1024: return _SMALL_CONFIG @@ -423,12 +431,24 @@ def warmup(device: torch.device | None = None) -> None: if torch.cuda.is_current_stream_capturing(): raise RuntimeError("warm the MFMA GEMM before HIP Graph capture") with torch.inference_mode(): - for k_size in (CHUNK_K, 2 * CHUNK_K + BLOCK_K, CHUNK_K + 8): + warmup_rows_and_configs = ( + (1, (_DECODE_CONFIG, _QWEN_LM_HEAD_DECODE_CONFIG)), + (4, (_QWEN_QKV_GATE_DECODE_CONFIG,)), + (SPLIT_SCHEDULE_MAX_ROWS + 1, (_SMALL_CONFIG,)), + (1025, (_LARGE_CONFIG,)), + ) + for k_size in ( + CHUNK_K, + 2 * CHUNK_K + BLOCK_K, + CHUNK_K + 8, + 4 * CHUNK_K, + ): b = torch.zeros((k_size, 64), dtype=torch.bfloat16, device=device) - for rows in (1, SPLIT_SCHEDULE_MAX_ROWS + 1, 1025): + for rows, configs in warmup_rows_and_configs: a = torch.zeros((rows, k_size), dtype=torch.bfloat16, device=device) - mfma_gemm(a, b) - mfma_gemm(a, b.t().contiguous().t()) + for config in configs: + mfma_gemm(a, b, config=config) + mfma_gemm(a, b.t().contiguous().t(), config=config) torch.cuda.synchronize(device) diff --git a/tests/test_rocm_mfma_gemm.py b/tests/test_rocm_mfma_gemm.py index 1f12c201..9ad7aa8a 100644 --- a/tests/test_rocm_mfma_gemm.py +++ b/tests/test_rocm_mfma_gemm.py @@ -45,6 +45,19 @@ def _configs(): ] +def test_qwen_tp4_decode_config_selection(): + default = M.MfmaGemmConfig(16, 32, 2, waves_per_eu=0, num_stages=2, group_m=1) + wide = M.MfmaGemmConfig(32, 64, 4, waves_per_eu=0, num_stages=2, group_m=1) + vocab = M.MfmaGemmConfig(16, 128, 4, waves_per_eu=2, num_stages=2, group_m=1) + + assert M.select_config(1, 1536, 4096) == default + assert M.select_config(4, 1536, 4096) == wide + assert M.select_config(4, 6144, 4096) == wide + assert M.select_config(4, 37984, 4096) == vocab + assert M.select_config(4, 4096, 1024) == default + assert M.select_config(4, 4096, 3072) == default + + @pytest.mark.parametrize("k_size,n_size", QWEN_TP4_SHAPES) def test_forward_matches_fp32_reference(k_size, n_size): a = _rand(300, k_size) From 8e4c8ec253ac148b4164ec7e0af9f2276dc44fe8 Mon Sep 17 00:00:00 2001 From: Codex Date: Fri, 11 Sep 2026 08:54:20 +0000 Subject: [PATCH 2/8] fix(cuda): restore VIME module ablation evidence --- examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py | 13 ++ .../vime_qwen3_8b_tp4_cp2_200/validate_run.py | 130 ++++++++---------- .../tis_metrics.py | 17 ++- rl_engine/distributed/collectives.py | 15 +- ...test_transport_deterministic_collective.py | 4 + tests/test_vime_tp4_example.py | 22 +++ 6 files changed, 120 insertions(+), 81 deletions(-) diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py b/examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py index cb07b8fb..def6a1fd 100755 --- a/examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py @@ -71,6 +71,9 @@ class Arm: # backend auto-selection. MEGATRON_ATTENTION_BACKEND = "fused" RL_KERNEL_LINEAR_LOGP_PROVIDER = "rl_engine.integrations.vime.linear_logp_provider.provider" +RL_KERNEL_MISMATCH_METRICS_HOOK = ( + "examples.vime_rocm_attention_ablation.tis_metrics.metrics_only_tis" +) MODEL_ARGS = ( "--swiglu", @@ -135,6 +138,14 @@ def _linear_logp_provider_args(arm: Arm) -> tuple[str, ...]: raise ValueError(f"unsupported training logp route: {arm.logp_case!r}") +def _mismatch_metrics_args() -> tuple[str, ...]: + return ( + "--get-mismatch-metrics", + "--custom-tis-function-path", + RL_KERNEL_MISMATCH_METRICS_HOOK, + ) + + def _path(value: str | None, label: str) -> Path: if not value: raise ValueError(f"{label} is required (argument or environment variable)") @@ -346,6 +357,7 @@ def main(argv: list[str] | None = None) -> int: "RL_KERNEL_FFN_CASE": arm.ffn_case, "RL_KERNEL_LOGP_CASE": arm.logp_case, "RL_KERNEL_READBACK_DIR": str(run_dir / "readbacks"), + "RL_KERNEL_MISMATCH_SIDECAR_DIR": str(run_dir / "mismatch-sidecars"), "RL_KERNEL_VLLM_REAL_VOCAB_SIZE": "151936", "RL_KERNEL_VLLM_PADDED_VOCAB_SIZE": "152064", "RL_KERNEL_VLLM_TEMPERATURE": "1.0", @@ -456,6 +468,7 @@ def main(argv: list[str] | None = None) -> int: str(TOPOLOGY["rollout_gpus_per_engine"]), "--vllm-gpu-memory-utilization", str(args.vllm_gpu_memory_utilization), + *_mismatch_metrics_args(), ] if arm.framework_use_rollout_logprobs: train_command.append("--use-rollout-logprobs") diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py b/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py index 70e56ca9..e12e7e45 100755 --- a/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py @@ -18,7 +18,6 @@ from rl_engine.integrations.runtime import _contains_triton, _runtime_platform - RECORD_RE = re.compile(r"\b(rollout|step|perf)\s+(\d+):\s+(\{.*\})\s*$") FRAMEWORKS = (("megatron", "training"), ("vllm", "rollout")) MODULES = ("attention", "ffn", "logp") @@ -40,14 +39,13 @@ "ffn": "ffn_case", "logp": "logp_case", } -RL_KERNEL_LINEAR_LOGP_PROVIDER = ( - "rl_engine.integrations.vime.linear_logp_provider.provider" -) +RL_KERNEL_LINEAR_LOGP_PROVIDER = "rl_engine.integrations.vime.linear_logp_provider.provider" VIME_NATIVE_LINEAR_LOGP_MARKER = ( "linear_logp native active: " "backend_id=vime.utils.ppo_utils.calculate_log_probs_and_entropy " "contract_id=vime.native.linear_logp.v1 route=unconfigured device=cuda" ) +RL_KERNEL_MISMATCH_SIDECAR_MARKER = "rlkernel mismatch sidecar active: logp_case=" CUDA_GRAPH_LAUNCHER_MARKERS = ( "required vLLM full-decode CUDA Graph capture sizes", "strict vLLM full-decode CUDA Graph capture sizes", @@ -61,6 +59,23 @@ def _load_json(path: Path) -> dict[str, Any]: return value +def _compare_mismatch_sidecars( + directory: Path, + *, + require_exact: bool, + tensor_parallel_size: int, + context_parallel_size: int, +) -> dict[str, Any]: + from examples.vime_rocm_attention_ablation.validate_artifacts import compare_train_rollout_logps + + return compare_train_rollout_logps( + directory, + require_exact=require_exact, + tensor_parallel_size=tensor_parallel_size, + context_parallel_size=context_parallel_size, + ) + + def _parse_runtime_records(log_text: str) -> dict[str, dict[int, dict[str, Any]]]: records: dict[str, dict[int, dict[str, Any]]] = { "rollout": {}, @@ -82,14 +97,11 @@ def _parse_runtime_records(log_text: str) -> dict[str, dict[int, dict[str, Any]] def _validate_cudagraph(log_text: str, manifest: Mapping[str, Any]) -> dict[str, Any]: execution = manifest.get("vllm_execution", {}) - capture_sizes = ( - execution.get("capture_sizes", []) if isinstance(execution, Mapping) else [] - ) + capture_sizes = execution.get("capture_sizes", []) if isinstance(execution, Mapping) else [] compact_sizes = "[" + ",".join(str(value) for value in capture_sizes) + "]" checks = { "launcher_marker": any( - f"{marker}: {compact_sizes}" in log_text - for marker in CUDA_GRAPH_LAUNCHER_MARKERS + f"{marker}: {compact_sizes}" in log_text for marker in CUDA_GRAPH_LAUNCHER_MARKERS ), "engine_mode": bool(re.search(r"cudagraph_mode.*FULL_DECODE_ONLY", log_text)), "not_eager": "enforce_eager=False" in log_text, @@ -161,19 +173,12 @@ def _validate_readbacks( records = [ value["operators"][module] for value in matching - if isinstance(value.get("operators"), Mapping) - and module in value["operators"] + if isinstance(value.get("operators"), Mapping) and module in value["operators"] ] - installed_count = sum( - module in value.get("installed_hooks", {}) for value in matching - ) + installed_count = sum(module in value.get("installed_hooks", {}) for value in matching) call_count = sum(int(record.get("call_count", 0)) for record in records) - implementations = sorted( - {str(record.get("implementation", "")) for record in records} - ) - backend_ids = sorted( - {str(record.get("backend_id", "")) for record in records} - ) + implementations = sorted({str(record.get("implementation", "")) for record in records}) + backend_ids = sorted({str(record.get("backend_id", "")) for record in records}) case_ids = sorted({str(record.get("case_id", "")) for record in records}) native_megatron_logp = ( framework == "megatron" @@ -182,15 +187,16 @@ def _validate_readbacks( and expected == "production" ) if native_megatron_logp: - marker_present = VIME_NATIVE_LINEAR_LOGP_MARKER in log_text + marker_present = ( + VIME_NATIVE_LINEAR_LOGP_MARKER in log_text + or f"{RL_KERNEL_MISMATCH_SIDECAR_MARKER}{case_id}" in log_text + ) if installed_count: errors.append( f"{label} production logp unexpectedly installed an RL-Kernel hook" ) if records: - errors.append( - f"{label} production logp unexpectedly entered provider readback" - ) + errors.append(f"{label} production logp unexpectedly entered provider readback") if not marker_present: errors.append( f"{label} production logp did not report Vime's native backend marker" @@ -203,9 +209,7 @@ def _validate_readbacks( "implementations": implementations, "backend_ids": backend_ids, "native_marker_present": marker_present, - "native_backend_id": ( - "vime.utils.ppo_utils.calculate_log_probs_and_entropy" - ), + "native_backend_id": ("vime.utils.ppo_utils.calculate_log_probs_and_entropy"), } continue if installed_count == 0: @@ -225,20 +229,17 @@ def _validate_readbacks( errors.append(f"{label} {module} did not report CUDA execution") if record.get("provenance", {}).get("fallback") is True: errors.append(f"{label} {module} provenance recorded fallback") - if expected == "rl_kernel" and not str( - record.get("backend_id", "") - ).startswith("rlkernel."): + if expected == "rl_kernel" and not str(record.get("backend_id", "")).startswith( + "rlkernel." + ): errors.append(f"{label} {module} did not use an RL-Kernel backend") provenance = record.get("provenance", {}) reported_backend_ids = _reported_backend_ids(record) - strict_execution = ( - isinstance(provenance, Mapping) - and ( - provenance.get("deterministic_linear_logp") is True - or ( - isinstance(provenance.get("execution"), Mapping) - and provenance["execution"].get("strict_backend") is True - ) + strict_execution = isinstance(provenance, Mapping) and ( + provenance.get("deterministic_linear_logp") is True + or ( + isinstance(provenance.get("execution"), Mapping) + and provenance["execution"].get("strict_backend") is True ) ) if expected == "production" and ( @@ -320,8 +321,7 @@ def _validate_runtime_logprobs( if len(rows) != expected_rounds: errors.append(f"observed {len(rows)} train steps, expected {expected_rounds}") bitwise_zero = bool(rows) and all( - row["bitwise_mismatch_count"] == 0.0 and row["max_abs_dlogp"] == 0.0 - for row in rows + row["bitwise_mismatch_count"] == 0.0 and row["max_abs_dlogp"] == 0.0 for row in rows ) if require_zero and not bitwise_zero: errors.append("R/R arm did not achieve bitwise-zero runtime metrics") @@ -334,9 +334,7 @@ def _validate_runtime_logprobs( ), "bitwise_zero": bitwise_zero, "rows": rows, - "total_active_token_exposure": sum( - row["active_token_count"] or 0.0 for row in rows - ), + "total_active_token_exposure": sum(row["active_token_count"] or 0.0 for row in rows), } @@ -345,9 +343,7 @@ def _inspect_offline_dumps(directory: Path) -> dict[str, Any]: comparable = 0 for path in paths: payload = torch.load(path, map_location="cpu", weights_only=False) - rollout_data = ( - payload.get("rollout_data", {}) if isinstance(payload, Mapping) else {} - ) + rollout_data = payload.get("rollout_data", {}) if isinstance(payload, Mapping) else {} if isinstance(rollout_data, Mapping) and "log_probs" in rollout_data: comparable += 1 return { @@ -357,7 +353,10 @@ def _inspect_offline_dumps(directory: Path) -> dict[str, Any]: "reason": ( None if paths and comparable == len(paths) - else "current VIME dump lacks captured training log_probs; runtime exact metrics are used" + else ( + "current VIME dump lacks captured training log_probs; " + "runtime exact metrics are used" + ) ), } @@ -371,21 +370,22 @@ def validate_run(run_dir: Path) -> dict[str, Any]: records = _parse_runtime_records(log_text) require_zero = all(str(arm[CASE_FIELDS[module]]) == "R/R" for module in MODULES) cudagraph = _validate_cudagraph(log_text, manifest) - readbacks = _validate_readbacks( - _load_readbacks(run_dir / "readbacks"), arm, log_text - ) - logprobs = _validate_runtime_logprobs( + readbacks = _validate_readbacks(_load_readbacks(run_dir / "readbacks"), arm, log_text) + runtime_logprobs = _validate_runtime_logprobs( records["step"], int(manifest["num_rollout"]), int(manifest["batching"]["global_batch_size"]), - require_zero, + False, + ) + logprobs = _compare_mismatch_sidecars( + run_dir / "mismatch-sidecars", + require_exact=require_zero, + tensor_parallel_size=int(manifest["topology"]["tp"]), + context_parallel_size=int(manifest["topology"]["cp"]), ) global_errors = [] algorithm = manifest.get("algorithm", {}) - if ( - not isinstance(algorithm, Mapping) - or algorithm.get("advantage_estimator") != "grpo" - ): + if not isinstance(algorithm, Mapping) or algorithm.get("advantage_estimator") != "grpo": global_errors.append("manifest does not explicitly select GRPO") train_command = manifest.get("train_command", []) expected_algorithm_pair = ["--advantage-estimator", "grpo"] @@ -433,9 +433,7 @@ def validate_run(run_dir: Path) -> dict[str, Any]: "production Megatron logp must not configure a linear_logp provider" ) if "--linear-logp-provider-mode" in train_command: - global_errors.append( - "production Megatron logp must not configure provider mode" - ) + global_errors.append("production Megatron logp must not configure provider mode") elif not has_provider or not has_strict_mode: global_errors.append( "RL-Kernel Megatron logp must configure the strict RL-Kernel provider" @@ -446,9 +444,7 @@ def validate_run(run_dir: Path) -> dict[str, Any]: "recompute_num_layers": 1, } if manifest.get("training_memory") != expected_recompute: - global_errors.append( - "manifest does not contain the required recompute configuration" - ) + global_errors.append("manifest does not contain the required recompute configuration") if re.search(r"fallback=true", log_text, re.IGNORECASE): global_errors.append("run log contains fallback=true") if "Traceback (most recent call last)" in log_text: @@ -458,15 +454,13 @@ def validate_run(run_dir: Path) -> dict[str, Any]: "run_id": manifest.get("run_id"), "group": arm.get("group"), "passed": bool( - cudagraph["passed"] - and readbacks["passed"] - and logprobs["passed"] - and not global_errors + cudagraph["passed"] and readbacks["passed"] and logprobs["passed"] and not global_errors ), "errors": global_errors, "cudagraph": cudagraph, "runtime_readbacks": readbacks, "train_rollout_logprob": logprobs, + "runtime_scalar_logprob": runtime_logprobs, "offline_tensor_comparison": _inspect_offline_dumps(run_dir / "train-data"), } return report @@ -488,9 +482,7 @@ def main(argv: list[str] | None = None) -> int: "passed": False, "errors": [f"{type(exc).__name__}: {exc}"], } - output.write_text( - json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8" - ) + output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") print(json.dumps(report, indent=2, sort_keys=True)) if args.seal and report["passed"]: (run_dir / "COMPLETE").touch(exist_ok=False) diff --git a/examples/vime_rocm_attention_ablation/tis_metrics.py b/examples/vime_rocm_attention_ablation/tis_metrics.py index c71e4e30..f866cc20 100644 --- a/examples/vime_rocm_attention_ablation/tis_metrics.py +++ b/examples/vime_rocm_attention_ablation/tis_metrics.py @@ -6,6 +6,7 @@ from __future__ import annotations import itertools +import logging import os import threading from pathlib import Path @@ -16,9 +17,12 @@ SIDECAR_SCHEMA_VERSION = "rlkernel.vime_rocm_attention_mismatch_sidecar.v1" SIDECAR_DIRECTORY_ENV = "RL_KERNEL_MISMATCH_SIDECAR_DIR" +NATIVE_LOGP_SIDECAR_MARKER = "rlkernel mismatch sidecar active: logp_case=" _CALL_COUNTER = itertools.count() _CALL_COUNTER_LOCK = threading.Lock() +_LOGGED_CASES: set[str] = set() +logger = logging.getLogger(__name__) def _cpu_vector(value: Any, *, label: str) -> torch.Tensor: @@ -76,8 +80,7 @@ def _write_sidecar( _cpu_vector(value, label="train_log_probs") for value in train_log_probs ], "rollout_log_probs": [ - _cpu_vector(value, label="rollout_log_probs") - for value in rollout_log_probs + _cpu_vector(value, label="rollout_log_probs") for value in rollout_log_probs ], "loss_masks": [_cpu_vector(value, label="loss_masks") for value in loss_masks], "total_lengths": [int(value) for value in total_lengths], @@ -92,6 +95,15 @@ def _write_sidecar( os.replace(temporary, path) +def _log_route_identity() -> None: + case_id = os.environ.get("RL_KERNEL_LOGP_CASE", "").strip() or "unknown" + with _CALL_COUNTER_LOCK: + if case_id in _LOGGED_CASES: + return + _LOGGED_CASES.add(case_id) + logger.info("%s%s", NATIVE_LOGP_SIDECAR_MARKER, case_id) + + def metrics_only_tis( args: Any, *, @@ -127,6 +139,7 @@ def metrics_only_tis( "tis_clipfrac": (clipped != ratio).to(dtype=ratio.dtype), "tis_abs": (ratio - 1).abs(), } + _log_route_identity() _write_sidecar( args, train_log_probs=train_log_probs, diff --git a/rl_engine/distributed/collectives.py b/rl_engine/distributed/collectives.py index 0b8ecb9a..8b7b99ef 100644 --- a/rl_engine/distributed/collectives.py +++ b/rl_engine/distributed/collectives.py @@ -32,9 +32,7 @@ def _deterministic_all_reduce_(input: torch.Tensor, collective_handle: int) -> N from rl_engine import _C if torch.version.hip is not None: - _C.deterministic_collective_rocm_ipc_all_reduce_input( - collective_handle, input, input - ) + _C.deterministic_collective_rocm_ipc_all_reduce_input(collective_handle, input, input) else: _C.deterministic_collective_all_reduce_fused(collective_handle, input, input) @@ -65,9 +63,7 @@ def _deterministic_staging_reserve_(staging: torch.Tensor, collective_handle: in from rl_engine import _C if torch.version.hip is not None: - _C.deterministic_collective_rocm_ipc_prepare_staged( - collective_handle, staging - ) + _C.deterministic_collective_rocm_ipc_prepare_staged(collective_handle, staging) else: _C.deterministic_collective_prepare_staged(collective_handle, staging) @@ -91,9 +87,7 @@ def _deterministic_staged_all_reduce( output = torch.empty_like(staging) if torch.version.hip is not None: - _C.deterministic_collective_rocm_ipc_all_reduce_staged( - collective_handle, staging, output - ) + _C.deterministic_collective_rocm_ipc_all_reduce_staged(collective_handle, staging, output) else: _C.deterministic_collective_all_reduce_staged(collective_handle, staging, output) return output @@ -140,6 +134,8 @@ class DeterministicCollective: host barrier and advance correctly during CUDA Graph replay. """ + backend_id = "cuda_ipc_fixed_tree" + def __init__( self, group: dist.ProcessGroup | None = None, @@ -626,7 +622,6 @@ def collective_for_group( TorchDistributedDeterministicCollective, ) - __all__ = [ "DETERMINISTIC_ALL_REDUCE_OP", "DeterministicCollective", diff --git a/tests/distributed/test_transport_deterministic_collective.py b/tests/distributed/test_transport_deterministic_collective.py index 23621ee0..ed5217d6 100644 --- a/tests/distributed/test_transport_deterministic_collective.py +++ b/tests/distributed/test_transport_deterministic_collective.py @@ -105,6 +105,10 @@ def test_public_exports_use_canonical_collectives_module() -> None: ) +def test_cuda_collective_has_stable_backend_identity() -> None: + assert collectives.DeterministicCollective.backend_id == "cuda_ipc_fixed_tree" + + def _make_collective( monkeypatch: pytest.MonkeyPatch, peer_inputs: list[torch.Tensor], diff --git a/tests/test_vime_tp4_example.py b/tests/test_vime_tp4_example.py index 6bed6327..f07424c5 100644 --- a/tests/test_vime_tp4_example.py +++ b/tests/test_vime_tp4_example.py @@ -11,9 +11,11 @@ MEGATRON_ATTENTION_BACKEND, RL_KERNEL_LINEAR_LOGP_PROVIDER, _linear_logp_provider_args, + _mismatch_metrics_args, ) from examples.vime_qwen3_8b_tp4_cp2_200.run_supplement_suite import specs from examples.vime_qwen3_8b_tp4_cp2_200.validate_run import ( + RL_KERNEL_MISMATCH_SIDECAR_MARKER, VIME_NATIVE_LINEAR_LOGP_MARKER, _validate_readbacks, ) @@ -147,6 +149,14 @@ def test_module_ablation_logp_provider_follows_training_route(): assert expected_module_groups == {key for key in ARMS if key.startswith("M")} +def test_module_ablation_enables_rlkernel_mismatch_sidecars(): + assert _mismatch_metrics_args() == ( + "--get-mismatch-metrics", + "--custom-tis-function-path", + "examples.vime_rocm_attention_ablation.tis_metrics.metrics_only_tis", + ) + + def test_supplement_suite_uses_short_module_and_three_seed_precision_designs(): module = specs("module") assert len(module) == 8 @@ -172,6 +182,18 @@ def test_validator_accepts_native_vime_logp_evidence_for_production_arm(): assert training_logp["call_count"] == 0 +def test_validator_accepts_rlkernel_sidecar_evidence_for_production_logp(): + report = _validate_readbacks( + _production_readbacks(), + asdict(ARMS["G10"]), + f"{RL_KERNEL_MISMATCH_SIDECAR_MARKER}P/P", + ) + + assert report["passed"] + training_logp = report["frameworks"]["megatron/training"]["modules"]["logp"] + assert training_logp["native_marker_present"] + + def test_validator_rejects_provider_readback_on_production_megatron_logp(): readbacks = _production_readbacks() contaminated = _production_operator("megatron", "training", "logp") From a79b94176376fea3c23e910be00d23febbb15b19 Mon Sep 17 00:00:00 2001 From: Codex Date: Fri, 11 Sep 2026 09:03:50 +0000 Subject: [PATCH 3/8] fix(megatron): keep runtime attention state out of checkpoints --- rl_engine/integrations/megatron_runtime.py | 19 +++---- tests/test_megatron_runtime_state.py | 65 ++++++++++++++++++++++ 2 files changed, 74 insertions(+), 10 deletions(-) create mode 100644 tests/test_megatron_runtime_state.py diff --git a/rl_engine/integrations/megatron_runtime.py b/rl_engine/integrations/megatron_runtime.py index 596b0645..58c956ec 100644 --- a/rl_engine/integrations/megatron_runtime.py +++ b/rl_engine/integrations/megatron_runtime.py @@ -93,7 +93,9 @@ def _save_megatron_layer_diagnostic( try: enabled_layers = {int(value.strip()) for value in requested_layers.split(",")} except ValueError as exc: - raise RuntimeError("RL_KERNEL_ALIGNMENT_LAYERS must be comma-separated integers") from exc + raise RuntimeError( + "RL_KERNEL_ALIGNMENT_LAYERS must be comma-separated integers" + ) from exc if layer not in enabled_layers: return if input_value.ndim == 3: @@ -136,8 +138,7 @@ def _save_megatron_layer_diagnostic( payload, output_dir / ( - f"megatron-pid{os.getpid()}-rank{rank:05d}-layer{layer:02d}-" - f"call{call_index:08d}.pt" + f"megatron-pid{os.getpid()}-rank{rank:05d}-layer{layer:02d}-" f"call{call_index:08d}.pt" ), ) @@ -252,8 +253,7 @@ def _strict_rocm_rope_positions( if cu_seqlens is None: raise RuntimeError("strict ROCm THD RoPE requires cu_seqlens") values = tuple( - int(value) - for value in cu_seqlens.detach().to(device="cpu", dtype=torch.int64).tolist() + int(value) for value in cu_seqlens.detach().to(device="cpu", dtype=torch.int64).tolist() ) if len(values) < 2 or values[0] != 0: raise RuntimeError("strict ROCm THD RoPE received invalid cu_seqlens") @@ -359,9 +359,7 @@ def strict_apply_rotary_pos_emb( def _install_torch_dist_object_compatibility() -> None: """Normalize the PyTorch DCP object shape expected by this Megatron revision.""" - strategy = importlib.import_module( - "megatron.core.dist_checkpointing.strategies.torch" - ) + strategy = importlib.import_module("megatron.core.dist_checkpointing.strategies.torch") original = strategy._replace_sharded_keys_with_state_dict_keys if getattr(original, "__rl_kernel_dcp_object_compatibility__", False): return @@ -766,8 +764,8 @@ def attention_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: ) setattr(qkv, _STRICT_ATTENTION_PROJECTION_MARKER, "qkv") setattr(projection, _STRICT_ATTENTION_PROJECTION_MARKER, "o_proj") - setattr(qkv, _STRICT_ATTENTION_CORE_MARKER, core_attention) - setattr(projection, _STRICT_ATTENTION_CORE_MARKER, core_attention) + object.__setattr__(qkv, _STRICT_ATTENTION_CORE_MARKER, core_attention) + object.__setattr__(projection, _STRICT_ATTENTION_CORE_MARKER, core_attention) if copy_to_tp is None: bind_collective_identity( qkv, @@ -793,6 +791,7 @@ def attention_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: callback_backend(reduce_from_tp), ) if hasattr(qkv, "layer_norm_weight"): + def te_qkv_forward(module: Any, input_value: torch.Tensor) -> Any: normalized = _fused_rms_norm_input(module, input_value, "linear_qkv") normalized = strict_tp_copy(module, core_attention, normalized) diff --git a/tests/test_megatron_runtime_state.py b/tests/test_megatron_runtime_state.py new file mode 100644 index 00000000..8959f8be --- /dev/null +++ b/tests/test_megatron_runtime_state.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import torch + +from rl_engine.integrations.megatron_runtime import _patch_strict_attention_projections + + +def test_strict_attention_runtime_core_does_not_pollute_state_dict(): + class CoreAttention(torch.nn.Module): + def get_extra_state(self): + return {"runtime": True} + + class ColumnLinear(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.eye(2)) + self.gather_output = False + self.skip_bias_add = False + self.bias = None + self.allreduce_dgrad = True + + def _forward_impl(self, input, weight, *args, **kwargs): + del args, kwargs + return input @ weight.t() + + class RowLinear(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.eye(2)) + self.input_is_parallel = True + self.skip_bias_add = False + self.return_bias = False + self.bias = None + + def _forward_impl(self, input, weight, *args, **kwargs): + del args, kwargs + return input @ weight.t() + + class SelfAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear_qkv = ColumnLinear() + self.linear_proj = RowLinear() + self.core_attention = CoreAttention() + + _patch_strict_attention_projections( + self_attention_cls=SelfAttention, + column_linear_cls=ColumnLinear, + row_linear_cls=RowLinear, + det_gemm=lambda lhs, rhs: lhs @ rhs, + copy_to_tp=lambda value: value, + reduce_from_tp=lambda value: value, + ) + + attention = SelfAttention() + + assert ( + getattr( + attention.linear_proj, + "__rl_kernel_strict_attention_core__", + ) + is attention.core_attention + ) + assert not any("__rl_kernel_strict_attention_core__" in key for key in attention.state_dict()) From 112d9d67c89899f1cf007b482f656e29843e3939 Mon Sep 17 00:00:00 2001 From: Codex Date: Fri, 11 Sep 2026 09:14:21 +0000 Subject: [PATCH 4/8] fix(vime): expose mismatch sidecars to Ray workers --- examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py | 11 +++++++---- examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py | 7 ++++++- tests/test_vime_tp4_example.py | 2 +- 3 files changed, 14 insertions(+), 6 deletions(-) diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py b/examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py index def6a1fd..b189d8d2 100755 --- a/examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py @@ -71,9 +71,7 @@ class Arm: # backend auto-selection. MEGATRON_ATTENTION_BACKEND = "fused" RL_KERNEL_LINEAR_LOGP_PROVIDER = "rl_engine.integrations.vime.linear_logp_provider.provider" -RL_KERNEL_MISMATCH_METRICS_HOOK = ( - "examples.vime_rocm_attention_ablation.tis_metrics.metrics_only_tis" -) +RL_KERNEL_MISMATCH_METRICS_HOOK = "vime_rocm_attention_ablation.tis_metrics.metrics_only_tis" MODEL_ARGS = ( "--swiglu", @@ -329,7 +327,12 @@ def main(argv: list[str] | None = None) -> int: run_dir = args.output_root.expanduser().resolve() / run_id run_dir.mkdir(parents=True, exist_ok=False) - pythonpath = [str(rl_kernel_root), str(vime_root), str(megatron_root)] + pythonpath = [ + str(rl_kernel_root / "examples"), + str(rl_kernel_root), + str(vime_root), + str(megatron_root), + ] pythonpath.extend(str(Path(item).expanduser().resolve()) for item in args.extra_pythonpath) if os.environ.get("PYTHONPATH"): pythonpath.extend(item for item in os.environ["PYTHONPATH"].split(os.pathsep) if item) diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py b/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py index e12e7e45..2e89356a 100755 --- a/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py @@ -66,7 +66,12 @@ def _compare_mismatch_sidecars( tensor_parallel_size: int, context_parallel_size: int, ) -> dict[str, Any]: - from examples.vime_rocm_attention_ablation.validate_artifacts import compare_train_rollout_logps + import sys + + examples_root = Path(__file__).resolve().parents[1] + if str(examples_root) not in sys.path: + sys.path.insert(0, str(examples_root)) + from vime_rocm_attention_ablation.validate_artifacts import compare_train_rollout_logps return compare_train_rollout_logps( directory, diff --git a/tests/test_vime_tp4_example.py b/tests/test_vime_tp4_example.py index f07424c5..bf8a205d 100644 --- a/tests/test_vime_tp4_example.py +++ b/tests/test_vime_tp4_example.py @@ -153,7 +153,7 @@ def test_module_ablation_enables_rlkernel_mismatch_sidecars(): assert _mismatch_metrics_args() == ( "--get-mismatch-metrics", "--custom-tis-function-path", - "examples.vime_rocm_attention_ablation.tis_metrics.metrics_only_tis", + "vime_rocm_attention_ablation.tis_metrics.metrics_only_tis", ) From d78ef9694350a89db4ee32d25e76198f3ef06532 Mon Sep 17 00:00:00 2001 From: Codex Date: Fri, 11 Sep 2026 11:06:17 +0000 Subject: [PATCH 5/8] fix(rocm): validate mixed VIME module routes --- .../validate_module_artifacts.py | 164 ++++++++++++++++++ tests/test_vime_rocm_module_validation.py | 93 ++++++++++ 2 files changed, 257 insertions(+) create mode 100644 examples/vime_rocm_attention_ablation/validate_module_artifacts.py create mode 100644 tests/test_vime_rocm_module_validation.py diff --git a/examples/vime_rocm_attention_ablation/validate_module_artifacts.py b/examples/vime_rocm_attention_ablation/validate_module_artifacts.py new file mode 100644 index 00000000..246a40ac --- /dev/null +++ b/examples/vime_rocm_attention_ablation/validate_module_artifacts.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Validate ROCm Vime matrices with independently selected operator routes.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +from .validate_artifacts import ( + CASE_IMPLEMENTATIONS, + FRAMEWORK_TARGETS, + _case_id, + _validate_production_record, + _validate_rlkernel_record, + _validate_strict_dense_record, +) + +MODULES = ("attention", "ffn", "logp") +RL_KERNEL_MISMATCH_SIDECAR_MARKER = "rlkernel mismatch sidecar active: logp_case=" + + +def validate_module_readbacks( + readbacks: Sequence[Mapping[str, Any]], + cases: Mapping[str, str], + *, + log_text: str = "", +) -> dict[str, Any]: + """Validate a ROCm module matrix with independently selected operator routes.""" + + normalized_cases = {module: _case_id(str(cases[module])) for module in MODULES} + errors: list[str] = [] + frameworks: dict[str, Any] = {} + + for framework, target in FRAMEWORK_TARGETS: + label = f"{framework}/{target}" + matching = [ + value + for value in readbacks + if value.get("framework") == framework and value.get("target") == target + ] + if not matching: + errors.append(f"missing {label} readback") + continue + + for value in matching: + plan = value.get("plan") + plan_cases = plan.get("cases") if isinstance(plan, Mapping) else None + if not isinstance(plan_cases, Mapping): + errors.append(f"{label}: readback does not contain an integration plan") + else: + for module, case_id in normalized_cases.items(): + item = plan_cases.get(module) + if not isinstance(item, Mapping) or item.get("case_id") != case_id: + errors.append(f"{label}: readback {module} plan is not {case_id}") + if value.get("fallbacks"): + errors.append(f"{label} recorded fallback: {value['fallbacks']}") + + module_summary: dict[str, Any] = {} + for module, case_id in normalized_cases.items(): + module_label = f"{label} {module}" + expected = CASE_IMPLEMENTATIONS[case_id][target] + records = [ + value["operators"][module] + for value in matching + if isinstance(value.get("operators"), Mapping) + and isinstance(value["operators"].get(module), Mapping) + ] + hook_count = sum( + isinstance(value.get("installed_hooks"), Mapping) + and bool(value["installed_hooks"].get(module)) + for value in matching + ) + call_count = sum(int(record.get("call_count", 0)) for record in records) + native_megatron_logp = ( + framework == "megatron" + and target == "training" + and module == "logp" + and expected == "production" + ) + if native_megatron_logp: + marker_present = f"{RL_KERNEL_MISMATCH_SIDECAR_MARKER}{case_id}" in log_text + if hook_count: + errors.append(f"{module_label} production route installed an RL-Kernel hook") + if records: + errors.append(f"{module_label} production route entered provider readback") + if not marker_present: + errors.append( + f"{module_label} production route lacks mismatch-sidecar evidence" + ) + module_summary[module] = { + "case_id": case_id, + "expected_implementation": expected, + "installed_processes": hook_count, + "call_count": call_count, + "implementations": [], + "backend_ids": [], + "native_marker_present": marker_present, + } + continue + + if hook_count == 0: + errors.append(f"{module_label} hook was not installed") + if not records: + errors.append(f"missing {module_label} execution record") + if call_count <= 0: + errors.append(f"{module_label} had zero calls") + + implementations: set[str] = set() + backend_ids: set[str] = set() + for record in records: + implementations.add(str(record.get("implementation", ""))) + backend_ids.add(str(record.get("backend_id", ""))) + if record.get("case_id") != case_id: + errors.append(f"{module_label} record has the wrong case_id") + if record.get("implementation") != expected: + errors.append( + f"{module_label} implementation={record.get('implementation')!r}, " + f"expected {expected!r}" + ) + continue + execution_mode = record.get("execution_mode", "eager") + if framework == "vllm": + if execution_mode not in {"eager", "compiled_hip_graph"}: + errors.append(f"{module_label} has invalid HIP execution mode") + elif execution_mode != "eager": + errors.append(f"{module_label} did not execute in eager mode") + if expected == "production": + _validate_production_record(record, label=module_label, errors=errors) + elif module == "attention": + _validate_rlkernel_record( + record, + label=module_label, + framework=framework, + errors=errors, + ) + else: + _validate_strict_dense_record( + record, + module=module, + framework=framework, + label=module_label, + errors=errors, + ) + + module_summary[module] = { + "case_id": case_id, + "expected_implementation": expected, + "installed_processes": hook_count, + "call_count": call_count, + "implementations": sorted(implementations), + "backend_ids": sorted(backend_ids), + } + frameworks[label] = { + "readback_count": len(matching), + "modules": module_summary, + } + + return { + "passed": not errors, + "errors": errors, + "frameworks": frameworks, + } diff --git a/tests/test_vime_rocm_module_validation.py b/tests/test_vime_rocm_module_validation.py new file mode 100644 index 00000000..75a6d820 --- /dev/null +++ b/tests/test_vime_rocm_module_validation.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from examples.vime_rocm_attention_ablation.validate_module_artifacts import ( + RL_KERNEL_MISMATCH_SIDECAR_MARKER, + validate_module_readbacks, +) + + +def _production_record( + framework: str, + module: str, + *, + call_count: int = 1, +) -> dict: + return { + "case_id": "P/P", + "implementation": "production", + "backend_id": f"{framework}.production.{module}", + "call_count": call_count, + "provenance": { + "runtime_platform": "rocm", + "fallback": False, + }, + } + + +def _production_readbacks(*, rollout_ffn_calls: int = 1) -> list[dict]: + cases = {module: {"case_id": "P/P"} for module in ("attention", "ffn", "logp")} + return [ + { + "framework": "megatron", + "target": "training", + "plan": {"cases": cases}, + "installed_hooks": { + "attention": "test.attention", + "ffn": "test.ffn", + }, + "operators": { + module: _production_record("megatron", module) for module in ("attention", "ffn") + }, + "fallbacks": [], + }, + { + "framework": "vllm", + "target": "rollout", + "plan": {"cases": cases}, + "installed_hooks": { + module: f"test.{module}" for module in ("attention", "ffn", "logp") + }, + "operators": { + "attention": _production_record("vllm", "attention"), + "ffn": _production_record( + "vllm", + "ffn", + call_count=rollout_ffn_calls, + ), + "logp": _production_record("vllm", "logp"), + }, + "fallbacks": [], + }, + ] + + +def test_rocm_module_validator_accepts_native_production_routes(): + report = validate_module_readbacks( + _production_readbacks(), + { + "attention": "P/P", + "ffn": "P/P", + "logp": "P/P", + }, + log_text=f"{RL_KERNEL_MISMATCH_SIDECAR_MARKER}P/P", + ) + + assert report["passed"] + + +def test_rocm_module_validator_reports_graph_capture_observability_gap(): + report = validate_module_readbacks( + _production_readbacks(rollout_ffn_calls=0), + { + "attention": "P/P", + "ffn": "P/P", + "logp": "P/P", + }, + log_text=f"{RL_KERNEL_MISMATCH_SIDECAR_MARKER}P/P", + ) + + assert not report["passed"] + assert report["errors"] == ["vllm/rollout ffn had zero calls"] From 84a28696a9bc5d69b695459142c3f89319050671 Mon Sep 17 00:00:00 2001 From: Codex Date: Fri, 11 Sep 2026 19:13:21 +0800 Subject: [PATCH 6/8] fix(rocm): allow independent VIME module routes --- .../vime_rocm_attention_ablation/launch_arm.sh | 14 +++++++++----- tests/test_vime_rocm_attention_topology.py | 6 ++++-- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/examples/vime_rocm_attention_ablation/launch_arm.sh b/examples/vime_rocm_attention_ablation/launch_arm.sh index b04a2294..c2a845c0 100644 --- a/examples/vime_rocm_attention_ablation/launch_arm.sh +++ b/examples/vime_rocm_attention_ablation/launch_arm.sh @@ -53,11 +53,15 @@ if [[ "${RL_KERNEL_ATTENTION_CASE:-}" != "${RLK_ABLATION_CASE_ID}" ]]; then echo "RL_KERNEL_ATTENTION_CASE disagrees with the arm ID" >&2 exit 2 fi -if [[ "${RL_KERNEL_FFN_CASE:-}" != "${RL_KERNEL_LOGP_CASE:-}" ]] || - [[ "${RL_KERNEL_FFN_CASE:-}" != "R/R" && "${RL_KERNEL_FFN_CASE:-}" != "P/P" ]]; then - echo "FFN and Logp must both use P/P or both use R/R" >&2 - exit 2 -fi +for module_case in "${RL_KERNEL_FFN_CASE:-}" "${RL_KERNEL_LOGP_CASE:-}"; do + case "${module_case}" in + P/P|R/R) ;; + *) + echo "FFN and Logp cases must each be P/P or R/R" >&2 + exit 2 + ;; + esac +done if [[ "${RL_KERNEL_VLLM_INTEGRATION:-}" != "1" ]]; then echo "RL_KERNEL_VLLM_INTEGRATION=1 is required for rollout route readback" >&2 exit 2 diff --git a/tests/test_vime_rocm_attention_topology.py b/tests/test_vime_rocm_attention_topology.py index 0d600be2..c8171acb 100644 --- a/tests/test_vime_rocm_attention_topology.py +++ b/tests/test_vime_rocm_attention_topology.py @@ -124,5 +124,7 @@ def test_launcher_uses_pr377_torch_dist_actor_load_without_reference_model(): assert "--linear-logp-provider" in launcher assert "rl_engine.integrations.vime.linear_logp_provider.provider" in launcher assert "--linear-logp-provider-mode strict" in launcher - assert '"${RL_KERNEL_FFN_CASE:-}" != "R/R"' in launcher - assert '"${RL_KERNEL_LOGP_CASE:-}" != "R/R"' in launcher + assert ( + 'for module_case in "${RL_KERNEL_FFN_CASE:-}" ' '"${RL_KERNEL_LOGP_CASE:-}"; do' + ) in launcher + assert "FFN and Logp cases must each be P/P or R/R" in launcher From a588f114c2ac7878246eb45f65db820d20cddede Mon Sep 17 00:00:00 2001 From: Codex Date: Fri, 11 Sep 2026 19:14:30 +0800 Subject: [PATCH 7/8] feat(rocm): support KL module ablations --- examples/vime_rocm_attention_ablation/launch_arm.sh | 9 +++++++++ tests/test_vime_rocm_attention_topology.py | 5 +++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/examples/vime_rocm_attention_ablation/launch_arm.sh b/examples/vime_rocm_attention_ablation/launch_arm.sh index c2a845c0..21676461 100644 --- a/examples/vime_rocm_attention_ablation/launch_arm.sh +++ b/examples/vime_rocm_attention_ablation/launch_arm.sh @@ -282,6 +282,14 @@ ROLLOUT_LOGPROBS_ARGS=() if [[ "${RLK_ABLATION_USE_ROLLOUT_LOGPROBS:-0}" == "1" ]]; then ROLLOUT_LOGPROBS_ARGS+=(--use-rollout-logprobs) fi +REFERENCE_MODEL_ARGS=() +if [[ "${RLK_ABLATION_USE_KL_LOSS:-0}" == "1" ]]; then + : "${RLK_ABLATION_KL_LOSS_COEF:?}" + REFERENCE_MODEL_ARGS+=( + --use-kl-loss + --kl-loss-coef "${RLK_ABLATION_KL_LOSS_COEF}" + ) +fi ray job submit \ --address="${ray_job_address}" \ --runtime-env-json="${RUNTIME_ENV_JSON}" \ @@ -297,6 +305,7 @@ ray job submit \ "${MODEL_ARGS[@]}" \ --hf-checkpoint "${RLK_ABLATION_MODEL_ROOT}" \ --ref-load "${RLK_ABLATION_REFERENCE_CHECKPOINT}" \ + "${REFERENCE_MODEL_ARGS[@]}" \ --load "${RLK_ABLATION_REFERENCE_CHECKPOINT}" \ --start-rollout-id 0 \ "${SAVE_ARGS[@]}" \ diff --git a/tests/test_vime_rocm_attention_topology.py b/tests/test_vime_rocm_attention_topology.py index c8171acb..6757e9a1 100644 --- a/tests/test_vime_rocm_attention_topology.py +++ b/tests/test_vime_rocm_attention_topology.py @@ -119,8 +119,9 @@ def test_launcher_uses_pr377_torch_dist_actor_load_without_reference_model(): assert '--load "${RLK_ABLATION_REFERENCE_CHECKPOINT}"' in launcher assert "--megatron-to-hf-mode" not in launcher - assert "--use-kl-loss" not in launcher - assert "--kl-loss-coef" not in launcher + assert 'RLK_ABLATION_USE_KL_LOSS:-0}" == "1"' in launcher + assert "--use-kl-loss" in launcher + assert "--kl-loss-coef" in launcher assert "--linear-logp-provider" in launcher assert "rl_engine.integrations.vime.linear_logp_provider.provider" in launcher assert "--linear-logp-provider-mode strict" in launcher From b23500ba35af3ca9803652ef72ac20ed4595f513 Mon Sep 17 00:00:00 2001 From: Codex Date: Fri, 11 Sep 2026 23:21:43 +0800 Subject: [PATCH 8/8] feat(vime): derive rollout topology from TP and CP --- examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py | 82 ++++++++++++++++--- .../run_supplement_suite.py | 6 ++ .../vime_qwen3_8b_tp4_cp2_200/validate_run.py | 59 +++++++++++-- tests/test_vime_tp4_example.py | 38 +++++++++ 4 files changed, 166 insertions(+), 19 deletions(-) diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py b/examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py index b189d8d2..34648710 100755 --- a/examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py @@ -62,10 +62,36 @@ class Arm: "colocate": True, "offload_train": False, "offload_rollout": True, + "rollout_tp": 4, + "rollout_cp": 1, "rollout_gpus_per_engine": 4, "rollout_engines": 2, } + +def _rollout_topology( + rollout_tp_size: int, + rollout_cp_size: int, +) -> dict[str, int | bool]: + if rollout_tp_size <= 0: + raise ValueError("--rollout-tp-size must be positive") + if rollout_cp_size <= 0: + raise ValueError("--rollout-cp-size must be positive") + rollout_gpus = int(TOPOLOGY["rollout_gpus"]) + gpus_per_engine = rollout_tp_size * rollout_cp_size + if rollout_gpus % gpus_per_engine: + raise ValueError( + "--rollout-tp-size * --rollout-cp-size must divide the configured " + f"rollout GPU count ({gpus_per_engine} does not divide {rollout_gpus})" + ) + topology = dict(TOPOLOGY) + topology["rollout_tp"] = rollout_tp_size + topology["rollout_cp"] = rollout_cp_size + topology["rollout_gpus_per_engine"] = gpus_per_engine + topology["rollout_engines"] = rollout_gpus // gpus_per_engine + return topology + + # CP2/P2P requires Transformer Engine's fused attention on this topology. # Pinning the choice keeps the production arms independent of host-specific # backend auto-selection. @@ -104,7 +130,10 @@ class Arm: def _max_engine_decode_batch( - rollout_batch_size: int, n_samples_per_prompt: int, router_policy: str + rollout_batch_size: int, + n_samples_per_prompt: int, + router_policy: str, + rollout_engines: int, ) -> int: """Largest decode batch one vLLM engine can hold under the active router. @@ -114,9 +143,8 @@ def _max_engine_decode_batch( """ concurrency = rollout_batch_size * n_samples_per_prompt - engines = TOPOLOGY["rollout_gpus"] // TOPOLOGY["rollout_gpus_per_engine"] - if router_policy == "round_robin" and engines > 1: - return -(-concurrency // engines) # ceil + if router_policy == "round_robin" and rollout_engines > 1: + return -(-concurrency // rollout_engines) # ceil return concurrency @@ -240,6 +268,26 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--rollout-batch-size", type=int, default=1) parser.add_argument("--n-samples-per-prompt", type=int, default=8) parser.add_argument("--global-batch-size", type=int, default=8) + parser.add_argument( + "--rollout-tp-size", + type=int, + default=4, + help=( + "vLLM rollout tensor-parallel size. The runner derives " + "--rollout-num-gpus-per-engine, router engine count, and per-engine " + "CUDA Graph capture sizes from this value." + ), + ) + parser.add_argument( + "--rollout-cp-size", + type=int, + default=1, + help=( + "vLLM rollout prefill context-parallel size. The runner combines " + "this with --rollout-tp-size when deriving GPUs per engine and " + "router engine count." + ), + ) parser.add_argument( "--use-kl-loss", action="store_true", @@ -295,6 +343,10 @@ def main(argv: list[str] | None = None) -> int: "coefficient requires --use-kl-loss" ) arm = ARMS[args.group] + topology = _rollout_topology( + args.rollout_tp_size, + args.rollout_cp_size, + ) script_dir = Path(__file__).resolve().parent rl_kernel_root = _path(args.rl_kernel_root, "RL-Kernel root") @@ -338,7 +390,10 @@ def main(argv: list[str] | None = None) -> int: pythonpath.extend(item for item in os.environ["PYTHONPATH"].split(os.pathsep) if item) max_engine_decode_batch = _max_engine_decode_batch( - args.rollout_batch_size, args.n_samples_per_prompt, args.router_policy + args.rollout_batch_size, + args.n_samples_per_prompt, + args.router_policy, + int(topology["rollout_engines"]), ) env_vars = { "RL_KERNEL_ROOT": str(rl_kernel_root), @@ -380,9 +435,9 @@ def main(argv: list[str] | None = None) -> int: "--actor-num-nodes", "1", "--actor-num-gpus-per-node", - str(TOPOLOGY["actor_gpus"]), + str(topology["actor_gpus"]), "--rollout-num-gpus", - str(TOPOLOGY["rollout_gpus"]), + str(topology["rollout_gpus"]), "--colocate", "--no-offload-train", "--offload-rollout", @@ -421,9 +476,9 @@ def main(argv: list[str] | None = None) -> int: str(args.global_batch_size), "--balance-data", "--tensor-model-parallel-size", - str(TOPOLOGY["tp"]), + str(topology["tp"]), "--context-parallel-size", - str(TOPOLOGY["cp"]), + str(topology["cp"]), "--cp-comm-type", "p2p", "--pipeline-model-parallel-size", @@ -468,7 +523,9 @@ def main(argv: list[str] | None = None) -> int: "--router-policy", args.router_policy, "--rollout-num-gpus-per-engine", - str(TOPOLOGY["rollout_gpus_per_engine"]), + str(topology["rollout_gpus_per_engine"]), + "--vllm-prefill-context-parallel-size", + str(topology["rollout_cp"]), "--vllm-gpu-memory-utilization", str(args.vllm_gpu_memory_utilization), *_mismatch_metrics_args(), @@ -517,7 +574,7 @@ def main(argv: list[str] | None = None) -> int: "num_rollout": args.num_rollout, "seed": args.seed, "rollout_seed": args.rollout_seed, - "topology": dict(TOPOLOGY), + "topology": topology, "batching": { "rollout_batch_size": args.rollout_batch_size, "n_samples_per_prompt": args.n_samples_per_prompt, @@ -541,6 +598,9 @@ def main(argv: list[str] | None = None) -> int: }, "rollout_routing": { "router_policy": args.router_policy, + "engine_count": topology["rollout_engines"], + "tensor_parallel_size": topology["rollout_tp"], + "prefill_context_parallel_size": topology["rollout_cp"], }, "training_memory": { "recompute_granularity": "full", diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/run_supplement_suite.py b/examples/vime_qwen3_8b_tp4_cp2_200/run_supplement_suite.py index bc18cbbe..1cdea68b 100644 --- a/examples/vime_qwen3_8b_tp4_cp2_200/run_supplement_suite.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/run_supplement_suite.py @@ -36,6 +36,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--ray-bin", type=Path, required=True) parser.add_argument("--extra-pythonpath", action="append", default=[]) parser.add_argument("--ld-library-path", required=True) + parser.add_argument("--rollout-tp-size", type=int, default=4) + parser.add_argument("--rollout-cp-size", type=int, default=1) parser.add_argument("--idle-memory-mib", type=int, default=1024) parser.add_argument("--idle-poll-seconds", type=int, default=60) return parser.parse_args() @@ -152,6 +154,10 @@ def run_one(args: argparse.Namespace, group: str, rounds: int, seed: int) -> Non "16", "--global-batch-size", "128", + "--rollout-tp-size", + str(args.rollout_tp_size), + "--rollout-cp-size", + str(args.rollout_cp_size), "--max-response-len", "7168", "--max-tokens-per-gpu", diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py b/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py index 2e89356a..44aabf84 100755 --- a/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py @@ -21,7 +21,7 @@ RECORD_RE = re.compile(r"\b(rollout|step|perf)\s+(\d+):\s+(\{.*\})\s*$") FRAMEWORKS = (("megatron", "training"), ("vllm", "rollout")) MODULES = ("attention", "ffn", "logp") -EXPECTED_TOPOLOGY = { +EXPECTED_FIXED_TOPOLOGY = { "gpus": 8, "actor_gpus": 8, "rollout_gpus": 8, @@ -31,8 +31,6 @@ "colocate": True, "offload_train": False, "offload_rollout": True, - "rollout_gpus_per_engine": 4, - "rollout_engines": 2, } CASE_FIELDS = { "attention": "attention_case", @@ -52,6 +50,40 @@ ) +def _validate_topology(value: Any) -> list[str]: + if not isinstance(value, Mapping): + return ["manifest does not contain the required TP4/CP2 colocated topology"] + errors = [ + f"manifest topology {key}={value.get(key)!r}, expected {expected!r}" + for key, expected in EXPECTED_FIXED_TOPOLOGY.items() + if value.get(key) != expected + ] + rollout_gpus = value.get("rollout_gpus") + rollout_gpus_per_engine = value.get("rollout_gpus_per_engine") + rollout_cp = value.get("rollout_cp", 1) + rollout_tp = value.get("rollout_tp", rollout_gpus_per_engine) + rollout_engines = value.get("rollout_engines") + if ( + not isinstance(rollout_gpus, int) + or not isinstance(rollout_tp, int) + or not isinstance(rollout_cp, int) + or not isinstance(rollout_gpus_per_engine, int) + or rollout_tp <= 0 + or rollout_cp <= 0 + or rollout_gpus_per_engine != rollout_tp * rollout_cp + or rollout_gpus % rollout_gpus_per_engine + ): + errors.append( + "manifest rollout_gpus_per_engine must equal rollout_tp * rollout_cp " + "and divide rollout_gpus" + ) + elif rollout_engines != rollout_gpus // rollout_gpus_per_engine: + errors.append( + "manifest rollout_engines does not match " "rollout_gpus // rollout_gpus_per_engine" + ) + return errors + + def _load_json(path: Path) -> dict[str, Any]: value = json.loads(path.read_text(encoding="utf-8")) if not isinstance(value, dict): @@ -399,15 +431,26 @@ def validate_run(run_dir: Path) -> dict[str, Any]: for index in range(max(0, len(train_command) - 1)) ): global_errors.append("train command does not explicitly select GRPO") - if manifest.get("topology") != EXPECTED_TOPOLOGY: - global_errors.append("manifest does not contain the required TP4/CP2 colocated topology") - required_command_pairs = ( + topology = manifest.get("topology") + global_errors.extend(_validate_topology(topology)) + topology = topology if isinstance(topology, Mapping) else {} + required_command_pairs = [ ("--actor-num-gpus-per-node", "8"), ("--rollout-num-gpus", "8"), ("--tensor-model-parallel-size", "4"), ("--context-parallel-size", "2"), - ("--rollout-num-gpus-per-engine", "4"), - ) + ( + "--rollout-num-gpus-per-engine", + str(topology.get("rollout_gpus_per_engine", "")), + ), + ] + if "rollout_cp" in topology: + required_command_pairs.append( + ( + "--vllm-prefill-context-parallel-size", + str(topology["rollout_cp"]), + ) + ) if isinstance(train_command, list): for flag, value in required_command_pairs: if not any( diff --git a/tests/test_vime_tp4_example.py b/tests/test_vime_tp4_example.py index bf8a205d..5a18b665 100644 --- a/tests/test_vime_tp4_example.py +++ b/tests/test_vime_tp4_example.py @@ -11,13 +11,16 @@ MEGATRON_ATTENTION_BACKEND, RL_KERNEL_LINEAR_LOGP_PROVIDER, _linear_logp_provider_args, + _max_engine_decode_batch, _mismatch_metrics_args, + _rollout_topology, ) from examples.vime_qwen3_8b_tp4_cp2_200.run_supplement_suite import specs from examples.vime_qwen3_8b_tp4_cp2_200.validate_run import ( RL_KERNEL_MISMATCH_SIDECAR_MARKER, VIME_NATIVE_LINEAR_LOGP_MARKER, _validate_readbacks, + _validate_topology, ) @@ -64,6 +67,41 @@ def test_tp4_formal_matrix_pins_the_vime_qwen3_attention_backend(): assert MEGATRON_ATTENTION_BACKEND == "fused" +def test_rollout_tp_cp_derive_router_engines_and_graph_batch(): + expected = { + (2, 1): (2, 4, 2), + (2, 2): (4, 2, 4), + (4, 1): (4, 2, 4), + (4, 2): (8, 1, 8), + (8, 1): (8, 1, 8), + } + for (rollout_tp, rollout_cp), ( + gpus_per_engine, + engines, + graph_batch, + ) in expected.items(): + topology = _rollout_topology(rollout_tp, rollout_cp) + assert topology["rollout_tp"] == rollout_tp + assert topology["rollout_cp"] == rollout_cp + assert topology["rollout_gpus_per_engine"] == gpus_per_engine + assert topology["rollout_engines"] == engines + assert ( + _max_engine_decode_batch( + 1, + 8, + "round_robin", + int(topology["rollout_engines"]), + ) + == graph_batch + ) + assert _validate_topology(topology) == [] + + legacy_topology = _rollout_topology(4, 1) + legacy_topology.pop("rollout_tp") + legacy_topology.pop("rollout_cp") + assert _validate_topology(legacy_topology) == [] + + def test_launcher_forces_cuda_graph_without_a_logp_provider(): root = Path(__file__).parents[1] launcher = root / "examples" / "vime_qwen3_8b_tp4_cp2_200" / "aligned_python_entrypoint.sh"