Skip to content

Fix packed THD interleaved RoPE fusion - #9

Draft
akhauriyash wants to merge 1 commit into
prodfrom
agent/fix-packed-thd-interleaved-rope
Draft

akhauriyash wants to merge 1 commit into
prodfrom
agent/fix-packed-thd-interleaved-rope

Conversation

@akhauriyash

Copy link
Copy Markdown
Collaborator

Summary

  • Forward rotary_interleaved into Transformer Engine's fused packed-THD RoPE path.
  • Use the TE >= 2.3.0 packed-THD API and fail closed when an older TE build is asked to run interleaved fused RoPE.
  • Add an unsupported-TE regression test and fused-vs-unfused GPU parity coverage for both rotary layouts.

Why

MoVA trains with rotary_interleaved=True and packed THD sequences. The existing fused THD path drops the interleaving flag, so the trainer and SGLang can apply different rotary layouts. That invalidates trainer/inference log-probability parity and importance-sampling ratios.

Retain --no-rope-fusion for MoVA training until this change is accepted by the full end-to-end parity gate and deployed in the pinned training image.

Validation

  • H200 Slurm job 2146011: PASS.
  • Focused pytest result: 3 passed, 0 failed, 0 errored, 0 skipped.
  • Evidence: /mnt/weka/shrd/k2pta/rl360/parity/2135686/megatron-thd-rope-patch-test-2146011/result.json.
  • Published diff is byte-for-byte identical to the tested patch (SHA-256 42096367b13bec5809f214b218c4588391dc7ca5d8e833c40e317ad963fe886e).
  • git diff --check and AST parsing of all three changed files pass.

Forward rotary_interleaved into Transformer Engine's packed THD fused path. Without it, MoVA trainer and SGLang apply different rotary layouts and produce invalid off-policy ratios.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant