[GRL] Land PR400 sparse/delta megatron-lm dependency

megatron/core/extensions/transformer_engine.py

@@ -81,6 +81,41 @@ def condition_init_method(config, init_method):
     return init_method if config.perform_initialization else (lambda w: None)
 
 
+def _te_dot_product_attention_accepts_argument(argument_name: str) -> bool:
+    """Return whether the installed TE DPA forward accepts a Megatron extension kwarg."""
+    return argument_name in inspect.signature(te.pytorch.DotProductAttention.forward).parameters
+
+
+def _packed_seq_kwargs_for_te(
+    packed_seq_params: PackedSeqParams,
+    kept_packed_seq_params: set[str],
+    supports_tree_metadata: bool,
+) -> dict[str, Any]:
+    """Build the TransformerEngine packed-sequence kwargs Megatron owns.
+
+    ``PackedSeqParams.tree_metadata`` is a Megatron-side extension field. It
+    must not leak as a ``None`` kwarg into ordinary TE attention calls. When
+    populated, it is a real tree-training request and therefore requires a TE
+    build whose DPA explicitly accepts ``tree_metadata``.
+    """
+    packed_seq_kwargs = {
+        key: getattr(packed_seq_params, key)
+        for key in kept_packed_seq_params
+        if key != "tree_metadata"
+    }
+    if packed_seq_params.tree_metadata is None:
+        return packed_seq_kwargs
+    if not supports_tree_metadata:
+        raise RuntimeError(
+            "PackedSeqParams.tree_metadata requires TransformerEngine "
+            "DotProductAttention.forward(..., tree_metadata=...), but the installed "
+            f"TransformerEngine v{get_te_version()} does not expose that argument. "
+            "Use a tree-attention-capable TransformerEngine build or disable tree training."
+        )
+    packed_seq_kwargs["tree_metadata"] = packed_seq_params.tree_metadata
+    return packed_seq_kwargs
+
+
 def split_te_layernorm_column_parallel_linear(
     fused_layer,
     config,
@@ -1006,6 +1041,10 @@ def __init__(
         self.kept_packed_seq_params = set(
             field.name for field in dataclasses.fields(PackedSeqParams)
         )
+        self.kept_packed_seq_params.discard("tree_metadata")
+        self._te_supports_tree_metadata = _te_dot_product_attention_accepts_argument(
+            "tree_metadata"
+        )
 
         if get_te_version() < PkgVersion("1.3.0"):
             # TE 1.3.0 introduces precomputing max_seqlen to remove unnecessary kernels and D2H
@@ -1077,7 +1116,11 @@ def forward(
             self.kept_packed_seq_params.discard("cp_group")
             self.kept_packed_seq_params.discard("local_cp_size")
         packed_seq_kwargs = (
-            {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params}
+            _packed_seq_kwargs_for_te(
+                packed_seq_params,
+                self.kept_packed_seq_params,
+                self._te_supports_tree_metadata,
+            )
             if packed_seq_params is not None
             else {}
         )

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

@@ -177,14 +177,22 @@ def get_emb(self, max_seq_len: int, offset: int = 0) -> Tensor:
 
     @internal_api
     def forward(
-        self, max_seq_len: int, offset: int = 0, packed_seq_params: Optional[PackedSeqParams] = None
+        self,
+        max_seq_len: int,
+        offset: int = 0,
+        packed_seq_params: Optional[PackedSeqParams] = None,
+        position_ids: Optional[Tensor] = None,
     ) -> Tensor:
         """Forward pass of RoPE embedding.
 
         Args:
             max_seq_len (int): Maximum size of sequence
             offset (int, optional): RoPE offset. Defaults to 0.
             packed_seq_params (PackedSeqParams, optional): Packed sequence params. Defaults to None.
+            position_ids (Tensor, optional): Per-token positional ids for
+                reindexing the emb table. When provided, ``emb[position_ids]``
+                is returned so each packed token gets the RoPE frequency for
+                its logical position. Defaults to None (sequential positions).
 
         Returns:
             Tensor: Embeddings after applying RoPE.
@@ -202,6 +210,20 @@ def forward(
             # and select the parition of the current CP rank
             emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group)
 
+        if position_ids is not None:
+            pos_flat = position_ids.squeeze(0) if position_ids.dim() == 2 else position_ids
+            if pos_flat.device != emb.device:
+                pos_flat = pos_flat.to(emb.device)
+            if pos_flat.numel() > 0:
+                max_pos = int(pos_flat.max().item())
+                if max_pos >= emb.shape[0]:
+                    raise RuntimeError(
+                        f"RotaryEmbedding: position_id {max_pos} exceeds "
+                        f"emb table length {emb.shape[0]}; rotary_seq_len must "
+                        f"be >= max(position_ids) + 1."
+                    )
+            emb = emb[pos_flat]
+
         return emb
 
     def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):

megatron/core/models/gpt/gpt_model.py

@@ -344,7 +344,9 @@ def _preprocess(
                     inference_context, self.decoder, decoder_input, self.config, packed_seq_params
                 )
                 rotary_pos_emb = self.rotary_pos_emb(
-                    rotary_seq_len, packed_seq_params=packed_seq_params
+                    rotary_seq_len,
+                    packed_seq_params=packed_seq_params,
+                    position_ids=position_ids,
                 )
         elif self.position_embedding_type == 'yarn':
             if self.training or not self.config.flash_decode:

megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py

@@ -4,6 +4,18 @@
 
 import torch
 
+# SparseRL-Sync integration: sparse_diff_context wraps the param.copy_() so the
+# attached SparseManager can snapshot pre-state, then diff against post-state to
+# build per-param sparse-update indices. Falls back to nullcontext when the
+# sparse_update package is not installed so the upstream behavior is unchanged.
+try:
+    from sparse_update import sparse_diff_context
+except ImportError:
+    from contextlib import nullcontext
+
+    def sparse_diff_context(*args, **kwargs):
+        return nullcontext()
+
 
 def _param_generator(cpu_optimizer):
     for group in cpu_optimizer.param_groups:
@@ -121,7 +133,8 @@ def param_copy_back_gpu_hook(optimizer, args, kwargs):
                 with torch.cuda.stream(self._h2d_stream):
                     for param in _param_generator(optimizer):
                         gpu_param = self.cpu_copys_map_gpu_param[param]
-                        gpu_param.data.copy_(param.data, non_blocking=True)
+                        with sparse_diff_context(gpu_param, param):
+                            gpu_param.data.copy_(param.data, non_blocking=True)
                 self._d2h_stream.record_event().wait(torch.cuda.current_stream())
 
             return param_copy_back_gpu_hook
@@ -137,7 +150,8 @@ def fp32_param_copy_back_gpu_hook(optimizer, args, kwargs):
 
                         if param in self.param_to_fp32_param:
                             fp32_param = self.param_to_fp32_param[param]
-                            param.data.copy_(fp32_param.data)
+                            with sparse_diff_context(param, fp32_param):
+                                param.data.copy_(fp32_param.data)
 
             return fp32_param_copy_back_gpu_hook
 

megatron/core/optimizer/distrib_optimizer.py

@@ -53,6 +53,22 @@
 from .optimizer import MixedPrecisionOptimizer, _zero_grad_group_helper, param_group_identifier_keys
 from .optimizer_config import OptimizerConfig
 
+# SparseRL-Sync integration: init_sparse_manager binds the optimizer's shard
+# views to a SparseManager that tracks per-rollout dp-local diff indices.
+# sparse_diff_context wraps the in-place shard_model_param.copy_() that
+# realises the new training weights. Both fall back to no-ops when
+# sparse_update is not importable so the upstream behaviour is preserved.
+try:
+    from sparse_update import init_sparse_manager, sparse_diff_context
+except ImportError:
+    from contextlib import nullcontext
+
+    def sparse_diff_context(*args, **kwargs):
+        return nullcontext()
+
+    def init_sparse_manager(*args, **kwargs):
+        return None
+
 logger = getLogger(__name__)
 
 
@@ -413,6 +429,15 @@ def _build_model_and_main_param_groups(
                     shard_float16_params_this_group.append(shard_model_param)
                     shard_fp32_from_float16_params_this_group.append(shard_main_param)
 
+                    # SparseRL-Sync: bind the shard views to a SparseManager so
+                    # later sparse_diff_context() calls can recover the owner.
+                    init_sparse_manager(
+                        model_param=model_param,
+                        shard_model_weight=shard_model_param,
+                        shard_main_weight=shard_main_param,
+                        param_range=param_range,
+                    )
+
                 # fp32 params.
                 elif model_param.type() == 'torch.cuda.FloatTensor':
                     shard_model_param = model_param.view(-1)[param_range.start : param_range.end]
@@ -424,6 +449,15 @@ def _build_model_and_main_param_groups(
                     if hasattr(model_param, 'shared'):
                         shard_model_param.shared = model_param.shared
 
+                    # SparseRL-Sync: fp32 params share the same shard view for
+                    # both "model" and "main" sides; pass it on both slots.
+                    init_sparse_manager(
+                        model_param=model_param,
+                        shard_model_weight=shard_model_param,
+                        shard_main_weight=shard_model_param,
+                        param_range=param_range,
+                    )
+
                 else:
                     raise TypeError(
                         'Wrapped parameters must be one of '
@@ -2456,7 +2490,11 @@ def copy_group_params(shard_main_groups, model_groups):
                         # FP8 params are quantized in the above "quantize_param_shard" function.
                         continue
                     else:
-                        shard_model_param.data.copy_(shard_main_param)
+                        # SparseRL-Sync: capture the pre/post state of the
+                        # shard copy so the per-param diff indices can be
+                        # computed by the manager.
+                        with sparse_diff_context(shard_model_param, shard_main_param):
+                            shard_model_param.data.copy_(shard_main_param)
 
         # Copy shard groups to model groups.
         copy_group_params(self.shard_fp32_from_float16_groups, self.model_float16_groups)

megatron/core/packed_seq_params.py

@@ -1,9 +1,13 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
 
 import torch.distributed as dist
 from torch import Tensor
 
+if TYPE_CHECKING:
+    from megatron.core.transformer.tree_metadata import TreeMetadata
+
 
 @dataclass
 class PackedSeqParams:
@@ -21,3 +25,9 @@ class PackedSeqParams:
     max_seqlen_kv: int = None
     local_cp_size: int = None
     cp_group: dist.ProcessGroup = None
+    # Tree-attention metadata for the native tree-training path. ``None`` for
+    # all existing callers; populated by slime's TreeDataIterator when
+    # ``--enable-tree-training`` is on. TEDotProductAttention passes this
+    # through to TE, which dispatches to TreeFlashAttention.
+    # See megatron.core.transformer.tree_metadata.
+    tree_metadata: Optional["TreeMetadata"] = None

megatron/core/pipeline_parallel/p2p_communication.py

@@ -26,22 +26,22 @@ def _batched_p2p_ops(
     ops = []
     if tensor_send_prev is not None:
         send_prev_op = torch.distributed.P2POp(
-            torch.distributed.isend, tensor_send_prev, prev_pipeline_rank,
+            torch.distributed.isend, tensor_send_prev, prev_pipeline_rank, group
         )
         ops.append(send_prev_op)
     if tensor_recv_prev is not None:
         recv_prev_op = torch.distributed.P2POp(
-            torch.distributed.irecv, tensor_recv_prev, prev_pipeline_rank,
+            torch.distributed.irecv, tensor_recv_prev, prev_pipeline_rank, group
         )
         ops.append(recv_prev_op)
     if tensor_send_next is not None:
         send_next_op = torch.distributed.P2POp(
-            torch.distributed.isend, tensor_send_next, next_pipeline_rank,
+            torch.distributed.isend, tensor_send_next, next_pipeline_rank, group
         )
         ops.append(send_next_op)
     if tensor_recv_next is not None:
         recv_next_op = torch.distributed.P2POp(
-            torch.distributed.irecv, tensor_recv_next, next_pipeline_rank,
+            torch.distributed.irecv, tensor_recv_next, next_pipeline_rank, group
         )
         ops.append(recv_next_op)
     if len(ops) > 0:

megatron/core/transformer/enums.py

@@ -65,6 +65,10 @@ class AttnBackend(enum.Enum):
     unfused = 3
     local = 4
     auto = 5
+    # Tree-attention backend: data-driven dispatch via tree_metadata on
+    # PackedSeqParams. TEDotProductAttention detects tree_metadata and
+    # routes to TE's TreeFlashAttention backend. ``cp_size == 1`` required.
+    tree = 6
 
 
 class CudaGraphScope(enum.Enum):

megatron/core/transformer/moe/moe_layer.py

@@ -39,6 +39,7 @@
 class MoESubmodules:
     """MoE Layer Submodule spec"""
 
+    router: Union[ModuleSpec, type] = TopKRouter
     experts: Union[ModuleSpec, type] = None
     shared_experts: Union[ModuleSpec, type] = None
 
@@ -126,7 +127,11 @@ def __init__(
         )
 
         # Initialize router
-        self.router = TopKRouter(config=self.config, pg_collection=pg_collection)
+        self.router = build_module(
+            self.submodules.router,
+            config=self.config,
+            pg_collection=pg_collection,
+        )
         self.tp_group = pg_collection.tp
         # Initialize token dispatcher
         if config.moe_token_dispatcher_type == "allgather":

megatron/core/transformer/moe/moe_utils.py

@@ -587,9 +587,6 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None):
         else:
             return torch.topk(scores, k=topk, dim=1)
 
-    from sirl.utils.routing_replay import get_routing_replay_compute_topk
-    compute_topk = get_routing_replay_compute_topk(compute_topk)
-
     if score_function == "softmax":
         if use_pre_softmax:
             scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits)

megatron/core/transformer/moe/router.py

@@ -201,9 +201,6 @@ def __init__(
             self.global_tokens_per_expert = None
             self.ga_steps = None
 
-        from sirl.utils.routing_replay import register_routing_replay
-        register_routing_replay(self)
-
     def _maintain_float32_expert_bias(self):
         """
         Maintain the expert bias in float32.

megatron/core/transformer/tree_metadata.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+"""Tree-attention metadata carried alongside packed THD sequences.
+
+Used by the tree-training native path to flow tree topology + precomputed FA3
+attention buffers + per-token positional ids from the data iterator through
+``PackedSeqParams.tree_metadata`` to TE's ``TreeFlashAttention`` backend.
+The per-token positional ids are also copied to
+``PackedSeqParams.position_ids`` for ``RotaryEmbedding`` to consume.
+
+The dataclass is deliberately lightweight: torch is imported only for
+``Tensor`` typing, no CUDA assumptions are made at definition time.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+
+from torch import Tensor
+
+
+@dataclass
+class TreeMetadata:
+    """Tree topology + per-token positional ids for one packed micro-batch.
+
+    Fields
+    ------
+    cu_node_lens
+        ``int32`` tensor of shape ``[num_nodes + 1]`` — cumulative token counts
+        per trie node. Identical contract to the FA3 tree kernel input.
+    node_parent
+        ``int32`` tensor of shape ``[num_nodes]`` — parent index per trie node;
+        ``-1`` marks roots.
+    tree_position_ids
+        ``int64`` tensor of shape ``[total_tokens]`` — per-token positional id.
+        Passed as ``position_ids`` through GPTModel.forward so
+        ``RotaryEmbedding.forward`` can gather rope freqs.
+    padded_size
+        Padded packed length (``cu_node_lens[-1]`` after padding).
+    num_nodes
+        Number of real trie nodes (``len(node_parent)``); convenience accessor
+        so consumers do not have to ``shape[0]`` the parent tensor.
+    precomputed
+        Opaque ``dict`` returned by ``flash_attn_interface.precompute_tree_metadata``.
+        Reused across all transformer layers in a forward / backward pass.
+    """
+
+    cu_node_lens: Tensor
+    node_parent: Tensor
+    tree_position_ids: Tensor
+    padded_size: int
+    num_nodes: int
+    precomputed: Optional[Dict[str, Any]] = field(default=None)