Add Megatron-Core integration for RoPE - #1325
Draft
PrathyushaPolepalli wants to merge 4 commits into
Draft
Conversation
Adds Liger's Triton RoPE to Megatron-Core, following the same one-line monkey-patch pattern as the RMSNorm (linkedin#1254) and cross-entropy (linkedin#1207, linkedin#1260) integrations. apply_liger_kernel_to_megatron(rope=True) reroutes Megatron's apply_rotary_pos_emb dispatcher — rebinding the symbol on every module that imported it by value (e.g. megatron.core.transformer.attention) — to Liger's kernel for the standard unfused bshd path. Fused (TE/Apex) RoPE, packed thd sequences, interleaved rotation, multi-latent attention, mscale scaling and per-batch/mRoPE freqs transparently fall back to the native implementation so numerics never silently change. Megatron applies RoPE once per tensor (query, then key) while Liger's kernel is fused q/k; LigerMegatronRopeFunction reuses the existing kernel by passing a one-head throwaway k (negligible extra work), avoiding a second Triton kernel. - src/liger_kernel/megatron/rope.py: adapter + single-tensor autograd op - monkey_patch.py: rope flag + _patch_apply_rotary_pos_emb (rebind-everywhere) - __init__.py: exports - test/megatron/test_rope.py: correctness vs native bshd reference - test/megatron/test_monkey_patch.py: RoPE patch-mechanism tests - benchmark/scripts/benchmark_megatron_rope.py: liger vs torch vs megatron - examples/megatron/README.md: rope=True docs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Accept and preserve Megatron's current mla_rotary_interleaved argument so MLA calls fall back to the native implementation instead of failing at the patched wrapper. Detect older Megatron signatures before forwarding the new argument to keep the integration backward compatible. Add regression coverage for both API generations and update the Megatron high-level API documentation for RoPE and vocab-parallel cross-entropy. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Use SBHD for Megatron's actual tensor order and BHSD for Liger's input order. Keep the public bshd-suffixed adapter name because it mirrors Megatron's upstream helper, and document that distinction explicitly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.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
1 participant
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.
Adds Liger's Triton RoPE to Megatron-Core
apply_liger_kernel_to_megatron(rope=True) reroutes Megatron's apply_rotary_pos_emb dispatcher — rebinding the symbol on every module that imported it by value (e.g. megatron.core.transformer.attention) — to Liger's kernel for the standard unfused bshd path. Fused (TE/Apex) RoPE, packed thd sequences, interleaved rotation, multi-latent attention, mscale scaling and per-batch/mRoPE freqs transparently fall back to the native implementation so numerics never silently change.
Megatron applies RoPE once per tensor (query, then key) while Liger's kernel is fused q/k; LigerMegatronRopeFunction reuses the existing kernel by passing a one-head throwaway k (negligible extra work), avoiding a second Triton kernel.
Summary
Testing Done
make testto ensure correctnessmake checkstyleto ensure code stylemake test-convergenceto ensure convergence