[feat] Init true on policy with qwen_dense

megatron/core/distributed/distributed_data_parallel.py

@@ -11,6 +11,7 @@
 from ..process_groups_config import ProcessGroupCollection
 from ..transformer.cuda_graphs import is_graph_capturing
 from ..transformer.transformer_config import TransformerConfig
+from miles_megatron_plugins.true_on_policy.contracts import resolve_true_on_policy_runtime_policy
 from ..utils import log_single_rank
 from .data_parallel_base import _BaseDataParallel
 from .distributed_data_parallel_config import DistributedDataParallelConfig
@@ -19,6 +20,29 @@
 logger = logging.getLogger(__name__)
 
 
+def _first_cp_comm_type(cp_comm_type):
+    if isinstance(cp_comm_type, list):
+        return cp_comm_type[0] if cp_comm_type else None
+    return cp_comm_type
+
+
+def _use_true_on_policy_ulysses_cp_gradient_scaling(config: TransformerConfig) -> bool:
+    """Return true when CP loss scaling is deferred to the CP gradient sum."""
+
+    return (
+        resolve_true_on_policy_runtime_policy(config).defer_ulysses_cp_loss_scaling_to_grad_sum
+        and getattr(config, "context_parallel_size", 1) > 1
+        and _first_cp_comm_type(getattr(config, "cp_comm_type", None)) == "a2a"
+        and not getattr(config, "calculate_per_token_loss", False)
+    )
+
+
+def _dense_gradient_scaling_factor(config: TransformerConfig, dp_cp_world_size: int) -> float:
+    if _use_true_on_policy_ulysses_cp_gradient_scaling(config):
+        return float(getattr(config, "context_parallel_size", 1)) / float(dp_cp_world_size)
+    return 1.0 / float(dp_cp_world_size)
+
+
 class DistributedDataParallel(_BaseDataParallel):
     """
     DDP wrapper which stores grads in contiguous buffers. Also has option of overlapping
@@ -127,7 +151,10 @@ def __init__(
                 expert_parallel_params.append(param)
 
         def _allocate_buffers_for_parameters(
-            input_params, data_parallel_group, gradient_scaling_factor
+            input_params,
+            data_parallel_group,
+            gradient_scaling_factor,
+            target_gradient_scaling_factor,
         ):
             param_and_grad_dtype_to_params = {}
             param_and_grad_dtype_to_offsets = {}
@@ -174,7 +201,6 @@ def _allocate_buffers_for_parameters(
                 param_and_grad_dtype_to_indices[(param_dtype, grad_dtype)] = indices
 
             if not config.calculate_per_token_loss:
-                target_gradient_scaling_factor = 1.0 / self.dp_cp_group.size()
                 if self.ddp_config.average_in_collective:
                     if self.ddp_config.num_distributed_optimizer_instances == 1:
                         # Collective is averaging gradients in collective with data_parallel_group.
@@ -288,12 +314,19 @@ def _allocate_buffers_for_parameters(
             else:
                 data_parallel_world_size = self.dp_cp_group.size()
 
-                gradient_scaling_factor = 1.0 / data_parallel_world_size
+                gradient_scaling_factor = _dense_gradient_scaling_factor(
+                    config, data_parallel_world_size
+                )
                 expert_gradient_scaling_factor = 1.0 / data_parallel_world_size
 
         # Allocate the param+grad buffers for dense params' grads.
         self.buffers, self.bucket_groups = _allocate_buffers_for_parameters(
-            dense_params, self.intra_dp_cp_group, gradient_scaling_factor=gradient_scaling_factor
+            dense_params,
+            self.intra_dp_cp_group,
+            gradient_scaling_factor=gradient_scaling_factor,
+            target_gradient_scaling_factor=_dense_gradient_scaling_factor(
+                config, self.dp_cp_group.size()
+            ),
         )
 
         # Allocate separate param+grad buffers for expert parallel params' grads.
@@ -302,6 +335,7 @@ def _allocate_buffers_for_parameters(
                 expert_parallel_params,
                 self.intra_expt_dp_group,
                 gradient_scaling_factor=expert_gradient_scaling_factor,
+                target_gradient_scaling_factor=1.0 / self.dp_cp_group.size(),
             )
         )
 

megatron/core/distributed/param_and_grad_buffer.py

@@ -3,6 +3,7 @@
 import functools
 import logging
 import math
+import os
 import warnings
 from contextlib import nullcontext
 from enum import Enum
@@ -17,6 +18,7 @@
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.rerun_state_machine import get_rerun_state_machine
 from megatron.core.utils import log_single_rank
+from miles_megatron_plugins.true_on_policy.debug import dump_param_and_grad_buffer_debug
 
 from ..fp8_utils import (
     is_float8tensor,
@@ -93,6 +95,7 @@ def __init__(
         gradient_scaling_factor: float,
         bucket_id: int,
         param_index_map: Dict[torch.nn.Parameter, tuple],
+        param_to_name: Dict[torch.nn.Parameter, str],
     ):
         self.params_list = params
         self.params = set(params)
@@ -106,6 +109,7 @@ def __init__(
         self.numel_unpadded = numel_unpadded
         self.gradient_scaling_factor = gradient_scaling_factor
         self.bucket_id = bucket_id
+        self.param_to_name = param_to_name
         # Derive bucket-local param offsets from the global param_index_map.
         self.param_to_index = {}
         for param in params:
@@ -199,6 +203,7 @@ def __init__(
         # or bucket.grad_data.
         self.cached_param_buffer_shard_list = [None] * len(self.buckets)
         self.cached_grad_buffer_shard_list = [None] * len(self.buckets)
+        self._grad_debug_dumped_buckets = set()
 
     def reset(self):
         """
