Support deterministic Megatron sequence parallelism for G11 consistency

[inaniloquentee]

Summary

• Enable strict Megatron sequence parallelism for QKV and output projections with deterministic TP all-gather/reduce-scatter and matching autograd communication.
• Support the strict TP LM head on sequence-sharded hidden states while retaining the existing canonical TP/CP gradient path and non-SP all-reduce order.
• Run the Qwen3-8B G11 example with training CP=2, TP=4, SP=4. Add an opt-in entropy coefficient for a nonzero-update consistency probe; the default remains 0.

Validation

• On H100 against this branch based on main (bea224b): 101 targeted tests passed (framework adapters, G11 example, canonical LM-head/CP gradients, Megatron runtime state).
• The earlier G11 end-to-end run on cb05cfc plus the same SP changes used two rollouts with CP=2, TP=SP=4 and entropy coefficient 0.01: 4,096 active train/rollout tokens matched bitwise; gradients were nonzero and weight versions changed. The default-coefficient two-rollout probe also had zero token mismatches, but zero gradient norms. The full 200-rollout suite has not been run on this PR branch.

Integration note

G11 requires a companion Vime hook change: gather sequence-sharded hidden states before the strict training-logprob provider, and preserve its expected token/logit layout. The H100 end-to-end result above used that local Vime change. RL-Kernel's SP training path is covered here, but the G11 end-to-end command cannot be reproduced with unmodified Vime until the companion hook is upstreamed.

Summary by CodeRabbit

• New Features
  • Training runs can use sequence parallelism with compatible attention and language-model output projections.
  • Added an --entropy-coef option to configure the entropy coefficient. It defaults to 0.0 and rejects negative or non-finite values.
  • The selected entropy coefficient is recorded in the run manifest.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: ae204fdb-76c5-441f-a919-93c490d1a487

---

📥 Commits

Reviewing files that changed from the base of the PR and between c698b53 and 2902456.

---

📒 Files selected for processing (2)

• rl_engine/integrations/megatron_runtime.py
• tests/test_framework_runtime_adapters.py

---

Included review availability: This review used your included allowance. Your plan provides up to 2 included reviews per hour; 0 remain after this review.

---

---

📝 Walkthrough

Walkthrough

The example runner now accepts and records an entropy coefficient and enables sequence parallelism. Megatron attention and the strict TP LM head add sequence-parallel gather and reduce-scatter behavior, with tests for collective operations and gradient layouts.

Changes

Sequence-parallel execution

Layer / File(s)	Summary
Sequence-parallel autograd collectives rl_engine/integrations/megatron_runtime.py, tests/test_framework_runtime_adapters.py	New autograd operations pair sequence-shard gathering with gradient reduce-scatter, and output reduce-scatter with gradient gathering. Tests check collective behavior and gradients.
Attention projection collectives rl_engine/integrations/megatron_runtime.py, examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py, tests/test_framework_runtime_adapters.py	Attention gathers QKV inputs and reduce-scatters output-projection results when sequence parallelism is active. Initialization requires matching QKV and output-projection settings and rejects callback overrides for sequence-parallel execution. The runner enables sequence parallelism. Tests check the projections and input gradients.
Sequence-parallel LM head rl_engine/integrations/megatron_runtime.py, tests/test_canonical_lm_head_backward.py, tests/test_framework_runtime_adapters.py	The strict TP LM head gathers its input and reduce-scatters its input gradient when sequence parallelism is active. Tests check gathered logits, gradient layouts, and weight gradients.

Entropy coefficient configuration

Layer / File(s)	Summary
Entropy coefficient option examples/vime_qwen3_8b_tp4_cp2_200/run_arm.py, tests/test_vime_tp4_example.py	The runner adds --entropy-coef, validates that the value is finite and nonnegative, passes it to training, and records it in the manifest. The parser test checks the default and an explicit value.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant MegatronAttention
  participant DeterministicTPCollective
  participant TPGroup
  MegatronAttention->>DeterministicTPCollective: Gather QKV sequence shards
  DeterministicTPCollective->>TPGroup: All-gather sequence shards
  TPGroup-->>DeterministicTPCollective: Return gathered QKV input
  MegatronAttention->>DeterministicTPCollective: Reduce-scatter output projection
  DeterministicTPCollective->>TPGroup: Reduce-scatter projection output
  TPGroup-->>DeterministicTPCollective: Return sequence-sharded output

Loading

---

Merge Risk

Merge Risk: 🔵 Low · up to 29024

This enables sequence-parallel training for strict Megatron attention and the LM head. Targeted tests pass, but the full 200-rollout suite has not been run on this branch, and the end-to-end command needs a companion Vime change that is not yet upstreamed.

Security Architecture Review

Security architecture risk: 🔵 Low · up to 29024

The inspected changes keep communication within the existing tensor-parallel group and retain explicit layout checks. No introduced security issue was substantiated. End-to-end integration, distributed interruption recovery, and concurrent operation ordering remain incompletely verified.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — The demonstrated communication exposure is bounded to activations and gradients among existing TP peers in the training job. The inspected change does not establish new tenant, credential, service, or environment authority, although complete security coverage was not supplied.

Trust Boundaries and Controls

• observed — Collectives validate rank signatures and tensor layouts, including dimension-zero divisibility for reduce-scatter. SP attention rejects fallback mappings that lack the required deterministic collective, and the LM head checks that backward restores the caller's input layout. These are execution-contract controls, not authentication controls.

Resilience and Maintainability Implications

• observed — The reused collective implementation locally serializes operations, synchronizes participating ranks, and preserves live borrowers when replacing a cached collective. These mechanisms predate this PR; they do not establish distributed abort/reset behavior or global ordering under concurrent calls. No introduced recovery defect was substantiated.

Hardening Proposals

• proposed — Before broader rollout, validate the companion hook and canonical CP/SP layout agreement across ranks, and establish an explicit job-wide abort or restart policy for interrupted collectives. These are containment proposals, not observed vulnerabilities.

---

---

---

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 4.65% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 43 functions across 5 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly summarizes the primary change: deterministic Megatron sequence parallelism for G11 consistency.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

---

---

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 2

📝 Generate docstrings 💡

• Commit to this branch
• Create a new PR

---

🛠️ Fix failing CI checks 💡

• Commit to this branch
• Create a new PR

---

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

---

---

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

_{Comment @coderabbitai help to get the list of available commands.}