From 3d102686769821b5273d9a989ffcae508d1942d3 Mon Sep 17 00:00:00 2001 From: racky-scitix Date: Wed, 20 May 2026 11:56:08 +0800 Subject: [PATCH 1/5] feat(tree-native): add tree-attention metadata + position_ids surface for tree training (#3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(tree-native): add Megatron PackedSeqParams.tree_metadata + TreeMetadata dataclass Stage 0.2 of the tree-training native migration. Carries trie topology (cu_node_lens, node_parent), per-token positional ids (tree_position_ids), and the FA3 precomputed attention buffers from slime's TreeDataIterator through PackedSeqParams to the upcoming TETreeDotProductAttention and TreePackedRotaryEmbedding. PackedSeqParams.tree_metadata defaults to None so all existing callers are untouched; TYPE_CHECKING-guarded import keeps the new dataclass off the import path until tree training is actually used. Refs #220. * feat(tree-native): add Megatron tree-attention extensions Stage 1.2 of the tree-training native migration: - TreePackedRotaryEmbedding (rotary_pos_embedding.py): subclass of RotaryEmbedding that gathers freqs by packed_seq_params.tree_metadata .tree_position_ids when present. apply_rotary_pos_emb (TE side) stays untouched — it is pure elementwise so once freqs are reordered the rest of the rope path is correct without TE modifications. - AttnBackend.tree (enums.py): new enum value to advertise the tree attention path. - TETreeDotProductAttention + _FA3TreeAttnFunc (extensions/transformer_engine.py): subclass of TEDotProductAttention that asserts cp_size == 1, forces qkv_format='thd', and routes through the FA3 tree kernel (flash_attn_interface.flash_attn_tree_func) directly. The direct kernel call is a transitional shape: once scitix/TransformerEngine adds TreeFlashAttention as a real TE backend (Stage 1.1), the kernel call here can be replaced by a TE backend dispatch. - get_gpt_layer_tree_spec (gpt_layer_specs.py): GPT layer spec that swaps core_attention for TETreeDotProductAttention and advertises AttnMaskType.arbitrary so the upstream block does not synthesize an unused mask. Refs #220. * refactor(tree-native): slim TETreeDotProductAttention to TE delegation Now that scitix/TransformerEngine ships TreeFlashAttention as a first-class backend in DotProductAttention's dispatch (scitix/TE commit deb33ef7), the Megatron-side tree wrapper no longer needs its own FA3 autograd wrapper or kernel call. - Remove _FA3TreeAttnFunc (moved into TE's backends.py). - TETreeDotProductAttention.forward now unpacks packed_seq_params.tree_metadata into tree_cu_node_lens / tree_node_parent / tree_precomputed kwargs and calls te.pytorch.DotProductAttention.forward directly (skipping the Megatron-side TEDotProductAttention.forward wrapper, which would drop the tree kwargs). TE's dispatch routes to TreeFlashAttention when those kwargs are present. Prerequisite for Stage 2 CP support: CP-aware tree attention will land inside TE's TreeFlashAttention.forward (via TE's existing cp_group/cp_comm_type infrastructure), and this wrapper will not need to change. Refs #220. * refactor(tree-native): data-driven tree dispatch, eliminate spec-level tree classes Tree routing is now fully data-driven: same model spec, same TEDotProductAttention instance, same RotaryEmbedding instance handle both tree and varlen micro-batches. When PackedSeqParams.tree_metadata is populated, both Megatron and TE route to the tree path automatically. Deleted: - TETreeDotProductAttention (Megatron) — tree detection folded into TEDotProductAttention.forward as a ~15-line early return - TreePackedRotaryEmbedding (Megatron) — tree-position gather folded into RotaryEmbedding.forward as a ~12-line suffix - get_gpt_layer_tree_spec (Megatron) — no longer needed - model_provider.py tree spec selection + MoE block spec post-process + RotaryEmbedding __class__ swap — all deleted This eliminates the 'MoE decoder block spec post-processing hack' gap from the Known Gaps list entirely. Refs #220. * refactor(tree-native): remove tree-specific code from TEDotProductAttention TE now accepts tree_metadata as a direct kwarg on DotProductAttention .forward (TE commit ebfa49a6). Megatron's packed_seq_kwargs dict already includes tree_metadata (it is a PackedSeqParams dataclass field), so it flows through to TE transparently via **packed_seq_kwargs. Delete the 15-line tree detection + manual TE dispatch block from TEDotProductAttention.forward — zero tree-specific code remains in Megatron's attention wrapper. Refs #220. * refactor(tree-native): generic position_ids on PackedSeqParams for RoPE Decouple RotaryEmbedding from tree-training: instead of reading tree_metadata.tree_position_ids directly, RotaryEmbedding.forward now checks the generic packed_seq_params.position_ids field. This follows the same pattern as MultimodalRotaryEmbedding accepting position_ids. Changes: - PackedSeqParams gains position_ids: Optional[Tensor] = None - RotaryEmbedding.forward gathers emb by position_ids when present — no 'tree' concept in the code, just 'custom per-token positions' - slime get_tree_batch sets both position_ids and tree_metadata on PackedSeqParams - TreeMetadata docstring updated: tree_position_ids is now consumed indirectly via PackedSeqParams.position_ids The position_ids field is reusable by any future feature needing non-sequential per-token RoPE positions. Refs #220. * refactor(tree-native): position_ids flows through model.forward, not PackedSeqParams Move per-token positional ids to the same level as other model.forward args instead of hiding them inside PackedSeqParams: - RotaryEmbedding.forward gains an explicit position_ids parameter (same pattern as MultimodalRotaryEmbedding). When provided, gathers emb by position_ids for non-sequential per-token RoPE. - GPTModel._preprocess passes position_ids through to self.rotary_pos_emb(..., position_ids=position_ids). - slime model.py: all three model() call sites changed from position_ids=None to position_ids=batch.get('position_ids'). For varlen batches, get() returns None (no behavior change). For tree batches, get_tree_batch already puts position_ids in the batch dict. - PackedSeqParams.position_ids field removed — no longer needed. - get_tree_batch: removed position_ids from PackedSeqParams construction. Refs #220. * refactor(tree-native): drop stray blank line in transformer_engine.py Leftover from the tree-specific code removal in 8b600e448; net diff vs sirl-dev base for this file is now empty. --------- Co-authored-by: racky-scitix Co-authored-by: racky-scitix --- .../common/embeddings/rotary_pos_embedding.py | 24 ++++++++- megatron/core/models/gpt/gpt_model.py | 4 +- megatron/core/packed_seq_params.py | 10 ++++ megatron/core/transformer/enums.py | 4 ++ megatron/core/transformer/tree_metadata.py | 51 +++++++++++++++++++ 5 files changed, 91 insertions(+), 2 deletions(-) create mode 100644 megatron/core/transformer/tree_metadata.py diff --git a/megatron/core/models/common/embeddings/rotary_pos_embedding.py b/megatron/core/models/common/embeddings/rotary_pos_embedding.py index 5d7b69cd34e..a511697ef9b 100644 --- a/megatron/core/models/common/embeddings/rotary_pos_embedding.py +++ b/megatron/core/models/common/embeddings/rotary_pos_embedding.py @@ -177,7 +177,11 @@ 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. @@ -185,6 +189,10 @@ def forward( 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): diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 1fd52f65a5d..1ff4b34a036 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/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: diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 08ebdac67d8..d91a8fdbbc6 100644 --- a/megatron/core/packed_seq_params.py +++ b/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 diff --git a/megatron/core/transformer/enums.py b/megatron/core/transformer/enums.py index d06d58d65f2..047a7de2435 100644 --- a/megatron/core/transformer/enums.py +++ b/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): diff --git a/megatron/core/transformer/tree_metadata.py b/megatron/core/transformer/tree_metadata.py new file mode 100644 index 00000000000..f825848b0a6 --- /dev/null +++ b/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) From d0b4183420f583a10d3bbc6f77ab6b6ff71a95cf Mon Sep 17 00:00:00 2001 From: tonic-scitix Date: Wed, 27 May 2026 10:42:27 +0800 Subject: [PATCH 2/5] feat(sparserl-sync): integrate SparseRL-Sync sparse diff tracking into optimizer (#5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wrap the in-place param.copy_() calls in DistributedOptimizer and HybridDeviceOptimizer with sparse_diff_context() so the attached SparseManager can snapshot pre/post state and build per-param sparse-update indices for the Trainer→Rollout weight broadcast. Also pass the communication group to P2POp constructors in p2p_communication.py (group was omitted; required for non-default PG usage). Co-authored-by: Claude --- .../cpu_offloading/hybrid_optimizer.py | 18 ++++++++- megatron/core/optimizer/distrib_optimizer.py | 40 ++++++++++++++++++- .../pipeline_parallel/p2p_communication.py | 8 ++-- 3 files changed, 59 insertions(+), 7 deletions(-) diff --git a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py index 28487c3b367..2c3123e229a 100644 --- a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py +++ b/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 diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index eac21a3ea8e..7b35417b6bc 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/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) diff --git a/megatron/core/pipeline_parallel/p2p_communication.py b/megatron/core/pipeline_parallel/p2p_communication.py index f18309217c3..ac839c21f18 100644 --- a/megatron/core/pipeline_parallel/p2p_communication.py +++ b/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: From 7ea2af438885a885ac822ff95b4422580dab72de Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 23 Jun 2026 17:17:49 +0800 Subject: [PATCH 3/5] feat(moe): allow router module injection --- megatron/core/transformer/moe/moe_layer.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 10d10f667fe..9f1a792d8ff 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/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": From 2f356f9c2f2d87b0e0da249cd782953de6d52dc3 Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 23 Jun 2026 17:45:06 +0800 Subject: [PATCH 4/5] fix(transformer): keep tree metadata out of regular TE attention --- .../core/extensions/transformer_engine.py | 45 +++++++++++++- .../transformer/test_attention_packed_seq.py | 59 +++++++++++++++++++ 2 files changed, 103 insertions(+), 1 deletion(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index d239db4ab0c..055682c48ba 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/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 {} ) diff --git a/tests/unit_tests/transformer/test_attention_packed_seq.py b/tests/unit_tests/transformer/test_attention_packed_seq.py index e6e2c847395..e397a751347 100644 --- a/tests/unit_tests/transformer/test_attention_packed_seq.py +++ b/tests/unit_tests/transformer/test_attention_packed_seq.py @@ -3,11 +3,13 @@ import pytest import torch +from megatron.core.extensions.transformer_engine import _packed_seq_kwargs_for_te from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.attention import SelfAttention from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.tree_metadata import TreeMetadata from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import is_te_min_version from tests.unit_tests.test_utilities import Utils @@ -44,6 +46,63 @@ def make_test_packed_padded_seq_params(sequence_length): return packed_seq_params +def test_packed_seq_kwargs_omit_empty_tree_metadata(): + packed_seq_params = PackedSeqParams( + cu_seqlens_q=torch.IntTensor([0, 1]), + cu_seqlens_kv=torch.IntTensor([0, 1]), + max_seqlen_q=1, + max_seqlen_kv=1, + qkv_format='thd', + ) + + kwargs = _packed_seq_kwargs_for_te( + packed_seq_params, + {"qkv_format", "cu_seqlens_q", "max_seqlen_q", "tree_metadata"}, + supports_tree_metadata=False, + ) + + assert kwargs["qkv_format"] == 'thd' + assert kwargs["max_seqlen_q"] == 1 + assert "tree_metadata" not in kwargs + + +def test_packed_seq_kwargs_reject_tree_metadata_without_te_support(): + tree_metadata = TreeMetadata( + cu_node_lens=torch.IntTensor([0, 1]), + node_parent=torch.IntTensor([-1]), + tree_position_ids=torch.LongTensor([0]), + padded_size=1, + num_nodes=1, + ) + packed_seq_params = PackedSeqParams(qkv_format='thd', tree_metadata=tree_metadata) + + with pytest.raises(RuntimeError, match="tree_metadata requires TransformerEngine"): + _packed_seq_kwargs_for_te( + packed_seq_params, + {"qkv_format", "tree_metadata"}, + supports_tree_metadata=False, + ) + + +def test_packed_seq_kwargs_pass_tree_metadata_when_te_supports_it(): + tree_metadata = TreeMetadata( + cu_node_lens=torch.IntTensor([0, 1]), + node_parent=torch.IntTensor([-1]), + tree_position_ids=torch.LongTensor([0]), + padded_size=1, + num_nodes=1, + ) + packed_seq_params = PackedSeqParams(qkv_format='thd', tree_metadata=tree_metadata) + + kwargs = _packed_seq_kwargs_for_te( + packed_seq_params, + {"qkv_format", "tree_metadata"}, + supports_tree_metadata=True, + ) + + assert kwargs["tree_metadata"] is tree_metadata + + class TestParallelAttentionWithPackedSequence: def setup_method(self, method): From 346081d8f1a46a84b5fba153f5514952a366a71e Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 24 Jun 2026 09:20:35 +0800 Subject: [PATCH 5/5] fix(moe): keep routing replay state out of megatron --- megatron/core/transformer/moe/moe_utils.py | 3 --- megatron/core/transformer/moe/router.py | 3 --- 2 files changed, 6 deletions(-) diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 71912617f12..28cff06f5ec 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/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) diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index ae238f93119..16fc9d9af8f 100644 --- a/megatron/core/transformer/moe/router.py +++ b/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.