Add opt-in Turbo fused cross entropy for TP1 training

.github/workflows/ci.yaml

@@ -42,7 +42,8 @@ env:
   # Includes fused expert wgrad output overwrite support (#534), MegaMoE-owned FlyDSL
   # kernels (#525), the repaired DeepEP intranode cheap fence, on by default (#539),
   # and MegaMoE SwiGLU-OAI activations (#538).
-  PRIMUS_TURBO_COMMIT: 8f19706aa1c656b11a146eb00db32a6c2b46822d
+  # Turbo main includes merged MoE gather fusion (#546) and fused cross entropy (#550).
+  PRIMUS_TURBO_COMMIT: 4e6bd90b4fd9aa40c7945674d83975ee150a7b68
   PRIMUS_TURBO_AITER_COMMIT: 0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e # AITER v0.1.14.post1 (tag commit) — required by Primus-Turbo main aiter_utils.py
   #ROCSHMEM_COMMIT: 17ff985c026f9f97f85068647e863ab541dd5645 # Update version to 3.2.0 for 7.2.0 rocm release (#351) (#355)
   UCCL_COMMIT: 5afb4117893c58cc0c8557d9286336141a301053 # [EP]: fix fp8 error of internode_ll on amd gfx950 arch. (#710)

docs/README.md

@@ -68,6 +68,7 @@ Deep technical topics for advanced users.
 - [Parallelism configuration](./04-technical-guides/parallelism-configuration.md): per-backend parallelism setup and batch size relationships
 - [Collective operations](./04-technical-guides/collective-operations.md): NCCL/RCCL operations and their role in each parallelism strategy
 - [Performance tuning](./04-technical-guides/performance-tuning.md): HipBLASLt, Primus-Turbo, FP8, MoE optimization
+- [Turbo cross entropy](./features/turbo_cross_entropy.md): TP1 fused loss and configuration
 - [MoE training deep-dive](./04-technical-guides/moe-training.md): bottlenecks and Primus-Turbo optimizations for Mixture-of-Experts models
 - [MegaMoE fused MoE layer](./04-technical-guides/mega-moe.md): FlyDSL-based fused MoE layer for EP-only bf16 training, setup and reproduction
 - [Data preparation](./04-technical-guides/data-preparation.md): tokenization, data formats, mock data

docs/features/turbo_cross_entropy.md

@@ -0,0 +1,42 @@
+# Opt-in Turbo cross entropy for Megatron training
+
+Install a Primus-Turbo version with `primus_turbo.pytorch.ops.cross_entropy`, then
+set the following Megatron training arguments in the experiment YAML:
+
+```yaml
+use_turbo_cross_entropy: true
+turbo_ce_overwrite_input: false
+cross_entropy_loss_fusion: true
+```
+
+These options are declared in `primus/configs/modules/megatron/primus_turbo.yaml`
+and can also be supplied as Primus overrides, for example
+`--use_turbo_cross_entropy true --turbo_ce_overwrite_input false`.
+
+The patch replaces
+Megatron's language-model loss for TP=1, BF16/FP32 logits when
+`cross_entropy_loss_fusion` is enabled. TP>1, other dtypes, and an unfused loss
+configuration use the existing configured implementation.
+
+The patch preserves Core's layout and reduction contract: sequence-first logits,
+batch-first labels and output per-token loss. Training still applies its mask,
+normalization and gradient scaling. Do not set `cross_entropy_fusion_impl=turbo`;
+the opt-in patch provides the integration while retaining the original Core
+configuration for fallback.
+
+`turbo_ce_overwrite_input: true` additionally permits backward to reuse
+contiguous logits storage. Enable it only when the language-model logits have
+no other consumer that needs their original values. Noncontiguous logits use
+the preserving path. Both arguments default to false. Set
+`use_turbo_cross_entropy: false` to restore the configured loss path without
+changing packages.
+
+This integration changes no installed Transformer Engine files. It requires no
+TE native rebuild; Triton JIT compilation occurs on first use. The operator
+supports one backward per forward and no higher-order gradients. Language-model recipes need
+numerical and convergence validation because its FP32 gradient arithmetic
+differs from older TE's intermediate BF16 rounding.
+
+The former `PRIMUS_TURBO_CROSS_ENTROPY` and `PRIMUS_TURBO_CE_OVERWRITE_INPUT`
+environment variables are no longer read by the patch. Launchers that retain
+environment overrides must resolve them into these YAML arguments.

