diff --git a/csrc/cuda/distributed/deterministic_collective.cu b/csrc/cuda/distributed/deterministic_collective.cu index 72d3f874..41cea442 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 @@ -1090,10 +1096,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/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