[GRL] Land PR400 sparse/delta megatron-lm dependency#16
Open
flowerfox-scitix wants to merge 5 commits into
Open
[GRL] Land PR400 sparse/delta megatron-lm dependency#16flowerfox-scitix wants to merge 5 commits into
flowerfox-scitix wants to merge 5 commits into
Conversation
… for tree training (#3) * 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 NVIDIA#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 NVIDIA#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 NVIDIA#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 NVIDIA#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 NVIDIA#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 NVIDIA#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 NVIDIA#220. * refactor(tree-native): drop stray blank line in transformer_engine.py Leftover from the tree-specific code removal in 8b600e4; net diff vs sirl-dev base for this file is now empty. --------- Co-authored-by: racky-scitix <yellowhuangjiamingyellow@gmail.com> Co-authored-by: racky-scitix <racky-scitix@users.noreply.github.com>
…o optimizer (#5) 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 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
This PR promotes the Megatron-side pieces consumed by slime-trainer PR400's sparse/delta and routed-training development branch into
sirl-stable. SGLang owns the delta applier endpoints; Megatron owns the producer/training-side state that makes those payloads and routed batches usable without leaking runtime-only metadata into trainer math.Modifications
Verification
Run locally on current PR head
346081d8f1a46a84b5fba153f5514952a366a71e:A broader
tests/unit_tests/transformer/test_attention_packed_seq.pyinvocation reaches CUDA model-parallel setup on this host and fails because no CUDA device is visible; the three targeted tree-metadata contract tests above avoid that host-specific CUDA path.slime-trainer Context
scitix/slime-trainer#400.codex/next-rl-framework-terminal-state.20260623T180705452306_v3psd28ncompleted successfully withlive_proof=truefor the development sparse/delta chain.scitix/sglang#29.The slime-trainer submodule pointer must not move to this PR head directly. It must move only after this PR is merged into
sirl-stable, the annotated stable-source tag exists, and the downstream sparse/delta readiness gate verifies the tagged stable dependency.Stable Release
Planned stable-source tag:
scitix-megatron-sirl-stable-2026.06.25.1if this PR lands on 2026-06-25 UTC; use the same tag family with the landing date / sequence number if it lands later.Stable branch target:
sirl-stable.Stable commit to tag after merge: the
sirl-stablemerge commit that contains PR head346081d8f1a46a84b5fba153f5514952a366a71e.Downstream consumer:
scitix/slime-trainerPR NVIDIA#400 / branchcodex/next-rl-framework-terminal-state. The downstream PR must updatevendors/megatron-lmonly after this PR is merged, the annotated stable-source tag exists, and the slime-trainer sparse/delta readiness gate verifies that the submodule gitlink points at the tagged stable commit.Stable release evidence to carry into the annotated tag and downstream Vendor Release Map:
WORLD_SIZE=1 RANK=0 LOCAL_RANK=0 MASTER_ADDR=127.0.0.1 MASTER_PORT=29501 PYTHONPATH=$PWD pytest tests/unit_tests/transformer/test_attention_packed_seq.py::test_packed_seq_kwargs_omit_empty_tree_metadata tests/unit_tests/transformer/test_attention_packed_seq.py::test_packed_seq_kwargs_reject_tree_metadata_without_te_support tests/unit_tests/transformer/test_attention_packed_seq.py::test_packed_seq_kwargs_pass_tree_metadata_when_te_supports_it -q(3 passed)20260623T180705452306_v3psd28ncompleted withlive_proof=truefor the sparse/delta development branch proof.