docs/sphinx/_toc.yml.in

@@ -90,6 +90,8 @@ subtrees:
         title: Data preparation
       - file: 04-technical-guides/checkpoint-management.md
         title: Checkpoint management
+      - file: features/turbo_cross_entropy.md
+        title: Turbo cross entropy
       - file: 04-technical-guides/moe-training.md
         title: MoE training
       - file: 04-technical-guides/mega-moe.md

primus/backends/megatron/patches/turbo/cross_entropy_patches.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# See LICENSE for license information.
+
+"""Opt-in Turbo TP=1 vocabulary CE without replacing installed TE files."""
+
+from functools import wraps
+
+import torch
+
+from primus.core.patches import PatchContext, get_args, register_patch
+from primus.core.utils.module_utils import log_rank_0
+
+
+def _enabled(ctx: PatchContext) -> bool:
+    return bool(getattr(get_args(ctx), "use_turbo_cross_entropy", False))
+
+
+@register_patch(
+    "megatron.turbo.cross_entropy",
+    backend="megatron",
+    phase="before_train",
+    description="Use Turbo fused CE with deferred backward for TP=1 BF16/FP32 logits",
+    condition=_enabled,
+)
+def patch_cross_entropy(ctx: PatchContext) -> None:
+    from megatron.core.models.common.language_module.language_module import (
+        LanguageModule,
+    )
+    from primus_turbo.pytorch.ops.cross_entropy import cross_entropy
+
+    original = LanguageModule.compute_language_model_loss
+    if getattr(original, "_primus_turbo_ce", False):
+        return
+    overwrite_input = bool(getattr(get_args(ctx), "turbo_ce_overwrite_input", False))
+
+    @wraps(original)
+    def compute_language_model_loss(self, labels, logits):
+        if (
+            self.config.tensor_model_parallel_size != 1
+            or logits.dtype not in (torch.bfloat16, torch.float32)
+            or not self.config.cross_entropy_loss_fusion
+        ):
+            return original(self, labels, logits)
+        # Core supplies sequence-first logits and batch-first targets. Keep
+        # the original per-token loss contract; training owns masks/scaling.
+        target = labels.transpose(0, 1).contiguous()
+        loss = cross_entropy(
+            logits,
+            target,
+            overwrite_input=overwrite_input and logits.is_contiguous(),
+        )
+        return loss.transpose(0, 1).contiguous()
+
+    compute_language_model_loss._primus_turbo_ce = True
+    LanguageModule.compute_language_model_loss = compute_language_model_loss
+    log_rank_0(
+        "[Patch:megatron.turbo.cross_entropy] Enabled Turbo CE for TP=1 BF16/FP32; "
+        f"overwrite_input={overwrite_input}; other cases retain the configured loss"
+    )

primus/configs/modules/megatron/primus_turbo.yaml

@@ -103,3 +103,11 @@ moe_permute_quant_fusion: false
 moe_backward_permute_quant_fusion: false
 # Preserve activation storage across identity expert sorting.
 moe_skip_identity_sort: true
+
+# ===== Cross entropy =====
+# Use Turbo fused cross entropy for TP=1 BF16/FP32 logits when
+# cross_entropy_loss_fusion is enabled. Other cases keep the configured loss.
+use_turbo_cross_entropy: false
+# Allow backward to overwrite contiguous logits with their gradient. Only enable
+# when no other consumer needs the original logits.
+turbo_ce_overwrite_input: false

tests/trainer/test_megatron_trainer.py

@@ -7,23 +7,13 @@
 
 import os
 import re
-import socket
 import subprocess
 import sys
 import unittest
 
 import pytest
 
-from tests.utils import PrimusUT, run_training_script
-
-
-def _find_free_port() -> int:
-    """Ask the kernel for a free TCP port (bind to 0); avoids EADDRINUSE from
-    guessing a port that overlaps the ephemeral range or a not-yet-released one."""
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("127.0.0.1", 0))
-        return s.getsockname()[1]
-
+from tests.utils import PrimusUT, isolated_torchrun_env, run_training_script
 
 _GFX_TO_PLATFORM = {
     "gfx942": "MI300X",
@@ -1031,11 +1021,7 @@ def _run_benchmark(self, suite: str, mode_flag: str) -> str:
             "--target-nodes",
             "2",
         ]
