diff --git a/csrc/cuda/distributed/deterministic_collective.cu b/csrc/cuda/distributed/deterministic_collective.cu index 72d3f874..e7d836a3 100644 --- a/csrc/cuda/distributed/deterministic_collective.cu +++ b/csrc/cuda/distributed/deterministic_collective.cu @@ -949,7 +949,10 @@ class DeterministicCollectiveState { has_staged_input_ = true; } - void all_reduce(torch::Tensor& output, cudaStream_t stream) { + void all_reduce( + torch::Tensor& output, + cudaStream_t stream, + bool allow_owner_path = true) { check_tensor(output, "output"); TORCH_CHECK(has_staged_input_, "stage() must be called before all_reduce()"); TORCH_CHECK( @@ -960,7 +963,10 @@ class DeterministicCollectiveState { "all-reduce output size must match the staged input size"); const int64_t element_count = output.numel(); - if (staged_owner_path_) { + // Graph-replayed calls must avoid the owner-push branch because it writes + // to remote IPC frames. Callers that use the fused ABI pass false here; + // direct staged collectives retain the eager owner optimization. + if (allow_owner_path && staged_owner_path_) { if (rank_ == 0) { // Logical rank 0 is the topology-favorable reader in both traced TP // engines. It evaluates the original fixed tree exactly once, then @@ -1022,6 +1028,58 @@ class DeterministicCollectiveState { return; } + // Small collectives are launch-bound. The fast kernel folds the peer + // stage wait, fixed-tree reduction, and completion publication into one + // launch while preserving the exact reduction order. + if (staged_fast_path_ && staged_bytes_ <= kSingleBlockFastPathMaxBytes) { + if (element_count > 0) { + switch (output.scalar_type()) { + case at::ScalarType::Float: + launch_all_reduce_fast( + peers_, + local_stage_sequence_, + local_done_sequence_, + static_cast(output.data_ptr()), + element_count, + world_size_, + stream); + break; + case at::ScalarType::Half: + launch_all_reduce_fast( + peers_, + local_stage_sequence_, + local_done_sequence_, + static_cast(output.data_ptr()), + element_count, + world_size_, + stream); + break; +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case at::ScalarType::BFloat16: + launch_all_reduce_fast( + peers_, + local_stage_sequence_, + local_done_sequence_, + static_cast(output.data_ptr()), + element_count, + world_size_, + stream); + break; +#endif + default: + TORCH_CHECK( + false, + "deterministic all-reduce supports float32, float16, and bfloat16; got ", + output.scalar_type()); + } + AT_CUDA_CHECK(cudaGetLastError()); + } else { + publish_done(stream); + } + has_staged_input_ = false; + return; + } + wait_for_staged_peers(stream); if (element_count == 0) { publish_done(stream); @@ -1090,10 +1148,11 @@ class DeterministicCollectiveState { output.numel() == input.numel(), "all-reduce output size must match the input size"); - // Route both ABI entry points through the same staged protocol so the - // topology-aware owner reduction is always applied. + // Use the graph-safe staged protocol for every message size. The fused + // two-slot protocol is intentionally kept available in the extension for + // experiments, but is not safe to replay across vLLM's many graph shapes. stage(input, stream); - all_reduce(output, stream); + all_reduce(output, stream, /*allow_owner_path=*/false); } void all_gather_fused( diff --git a/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh b/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh index 4b3a601c..eb1cf61b 100755 --- a/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh +++ b/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh @@ -12,14 +12,37 @@ if [[ "${1:-}" == "train.py" || "${1:-}" == */train.py ]]; then # vLLM hooks; otherwise it silently falls back to native attention/FFN and # loses the R/R performance path. Preserve explicit ablation selections. strict_linear_logp=0 + rollout_batch_size="" + n_samples_per_prompt="1" + explicit_vllm_execution_config=0 previous_arg="" for current_arg in "$@"; do if [[ "${previous_arg}" == "--linear-logp-provider-mode" && "${current_arg}" == "strict" ]]; then strict_linear_logp=1 - break + elif [[ "${previous_arg}" == "--rollout-batch-size" ]]; then + rollout_batch_size="${current_arg}" + elif [[ "${previous_arg}" == "--n-samples-per-prompt" ]]; then + n_samples_per_prompt="${current_arg}" fi + + case "${current_arg}" in + --linear-logp-provider-mode=strict) + strict_linear_logp=1 + ;; + --rollout-batch-size=*) + rollout_batch_size="${current_arg#*=}" + ;; + --n-samples-per-prompt=*) + n_samples_per_prompt="${current_arg#*=}" + ;; + --vllm-enforce-eager|--vllm-optimization-level|--vllm-optimization-level=*|--vllm-compilation-config|--vllm-compilation-config=*) + explicit_vllm_execution_config=1 + ;; + esac previous_arg="${current_arg}" done + + strict_cudagraph_args=() if [[ "${strict_linear_logp}" == "1" ]]; then export RL_KERNEL_VLLM_INTEGRATION="${RL_KERNEL_VLLM_INTEGRATION:-1}" export RL_KERNEL_CUDA_ONLY="${RL_KERNEL_CUDA_ONLY:-1}" @@ -27,6 +50,42 @@ if [[ "${1:-}" == "train.py" || "${1:-}" == */train.py ]]; then export RL_KERNEL_ATTENTION_CASE="${RL_KERNEL_ATTENTION_CASE:-R/R}" export RL_KERNEL_FFN_CASE="${RL_KERNEL_FFN_CASE:-R/R}" export RL_KERNEL_LOGP_CASE="${RL_KERNEL_LOGP_CASE:-R/R}" + + # Strict rollout kernels preserve their arithmetic order under CUDA Graph. + # Capturing the complete decode graph removes the per-layer host-launch + # gaps that otherwise dominate small decode batches. Capture every exact + # batch size: padding a strict custom kernel to a larger sparse graph can + # access invalid slots and, more importantly, changes the tested contract. + # Explicit vLLM execution flags always win so callers can opt out. + if [[ "${explicit_vllm_execution_config}" == "0" ]]; then + if [[ "${rollout_batch_size}" =~ ^[1-9][0-9]*$ && "${n_samples_per_prompt}" =~ ^[1-9][0-9]*$ ]]; then + max_capture_size=$((rollout_batch_size * n_samples_per_prompt)) + if [[ -n "${RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE:-}" ]]; then + max_capture_size="${RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE}" + fi + if ! [[ "${max_capture_size}" =~ ^[1-9][0-9]*$ ]]; then + echo "RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE must be a positive integer" >&2 + exit 2 + fi + + capture_sizes="[" + for ((batch_size = 1; batch_size <= max_capture_size; batch_size++)); do + if ((batch_size > 1)); then + capture_sizes+="," + fi + capture_sizes+="${batch_size}" + done + capture_sizes+="]" + compilation_config="{\"cudagraph_mode\":\"FULL_DECODE_ONLY\",\"cudagraph_capture_sizes\":${capture_sizes},\"max_cudagraph_capture_size\":${max_capture_size}}" + strict_cudagraph_args=( + --vllm-optimization-level 0 + --vllm-compilation-config "${compilation_config}" + ) + echo "[RL-Kernel] strict vLLM full-decode CUDA Graph capture sizes: ${capture_sizes}" >&2 + else + echo "[RL-Kernel] strict CUDA Graph disabled: rollout batch size is unavailable" >&2 + fi + fi fi exec "${REAL_PYTHON}" "$@" \ --seed 1234 \ @@ -35,7 +94,8 @@ if [[ "${1:-}" == "train.py" || "${1:-}" == */train.py ]]; then --vllm-attention-backend flash_attn \ --vllm-disable-custom-all-reduce \ --deterministic-mode \ - --accumulate-allreduce-grads-in-fp32 + --accumulate-allreduce-grads-in-fp32 \ + "${strict_cudagraph_args[@]}" fi exec "${REAL_PYTHON}" "$@" diff --git a/rl_engine/integrations/megatron_runtime.py b/rl_engine/integrations/megatron_runtime.py index e296e824..72ac8919 100644 --- a/rl_engine/integrations/megatron_runtime.py +++ b/rl_engine/integrations/megatron_runtime.py @@ -6,6 +6,7 @@ from __future__ import annotations import importlib +import os from collections.abc import Callable, Iterable from types import MethodType from typing import Any @@ -290,7 +291,48 @@ def initialize_from_environment(_args: Any = None) -> MegatronIntegration: from rl_engine.integrations.ablation import integration_plan_from_environment - return install_megatron_integration(integration_plan_from_environment()) + plan = integration_plan_from_environment() + integration = install_megatron_integration(plan) + if plan.implementation_for("attention", "training") is Implementation.RL_KERNEL: + _precompile_strict_attention_training(_args) + return integration + + +def _precompile_strict_attention_training(args: Any) -> None: + """Warm FA4 CuTe fwd/bwd JIT outside Vime's actor_train timer.""" + + if args is None or torch.version.hip is not None or not torch.cuda.is_available(): + return + if os.getenv("RL_KERNEL_PRECOMPILE_FA4", "1") == "0": + return + + from rl_engine.kernels.ops.cuda.attention.flash_attn import StrictFlashAttention4Core + + attention_heads = int(getattr(args, "num_attention_heads", 0) or 0) + query_groups = int(getattr(args, "num_query_groups", 0) or attention_heads) + tp_size = int(getattr(args, "tensor_model_parallel_size", 1) or 1) + head_dim = int( + getattr(args, "kv_channels", 0) + or (int(getattr(args, "hidden_size", 0) or 0) // attention_heads) + ) + if ( + attention_heads <= 0 + or query_groups <= 0 + or tp_size <= 0 + or attention_heads % tp_size + or query_groups % tp_size + or head_dim <= 0 + ): + return + + params_dtype = getattr(args, "params_dtype", None) + dtype = params_dtype if params_dtype in (torch.float16, torch.bfloat16) else torch.bfloat16 + StrictFlashAttention4Core.precompile_training( + q_heads=attention_heads // tp_size, + kv_heads=query_groups // tp_size, + head_dim=head_dim, + dtype=dtype, + ) __all__ = ["initialize_from_environment", "install_megatron_integration"] diff --git a/rl_engine/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py index c4da28be..ade13cab 100644 --- a/rl_engine/integrations/vllm_runtime.py +++ b/rl_engine/integrations/vllm_runtime.py @@ -12,7 +12,11 @@ import torch -from rl_engine.distributed.collectives import DETERMINISTIC_ALL_REDUCE_OP +from rl_engine.distributed.collectives import ( + DETERMINISTIC_ALL_REDUCE_OP, + collective_for_group, + deterministic_all_reduce_inplace, +) from rl_engine.integrations.ablation import ( Implementation, IntegrationPlan, @@ -40,6 +44,8 @@ _STRICT_RMS_NORM_INIT_MARKER = "__rl_kernel_original_strict_rms_norm_init__" _STRICT_ROTARY_INIT_MARKER = "__rl_kernel_original_strict_rotary_init__" _STRICT_LM_HEAD_LINEAR_PATCH_MARKER = "__rl_kernel_original_lm_head_linear_apply__" +_STRICT_O_PROJ_COLLECTIVE_MARKER = "__rl_kernel_o_proj_collective__" +_STRICT_ROW_PARALLEL_PATCH_MARKER = "__rl_kernel_original_row_parallel_forward__" _RLK_ATTENTION_BACKEND: type[Any] | None = None _RLK_ATTENTION_IMPL: type[Any] | None = None _RLK_ATTENTION_BUILDER: type[Any] | None = None @@ -367,6 +373,7 @@ def _patch_qwen3_strict_model( rotary_cls: type[Any] | None = None, linear_method_cls: type[Any] | None = None, attention_cls: type[Any] | None = None, + row_parallel_cls: type[Any] | None = None, det_gemm: Any | None = None, ) -> None: """Align vLLM's RMSNorm and Attention projections with Megatron.""" @@ -374,7 +381,7 @@ def _patch_qwen3_strict_model( production_classes = rms_norm_cls is None or linear_method_cls is None or attention_cls is None if production_classes: from vllm.model_executor.layers.layernorm import RMSNorm - from vllm.model_executor.layers.linear import UnquantizedLinearMethod + from vllm.model_executor.layers.linear import RowParallelLinear, UnquantizedLinearMethod from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding from vllm.model_executor.models.qwen3 import Qwen3Attention @@ -382,6 +389,7 @@ def _patch_qwen3_strict_model( rotary_cls = RotaryEmbedding linear_method_cls = UnquantizedLinearMethod attention_cls = Qwen3Attention + row_parallel_cls = RowParallelLinear assert rms_norm_cls is not None assert linear_method_cls is not None assert attention_cls is not None @@ -435,6 +443,62 @@ def strict_rms_norm_forward_cuda( eps=instance.variance_epsilon, ) + def bind_o_proj_collective(module: Any) -> None: + if int(getattr(module, "tp_size", 1)) <= 1: + return + from vllm.distributed.parallel_state import get_tp_group + + coordinator = get_tp_group() + group = getattr(coordinator, "device_group", coordinator) + collective = collective_for_group(group) + if collective is None: + raise RuntimeError("strict rollout o_proj requires an initialized TP process group") + setattr(module, _STRICT_O_PROJ_COLLECTIVE_MARKER, collective) + + if row_parallel_cls is not None and not hasattr( + row_parallel_cls, _STRICT_ROW_PARALLEL_PATCH_MARKER + ): + row_parallel_forward = row_parallel_cls.forward + + def strict_row_parallel_forward(instance: Any, input_: torch.Tensor) -> Any: + collective = getattr(instance, _STRICT_O_PROJ_COLLECTIVE_MARKER, None) + if collective is None: + return row_parallel_forward(instance, input_) + + if instance.input_is_parallel: + input_parallel = input_ + else: + from vllm.distributed import split_tensor_along_last_dim + + input_parallel = split_tensor_along_last_dim( + input_, num_partitions=instance.tp_size + )[instance.tp_rank].contiguous() + + assert instance.quant_method is not None + bias_ = ( + None + if (instance.tp_rank > 0 or instance.skip_bias_add) + else instance.bias + ) + output_parallel = instance.quant_method.apply(instance, input_parallel, bias_) + + if instance.reduce_results and instance.tp_size > 1: + deterministic_all_reduce_inplace( + output_parallel, + collective_handle=int(collective._handle), + ) + output = output_parallel + else: + output = output_parallel + + if not instance.return_bias: + return output + output_bias = instance.bias if instance.skip_bias_add else None + return output, output_bias + + setattr(row_parallel_cls, _STRICT_ROW_PARALLEL_PATCH_MARKER, row_parallel_forward) + row_parallel_cls.forward = strict_row_parallel_forward + if not hasattr(rms_norm_cls, _STRICT_RMS_NORM_INIT_MARKER): rms_norm_init = rms_norm_cls.__init__ @@ -463,6 +527,7 @@ def attention_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: attention_init(instance, *args, **kwargs) setattr(instance.qkv_proj, _STRICT_PROJECTION_MARKER, "qkv") setattr(instance.o_proj, _STRICT_PROJECTION_MARKER, "o_proj") + bind_o_proj_collective(instance.o_proj) setattr(attention_cls, _STRICT_MODEL_PATCH_MARKER, attention_init) linear_method_cls.apply = deterministic_linear_apply diff --git a/rl_engine/kernels/ops/cuda/attention/flash_attn.py b/rl_engine/kernels/ops/cuda/attention/flash_attn.py index cec8d2fe..9ad510b3 100644 --- a/rl_engine/kernels/ops/cuda/attention/flash_attn.py +++ b/rl_engine/kernels/ops/cuda/attention/flash_attn.py @@ -109,6 +109,82 @@ def __init__( self._op = op self._paged_op = paged_op + @classmethod + def precompile_training( + cls, + *, + q_heads: int, + kv_heads: int, + head_dim: int, + device: torch.device | None = None, + dtype: torch.dtype = torch.bfloat16, + sequence_length: int = 512, + ) -> None: + """Compile strict training FA4 forward/backward before timing. + + FA4 CuTe compiles its forward kernel on the first invocation and its + deterministic backward kernel on the first autograd backward. A + one-step RL workload would otherwise charge both compilations to + Vime's ``actor_train`` timer. These isolated tensors exercise the + same Qwen-style GQA and multi-block shape class without touching model + tensors, RNG state, or distributed collectives. + """ + if torch.version.hip is not None: + raise StrictFlashAttentionUnavailable( + "FA4 CUDA precompile is unavailable on ROCm" + ) + if not torch.cuda.is_available(): + raise StrictFlashAttentionUnavailable( + "FA4 CUDA precompile requires an available CUDA device" + ) + if dtype not in (torch.float16, torch.bfloat16): + raise ValueError("strict FA4 training precompile requires FP16 or BF16") + if q_heads <= 0 or kv_heads <= 0 or q_heads % kv_heads != 0: + raise ValueError("Q/KV head counts must be positive and GQA-compatible") + if head_dim <= 0 or sequence_length <= 0: + raise ValueError("head_dim and sequence_length must be positive") + + target = ( + torch.device("cuda", torch.cuda.current_device()) + if device is None + else device + ) + if target.type != "cuda": + raise ValueError("strict FA4 training precompile requires a CUDA device") + + core = cls() + # Zeros deliberately avoid consuming RNG state. The tensors are + # independent from the model and only populate FA4's process-local + # JIT caches. + q = torch.zeros( + (1, sequence_length, q_heads, head_dim), + dtype=dtype, + device=target, + requires_grad=True, + ) + k = torch.zeros( + (1, sequence_length, kv_heads, head_dim), + dtype=dtype, + device=target, + requires_grad=True, + ) + v = torch.zeros_like(k, requires_grad=True) + positions = torch.arange(sequence_length, dtype=torch.int64, device=target).expand(1, -1) + with torch.enable_grad(): + result = core.forward_bshd_with_lse( + q, + k, + v, + causal=True, + scale=head_dim**-0.5, + query_position_ids=positions, + key_position_ids=positions, + output_dtype=dtype, + ) + result.out.sum().backward() + torch.cuda.synchronize(target) + del result, q, k, v, positions, core + @staticmethod def _validate_api(op: Callable[..., Any]) -> None: try: