From 708bee7493d5694eaea80bb59a7d33be658c95c0 Mon Sep 17 00:00:00 2001 From: maocheng23 Date: Tue, 26 May 2026 00:10:16 -0700 Subject: [PATCH] Add deterministic dense math helpers --- .../true_on_policy/matmul.py | 58 +++++++- miles_megatron_plugins/true_on_policy/norm.py | 132 +++++++++++++++++- 2 files changed, 180 insertions(+), 10 deletions(-) diff --git a/miles_megatron_plugins/true_on_policy/matmul.py b/miles_megatron_plugins/true_on_policy/matmul.py index c53a1fabfa4..12f9c4f95e1 100644 --- a/miles_megatron_plugins/true_on_policy/matmul.py +++ b/miles_megatron_plugins/true_on_policy/matmul.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from typing import Iterable, List, Optional import torch @@ -30,6 +31,29 @@ def _fixed_tree_sum_tensors(tensors: Iterable[torch.Tensor]) -> torch.Tensor: return partials[0] +def _sglang_tp_inv_matmul_forward( + input_2d: torch.Tensor, weight_t: torch.Tensor, bias: Optional[torch.Tensor] +) -> Optional[torch.Tensor]: + if input_2d.shape[-1] % _ROW_LINEAR_INV_BLOCK_K != 0: + return None + try: + import sglang.srt.tp_invariant_ops # noqa: F401 + except Exception: + return None + + try: + with torch.no_grad(): + return torch.ops.tp_inv_ops.matmul_tp_inv(input_2d, weight_t, bias) + except Exception: + return None + + +def _with_native_grad(exact: Optional[torch.Tensor], native: torch.Tensor) -> torch.Tensor: + if exact is None: + return native + return exact + (native - native.detach()) + + def _safe_group_size(group: Optional[torch.distributed.ProcessGroup]) -> int: if group is not None: return group.size() @@ -42,6 +66,15 @@ def _safe_group_size(group: Optional[torch.distributed.ProcessGroup]) -> int: def _safe_tensor_context_parallel_size() -> int: + rollout_tp_size = os.environ.get("MILES_TRUE_ON_POLICY_ROLLOUT_TP_SIZE") + if rollout_tp_size: + try: + parsed = int(rollout_tp_size) + except ValueError: + parsed = 0 + if parsed > 0: + return parsed + try: from megatron.core.parallel_state import get_tensor_and_context_parallel_world_size @@ -94,6 +127,8 @@ def _sglang_row_parallel_matmul( output = _fixed_tree_sum_tensors(partials).to(input_.dtype) if bias is not None: output = output + bias + exact = _sglang_tp_inv_matmul_forward(input_2d, weight_t, bias) + output = _with_native_grad(exact, output) return output.reshape(*input_shape[:-1], weight.shape[0]) @@ -120,7 +155,13 @@ def _sglang_rollout_partition_row_parallel_matmul( for start in range(0, input_2d.shape[1], rollout_partition_k): end = start + rollout_partition_k - partials.append(input_2d[:, start:end] @ weight_t[start:end, :]) + partials.append( + _sglang_row_parallel_matmul( + input_2d[:, start:end], + weight[:, start:end], + bias=None, + ) + ) output = _fixed_tree_sum_tensors(partials).to(input_.dtype) if bias is not None: @@ -164,12 +205,17 @@ def sglang_reference_matmul( if bias is not None and bias.dtype != weight.dtype: bias = bias.to(weight.dtype) - if _should_use_sglang_tp_invariant_row_linear(input_, row_parallel, tp_group): - return _sglang_row_parallel_matmul(input_, weight, bias) if row_parallel: - return _sglang_rollout_partition_row_parallel_matmul( - input_, weight, bias, tp_group=tp_group - ) + rollout_partition_k = _rollout_row_parallel_partition_k(input_, tp_group) + if ( + 0 < rollout_partition_k < input_.shape[-1] + and input_.shape[-1] % rollout_partition_k == 0 + ): + return _sglang_rollout_partition_row_parallel_matmul( + input_, weight, bias, tp_group=tp_group + ) + if _should_use_sglang_tp_invariant_row_linear(input_, row_parallel, tp_group): + return _sglang_row_parallel_matmul(input_, weight, bias) if weight.requires_grad: return linear_with_grad_accumulation_and_async_allreduce( diff --git a/miles_megatron_plugins/true_on_policy/norm.py b/miles_megatron_plugins/true_on_policy/norm.py index 0ce2da3514e..cd746154545 100644 --- a/miles_megatron_plugins/true_on_policy/norm.py +++ b/miles_megatron_plugins/true_on_policy/norm.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from typing import Optional import torch @@ -9,6 +10,75 @@ from .contracts import resolve_true_on_policy_runtime_policy +_FUSED_RMSNORM_ENV = "SGLANG_TRUE_ON_POLICY_FUSED_RMSNORM" + + +def _env_flag_enabled(name: str) -> bool: + return os.environ.get(name, "").lower() in {"1", "true", "yes", "on"} + + +def _should_use_fused_rms_norm(*tensors: Optional[torch.Tensor]) -> bool: + return _env_flag_enabled(_FUSED_RMSNORM_ENV) and any( + tensor is not None and tensor.is_cuda for tensor in tensors + ) + + +def _needs_native_grad(*tensors: Optional[torch.Tensor]) -> bool: + return torch.is_grad_enabled() and any( + tensor is not None and tensor.requires_grad for tensor in tensors + ) + + +def _with_native_grad( + exact: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + native: torch.Tensor | tuple[torch.Tensor, torch.Tensor], +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if isinstance(exact, tuple): + exact_output, exact_residual = exact + native_output, native_residual = native + return ( + exact_output.detach() + (native_output - native_output.detach()), + exact_residual.detach() + (native_residual - native_residual.detach()), + ) + return exact.detach() + (native - native.detach()) + + +def _maybe_true_on_policy_fused_rms_norm( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + *, + residual: Optional[torch.Tensor] = None, + post_residual_addition: Optional[torch.Tensor] = None, + cast_x_before_out_mul: bool = True, + norm_cast_dtype: torch.dtype, + weight_cast_dtype: torch.dtype, + output_dtype: Optional[torch.dtype] = None, + residual_dtype: Optional[torch.dtype] = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None: + if not _should_use_fused_rms_norm(x, weight, residual, post_residual_addition): + return None + + try: + from sglang.srt.batch_invariant_ops import true_on_policy_rms_norm + except Exception: + return None + + output = true_on_policy_rms_norm( + x, + weight, + eps, + residual=residual, + post_residual_addition=post_residual_addition, + cast_x_before_out_mul=cast_x_before_out_mul, + norm_cast_dtype=norm_cast_dtype, + weight_cast_dtype=weight_cast_dtype, + output_dtype=output_dtype, + residual_dtype=residual_dtype, + ) + return output + + class SGLangNorm(torch.nn.Module): """Norm wrapper with Megatron-compatible parameters and SGLang backend identity.""" @@ -84,13 +154,43 @@ def forward( weight = self.weight + 1 if self.zero_centered_gamma else self.weight return F.layer_norm(x, self.hidden_size, weight, self.bias, self.eps) + fused = None + if residual is None: + fused = _maybe_true_on_policy_fused_rms_norm( + x, + self.weight.float(), + self.eps, + residual=residual, + post_residual_addition=post_residual_addition, + cast_x_before_out_mul=self.cast_x_before_out_mul, + norm_cast_dtype=self.override_orig_dtype or x.dtype, + weight_cast_dtype=torch.float32, + residual_dtype=x.dtype, + ) + if fused is not None: + if not _needs_native_grad(x, self.weight, residual, post_residual_addition): + return fused + return _with_native_grad( + fused, + self._forward_rmsnorm_native(x, residual, post_residual_addition), + ) + + return self._forward_rmsnorm_native(x, residual, post_residual_addition) + + def _forward_rmsnorm_native( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + post_residual_addition: Optional[torch.Tensor] = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: orig_dtype = self.override_orig_dtype or x.dtype x_float = x.float() if residual is not None: x_float = x_float + residual.float() if post_residual_addition is not None: x_float = x_float + post_residual_addition.float() - residual = x_float.to(orig_dtype) + residual = x_float.to(x.dtype) + x_float = residual.float() output = x_float * torch.rsqrt(x_float.pow(2).mean(-1, keepdim=True) + self.eps) if self.cast_x_before_out_mul: @@ -133,6 +233,22 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: if not x.is_contiguous(): x = x.contiguous() + fused = _maybe_true_on_policy_fused_rms_norm( + x, + self.weight.float(), + self.eps, + cast_x_before_out_mul=self.cast_x_before_out_mul, + norm_cast_dtype=x.dtype, + weight_cast_dtype=torch.float32, + ) + if fused is not None: + if not _needs_native_grad(x, self.weight): + return fused + return _with_native_grad(fused, self._forward_native(x)) + + return self._forward_native(x) + + def _forward_native(self, x: torch.Tensor) -> torch.Tensor: orig_dtype = x.dtype x_float = x.to(torch.float32) x_float = x_float * torch.rsqrt(x_float.pow(2).mean(dim=-1, keepdim=True) + self.eps) @@ -158,13 +274,13 @@ def __init__( ) -> None: super().__init__() - del config del persist_layer_norm del zero_centered_gamma del normalization self.hidden_size = (hidden_size,) self.eps = eps + self.source_truth_orig_dtype = getattr(config, "params_dtype", None) self.weight = torch.nn.Parameter(torch.ones(hidden_size, dtype=torch.float32)) def forward( @@ -176,7 +292,15 @@ def forward( if not x.is_contiguous(): x = x.contiguous() - orig_dtype = x.dtype + return self._forward_native(x, residual, post_residual_addition) + + def _forward_native( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + post_residual_addition: Optional[torch.Tensor] = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + orig_dtype = self.source_truth_orig_dtype or x.dtype if residual is not None: x = x + residual if post_residual_addition is not None: @@ -185,7 +309,7 @@ def forward( x_float = x.to(torch.float32) x_float = x_float * torch.rsqrt(x_float.pow(2).mean(dim=-1, keepdim=True) + self.eps) - output = self.weight * x_float.to(orig_dtype) + output = self.weight.to(orig_dtype) * x_float.to(orig_dtype) if residual is not None: return output, residual