Add opt-in Turbo fused cross entropy for TP1 training

[suachong]

Megatron TP1 training can opt into Turbo fused cross entropy through the declared use_turbo_cross_entropy training argument. The patch preserves Core's sequence/batch layout and per-token loss contract; unsupported tensor parallelism, dtypes, and unfused configurations retain the configured implementation.

turbo_ce_overwrite_input additionally permits backward to reuse exclusively owned contiguous logits storage. Both arguments default false in primus/configs/modules/megatron/primus_turbo.yaml and support experiment YAML and Primus CLI overrides. The patch no longer reads PRIMUS_TURBO_CROSS_ENTROPY or PRIMUS_TURBO_CE_OVERWRITE_INPUT. Documentation is linked from the docs index and site navigation.

The Turbo dependencies #546 and #550 are merged. CI pins Turbo main at 4e6bd90b4fd9aa40c7945674d83975ee150a7b68, which includes both MoE gather fusion and fused cross entropy.

Training subprocess tests now use independent c10d rendezvous ports allocated and held by torchrun. This fixes inherited MASTER_PORT collisions in pipeline-warmup and projection tests. A regression test launches two concurrent distributed jobs while the inherited port is occupied. Multi-node test settings are preserved.

Validation:

• Full CI passed for 99439f8d4: https://github.com/AMD-AGI/Primus/actions/runs/38112167456
• MI300X PyTorch: 3,219 core tests passed (1,039 skipped, 8 deselected), including all 13 cross-entropy tests and the new concurrent-rendezvous regression. All 19 Megatron and 7 TorchTitan training tests passed, as did CLI smoke checks.
• TorchTitan Llama FP8 reported finite losses (12.26768, 22.06730, 22.74743) and gradients for all three steps. An earlier MI300X NaN failure did not recur; the rendezvous fix addresses port collisions, and no numerical workaround or relaxed metric assertion was added.
• All 25 focused local rendezvous, finite-metric, and cross-entropy configuration tests passed.
• Five local eight-GPU MI355X runs of the three-step TorchTitan Llama FP8 recipe passed with finite metrics, including fixed seeds 0/1/2 and a run through the updated training helper while the inherited port was deliberately occupied.
• JAX, lint, image build, dependency review, coverage summary, CodeQL, and documentation checks passed.

[Copilot]

🔵 Needs a closer look

The GPU-specific autograd path depends on an external Turbo pull request that remains unmerged and requires final human validation.

0 open findings

What changed in this PR

Adds opt-in Primus-Turbo fused cross-entropy dispatch for Megatron TP=1 training.

Changes:

• Wraps Megatron’s language-model loss with supported-case dispatch and fallback.
• Adds GPU integration tests and usage/performance documentation.
• Pins CI to the required Turbo operator commit.

File	Description
.github/​workflows/​ci.yaml	Pins the Turbo dependency for GPU tests.
primus/​backends/​megatron/​patches/​turbo/​cross_entropy_patches.py	Implements opt-in CE dispatch.
tests/​unit_tests/​test_turbo_cross_entropy_patch.py	Tests layout, gradients, overwrite mode, and fallback.
docs/​features/​turbo_cross_entropy.md	Documents configuration, constraints, and validation.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🔵 Needs a closer look

It modifies core loss and gradient behavior using an unmerged external dependency without full-model validation on the current parent.

0 open findings

Previously missed (1)

In code that hasn't changed since last review

Add the new documentation page to navigation and TOC

docs/​features/​turbo_cross_entropy.md:1

This page is not included in the documentation navigation, so the new usage instructions will not be published/discoverable through the rendered docs. The site uses the explicit list in docs/sphinx/_toc.yml.in, and docs/README.md also has no link to this sole docs/features/ page. Please move/link it under an existing section (for example Technical guides) and add it to the TOC.

🧠 Review effort: Balanced

[Copilot]

🔵 Needs a closer look

The GPU kernel integration depends on an unmerged, currently blocked Primus-Turbo pull request and requires final GPU CI confirmation.

0 open findings

🧠 Review effort: Balanced

[wenxie-amd]

Don't use an env to enable the feature, add a new argument to megatron and make it visible to more users.

[Copilot]

🔵 Needs a closer look

The fused autograd dependency remains unmerged and current GPU CI and full-model validation are still pending.

0 open findings

🧠 Review effort: Balanced

[Copilot]

🔵 Needs a closer look

Fresh GPU CI is still pending for the updated external kernel revision and merge commit.

0 open findings

🧠 Review effort: Balanced

[Copilot]

🔵 Needs a closer look

The core GPU autograd path is well covered structurally, but fresh GPU CI for the updated Turbo pin remains pending.

0 open findings

🧠 Review effort: Balanced

[Copilot]

🔵 Needs a closer look

Full CI remains in progress, and the previously observed MI300X non-finite-metrics failure has not yet been resolved or reproduced.

0 open findings

🧠 Review effort: Balanced