@@ -218,8 +223,16 @@ def check_grads(self, check_for_nan_or_inf, check_for_large):
         all-reduce / reduce-scatter.
         """
         rerun_state_machine = get_rerun_state_machine()
+        grad_debug_dir = os.environ.get("MILES_GRAD_DEBUG_DIR")
         for i in range(len(self.buckets)):
             grad_norm = self.buckets[i].grad_data.norm(p=2)
+            if grad_debug_dir and not torch.isfinite(grad_norm).item():
+                dump_param_and_grad_buffer_debug(
+                    self,
+                    bucket_index=i,
+                    grad_norm=grad_norm,
+                    grad_debug_dir=grad_debug_dir,
+                )
             # check for NaN, Inf and unexpectedly large grads
             if check_for_nan_or_inf:
                 rerun_state_machine.validate_result(
@@ -616,6 +629,7 @@ def __init__(
         self.ddp_config = ddp_config
         self.params = params
         self.param_indices = param_indices
+        self.param_to_name = param_to_name
 
         # Check that params are unique.
         unique_params = set()
@@ -974,6 +988,7 @@ def _new_bucket(
             gradient_scaling_factor=self.gradient_scaling_factor,
             bucket_id=bucket_id,
             param_index_map=self.param_index_map,
+            param_to_name=self.param_to_name,
         )
         for bucket_param in bucket_params:
             assert bucket_param not in self.param_to_bucket

megatron/core/extensions/sglang.py

@@ -0,0 +1,6 @@
+"""Compatibility imports for the SGLang-compatible true-on-policy backend.
+
+New code should import from :mod:`miles_megatron_plugins.true_on_policy.sglang_backend`.
+"""
+
+from miles_megatron_plugins.true_on_policy.sglang_backend import *  # noqa: F403

megatron/core/models/common/embeddings/rope_utils.py

@@ -119,11 +119,12 @@ def _apply_rotary_pos_emb_bshd(
 
     # first part is cosine component
     # second part is sine component, need to change signs with _rotate_half method
-    cos_ = (torch.cos(freqs) * mscale).to(t.dtype)
-    sin_ = (torch.sin(freqs) * mscale).to(t.dtype)
+    orig_dtype = t.dtype
+    cos_ = (torch.cos(freqs) * mscale).float()
+    sin_ = (torch.sin(freqs) * mscale).float()
+    t = (t.float() * cos_) + (_rotate_half(t, rotary_interleaved).float() * sin_)
 
-    t = (t * cos_) + (_rotate_half(t, rotary_interleaved) * sin_)
-    return torch.cat((t, t_pass), dim=-1)
+    return torch.cat((t.to(orig_dtype), t_pass), dim=-1)
 
 
 def _get_thd_freqs_on_this_cp_rank(
@@ -183,6 +184,7 @@ def _apply_rotary_pos_emb_thd(
     multi_latent_attention: bool = False,
     mscale: float = 1.0,
     cp_group: torch.distributed.ProcessGroup = None,
+    ulysses_cp: bool = False,
 ) -> Tensor:
     """A baseline implementation of applying RoPE for `thd` format.
 