-        env = os.environ.copy()
-        # Kernel-assigned free port so back-to-back GPU benches don't hit
-        # EADDRINUSE (a fixed/random port can clash with a not-yet-released
-        # previous run or the ephemeral range).
-        env["MASTER_PORT"] = str(_find_free_port())
+        env = isolated_torchrun_env()
         return self._check(subprocess.run(cmd, capture_output=True, text=True, env=env))
 
     def test_performance_simulate_dense(self):

tests/unit_tests/backends/megatron/test_pp_warmup_patches.py

@@ -56,6 +56,8 @@
 from pathlib import Path
 from typing import List, Optional, Tuple
 
+from tests.utils import isolated_torchrun_env
+
 _MBS = 1
 _GBS = 16
 _PRIMUS_TOTAL_LAYERS = 4
@@ -210,7 +212,7 @@ def _build_training_cmd(
         "False",
     ]
 
-    env = os.environ.copy()
+    env = isolated_torchrun_env()
     env.update(
         {
             "MBS": str(_MBS),

tests/unit_tests/test_torchrun_rendezvous.py

@@ -0,0 +1,74 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# See LICENSE for license information.
+
+"""Training test launches must not reuse a parent process's rendezvous port."""
+
+import os
+import re
+import socket
+import subprocess
+import sys
+from concurrent.futures import ThreadPoolExecutor
+
+from tests.utils import isolated_torchrun_env
+
+
+def test_concurrent_launches_with_occupied_inherited_port(tmp_path):
+    worker = tmp_path / "worker.py"
+    worker.write_text(
+        "import os, torch\n"
+        "import torch.distributed as dist\n"
+        "dist.init_process_group('gloo')\n"
+        "value = torch.tensor([dist.get_rank() + int(os.environ['TEST_JOB']) * 10])\n"
+        "dist.all_reduce(value)\n"
+        "assert value.item() == 1 + int(os.environ['TEST_JOB']) * 20\n"
+        "print('DIST_PORT=' + os.environ['MASTER_PORT'], flush=True)\n"
+        "dist.destroy_process_group()\n"
+    )
+    with socket.socket() as occupied:
+        occupied.bind(("127.0.0.1", 0))
+        occupied.listen()
+        port = occupied.getsockname()[1]
+        parent_env = dict(os.environ, NNODES="1", MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port))
+        original_env = parent_env.copy()
+
+        def launch(job):
+            env = isolated_torchrun_env(parent_env)
+            env["TEST_JOB"] = str(job)
+            return subprocess.run(
+                [
+                    sys.executable,
+                    "-m",
+                    "torch.distributed.run",
+                    "--nnodes=1",
+                    "--nproc_per_node=2",
+                    "--master_addr=127.0.0.1",
+                    f"--master_port={port}",
+                    str(worker),
+                ],
+                env=env,
+                capture_output=True,
+                text=True,
+                timeout=90,
+            )
+
+        with ThreadPoolExecutor(max_workers=2) as pool:
+            results = list(pool.map(launch, (0, 1)))
+
+        for result in results:
+            assert result.returncode == 0, result.stdout + result.stderr
+            ports = re.findall(r"DIST_PORT=(\d+)", result.stdout)
+            assert len(ports) == 2, result.stdout
+            assert len(set(ports)) == 1
+            assert ports[0] != str(port)
+        assert parent_env == original_env
+
+
+def test_multinode_rendezvous_is_preserved():
+    env = {"NNODES": "2", "MASTER_ADDR": "node0", "MASTER_PORT": "29500"}
+    assert isolated_torchrun_env(env) == env
+
+
+def test_slurm_multinode_rendezvous_is_preserved():
+    env = {"SLURM_NNODES": "2", "MASTER_ADDR": "node0", "MASTER_PORT": "29500"}
+    assert isolated_torchrun_env(env) == env

tests/unit_tests/test_turbo_cross_entropy_config.py

@@ -0,0 +1,92 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# See LICENSE for license information.
+
+"""CPU dispatch tests; numerical GPU coverage lives in the integration suite."""
+
+import importlib.util
+import sys
+from collections import defaultdict
+from pathlib import Path
+from types import ModuleType, SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from primus.core.patches import PatchContext
+from primus.core.patches.patch_registry import PatchRegistry
+
+
+def context(**kwargs):
+    return PatchContext(
+        backend="megatron",
+        phase="before_train",
+        extra={"module_config": SimpleNamespace(params=SimpleNamespace(**kwargs))},
+    )
+
+
+@pytest.fixture
+def patched(monkeypatch):
+    monkeypatch.setattr(PatchRegistry, "_all_patches", [])
+    monkeypatch.setattr(PatchRegistry, "_patches_by_backend_phase", defaultdict(lambda: defaultdict(list)))
+    path = Path(__file__).parents[2] / "primus/backends/megatron/patches/turbo/cross_entropy_patches.py"
+    spec = importlib.util.spec_from_file_location("ce_config_under_test", path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    monkeypatch.setattr(mod, "log_rank_0", lambda *args: None)
+    original = Mock(return_value="fallback")
+    original._primus_turbo_ce = False
+    cls = type("LanguageModule", (), {"compute_language_model_loss": original})
+    lm = ModuleType("megatron.core.models.common.language_module.language_module")
+    lm.LanguageModule = cls
+    monkeypatch.setitem(sys.modules, lm.__name__, lm)
+    ce = ModuleType("primus_turbo.pytorch.ops.cross_entropy")
+    ce.cross_entropy = Mock(side_effect=lambda logits, target, **kw: torch.zeros_like(target))
+    monkeypatch.setitem(sys.modules, ce.__name__, ce)
+    return mod, cls, original, ce.cross_entropy
+
+
+def test_explicit_arguments_override_legacy_environment(patched, monkeypatch):
+    mod, *_ = patched
+    monkeypatch.setenv("PRIMUS_TURBO_CROSS_ENTROPY", "1")
+    assert not mod._enabled(context())
+    assert not mod._enabled(context(use_turbo_cross_entropy=False))
+    monkeypatch.setenv("PRIMUS_TURBO_CROSS_ENTROPY", "0")
+    assert mod._enabled(context(use_turbo_cross_entropy=True))
+
+
+@pytest.mark.parametrize("overwrite", [False, True])
+@pytest.mark.parametrize("contiguous", [False, True])
+def test_overwrite_argument_and_layout(patched, monkeypatch, overwrite, contiguous):
+    mod, cls, original, ce = patched
+    monkeypatch.setenv("PRIMUS_TURBO_CE_OVERWRITE_INPUT", str(int(not overwrite)))
+    mod.patch_cross_entropy(context(turbo_ce_overwrite_input=overwrite))
+    installed = cls.compute_language_model_loss
+    mod.patch_cross_entropy(context(turbo_ce_overwrite_input=overwrite))
+    assert cls.compute_language_model_loss is installed
+    logits = torch.empty(7, 3, 11) if contiguous else torch.empty(3, 7, 11).transpose(0, 1)
+    labels = torch.zeros(3, 7, dtype=torch.long)
+    model = SimpleNamespace(
+        config=SimpleNamespace(tensor_model_parallel_size=1, cross_entropy_loss_fusion=True)
+    )
+    assert cls.compute_language_model_loss(model, labels, logits).shape == labels.shape
+    assert ce.call_args.kwargs["overwrite_input"] == (overwrite and contiguous)
+    torch.testing.assert_close(ce.call_args.args[1], labels.T.contiguous())
+    original.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "tp,fused,dtype", [(2, True, torch.float32), (1, False, torch.float32), (1, True, torch.float16)]
+)
+def test_fallback(patched, tp, fused, dtype):
+    mod, cls, original, ce = patched
+    mod.patch_cross_entropy(context())
+    model = SimpleNamespace(
+        config=SimpleNamespace(tensor_model_parallel_size=tp, cross_entropy_loss_fusion=fused)
+    )
+    assert (
+        cls.compute_language_model_loss(model, torch.zeros(3, 7), torch.empty(7, 3, 11, dtype=dtype))
+        == "fallback"
+    )
+    original.assert_called_once()
+    ce.assert_not_called()