@@ -199,9 +201,23 @@ def _apply_rotary_pos_emb_thd(
 
     if cp_group is None:
         raise ValueError("cp_group must be provided for THD format RoPE")
-    cp_size = cp_group.size()
-    cp_rank = cp_group.rank()
-    seqlens = ((cu_seqlens[1:] - cu_seqlens[:-1]) // cp_size).tolist()
+    full_seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
+    full_token_count = int(full_seqlens.sum().item())
+    local_token_count = t.size(0)
+
+    # Ulysses CP can appear in two layouts before RoPE:
+    # - full-sequence layout, where each rank has all packed tokens and RoPE
+    #   should behave like CP size 1;
+    # - local zigzag sequence-shard layout, used by the SGLang attention path
+    #   before its all-to-all head redistribution. In that case RoPE must use
+    #   the real CP rank/size to select the correct positional slices.
+    if ulysses_cp and local_token_count == full_token_count:
+        cp_size = 1
+        cp_rank = 0
+    else:
+        cp_size = cp_group.size()
+        cp_rank = cp_group.rank()
+    seqlens = (full_seqlens // cp_size).tolist()
 
     # Handle two different frequency tensor formats:
     # 1. If freqs.size(0) == cu_seqlens[-1]: freqs contains all positions across all sequences
@@ -254,6 +270,7 @@ def apply_rotary_pos_emb(
     cu_seqlens: Optional[Tensor] = None,
     mscale: float = 1.0,
     cp_group: torch.distributed.ProcessGroup = None,
+    ulysses_cp: bool = False,
 ):
     """
     Reroute to the appropriate apply_rotary_pos_emb function depending on
@@ -287,8 +304,16 @@ def apply_rotary_pos_emb(
                 return fused_apply_rotary_pos_emb(t, freqs, interleaved=config.rotary_interleaved)
         else:
             assert fused_apply_rotary_pos_emb_thd is not None, "apply_rope_fusion is not available."
+            full_token_count = int((cu_seqlens[1:] - cu_seqlens[:-1]).sum().item())
+            local_token_count = t.size(0)
+            if ulysses_cp and local_token_count == full_token_count:
+                cp_size = 1
+                cp_rank = 0
+            else:
+                cp_size = cp_group.size()
+                cp_rank = cp_group.rank()
             return fused_apply_rotary_pos_emb_thd(
-                t, cu_seqlens, freqs, cp_size=cp_group.size(), cp_rank=cp_group.rank()
+                t, cu_seqlens, freqs, cp_size=cp_size, cp_rank=cp_rank
             )
     # use unfused implementation
     if cu_seqlens is None:
@@ -308,6 +333,7 @@ def apply_rotary_pos_emb(
             multi_latent_attention=config.multi_latent_attention,
             mscale=mscale,
             cp_group=cp_group,
+            ulysses_cp=ulysses_cp,
         )