From 5da8ae3839602970f226b9c5434128ed9882599e Mon Sep 17 00:00:00 2001
From: Hejian Sang <sanghj0923@gmail.com>
Date: Thu, 22 Oct 2026 13:26:42 -0700
Subject: [PATCH 1/9] Compute log probs without Megatron when it is not
 installed

compute_log_probs, which FSDP and torchtitan share with Megatron, always called
Megatron's fused_vocab_parallel_cross_entropy, so neither backend could train
on an install without Megatron (a pip wheel). When Megatron is absent it now
uses a plain-torch port of the same vocab-parallel math; with Megatron
installed the fused kernel runs exactly as before, and a broken Megatron
install still fails at its import.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../training_utils/loss/hub/math_utils.py     |  54 ++++++++-
 .../test_vocab_parallel_cross_entropy.py      | 111 ++++++++++++++++++
 2 files changed, 163 insertions(+), 2 deletions(-)
 create mode 100644 tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py

diff --git a/miles/backends/training_utils/loss/hub/math_utils.py b/miles/backends/training_utils/loss/hub/math_utils.py
index 8a151440441..9da07ebb98c 100644
--- a/miles/backends/training_utils/loss/hub/math_utils.py
+++ b/miles/backends/training_utils/loss/hub/math_utils.py
@@ -1,6 +1,7 @@
 # Adapt from https://github.com/OpenRLHF/OpenRLHF/blob/10c733694ed9fbb78a0a2ff6a05efc7401584d46/openrlhf/models/utils.py
 # and https://github.com/OpenRLHF/OpenRLHF/blob/10c733694ed9fbb78a0a2ff6a05efc7401584d46/openrlhf/trainer/ppo_utils/experience_maker.py
 
+import importlib.util
 from argparse import Namespace
 from pathlib import Path
 
@@ -17,6 +18,9 @@
 from miles.backends.training_utils.parallel import get_parallel_state
 
 _LOG_RATIO_EXP_CLAMP = 20.0
+# A pip install for the FSDP or torchtitan backend has no Megatron; a broken Megatron install
+# still fails loudly at its import instead of silently switching implementations.
+_MEGATRON_INSTALLED = importlib.util.find_spec("megatron") is not None
 
 
 def _safe_clamp_log_ratio(log_ratio: torch.Tensor) -> torch.Tensor:
@@ -294,8 +298,10 @@ def compute_log_probs(
         log_probs = torch.log_softmax(full_logits, dim=-1)
         return log_probs.gather(dim=-1, index=tokens.unsqueeze(-1)).squeeze(-1)
 
-    # TODO: when megatron is not installed, fall back to naive implementation
-    from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy
+    if _MEGATRON_INSTALLED:
+        from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy
+    else:
+        fused_vocab_parallel_cross_entropy = vocab_parallel_cross_entropy
 
     # convert to [seq_len, batch_size, vocab_size] as expected by fused_vocab_parallel_cross_entropy
     logits = logits.unsqueeze(1)
@@ -311,6 +317,50 @@ def compute_log_probs(
     return -fused_vocab_parallel_cross_entropy(logits, tokens, process_group)
 
 
+class _VocabParallelCrossEntropy(torch.autograd.Function):
+    """Megatron's ``fused_vocab_parallel_cross_entropy`` math in plain torch, for installs without Megatron.
+
+    Logits are ``[seq, batch, vocab / tp]``, split over ``group`` along the vocabulary. Unlike
+    Megatron's, it leaves its input untouched and returns the gradient in the input's dtype.
+    """
+
+    @staticmethod
+    def forward(ctx, logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:
+        ctx.input_dtype = logits.dtype
+        logits = logits.float()
+        logits_max = logits.max(dim=-1)[0]
+        dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=group)
+        logits = logits - logits_max.unsqueeze(-1)
+
+        partition_vocab_size = logits.size(-1)
+        vocab_start = group.rank() * partition_vocab_size
+        target_mask = (target < vocab_start) | (target >= vocab_start + partition_vocab_size)
+        masked_target = (target - vocab_start).masked_fill(target_mask, 0)
+        predicted_logits = logits.gather(-1, masked_target.unsqueeze(-1)).squeeze(-1).masked_fill(target_mask, 0.0)
+        exp_logits = logits.exp()
+        sum_exp_logits = exp_logits.sum(dim=-1)
+
+        # One all-reduce for both, as Megatron batches them.
+        predicted_and_sum = torch.cat((predicted_logits, sum_exp_logits))
+        dist.all_reduce(predicted_and_sum, op=dist.ReduceOp.SUM, group=group)
+        predicted_logits, sum_exp_logits = predicted_and_sum.split(predicted_logits.size(0))
+
+        softmax = exp_logits.div_(sum_exp_logits.unsqueeze(-1))
+        ctx.save_for_backward(softmax, target_mask, masked_target)
+        return torch.log(sum_exp_logits) - predicted_logits
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        softmax, target_mask, masked_target = ctx.saved_tensors
+        # d(loss)/d(logits) = softmax - one_hot(target), the one-hot only on the rank owning the target.
+        grad_input = softmax.scatter_add(-1, masked_target.unsqueeze(-1), -(~target_mask).unsqueeze(-1).float())
+        return (grad_input * grad_output.unsqueeze(-1)).to(ctx.input_dtype), None, None
+
+
+def vocab_parallel_cross_entropy(logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup):
+    return _VocabParallelCrossEntropy.apply(logits, target, group)
+
+
 def _apply_sampling_mask(
     logits: torch.Tensor,
     sampling_mask: torch.Tensor | None,
diff --git a/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
new file mode 100644
index 00000000000..10b6f9b1160
--- /dev/null
+++ b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
@@ -0,0 +1,111 @@
+"""The torch-only vocab-parallel cross entropy that stands in for Megatron's fused one when it is absent."""
+
+import sys
+
+import pytest
+import torch
+import torch.distributed as dist
+from tests.fast.dist_utils import init_gloo, run_multiprocess
+
+from miles.backends.training_utils.loss.hub import math_utils
+from miles.backends.training_utils.loss.hub.math_utils import compute_log_probs, vocab_parallel_cross_entropy
+
+SEQ, BATCH, VOCAB = 5, 2, 12
+
+
+def _full_logits(dtype: torch.dtype) -> torch.Tensor:
+    return (torch.randn(SEQ, BATCH, VOCAB, generator=torch.Generator().manual_seed(7)) * 4).to(dtype)
+
+
+def _target() -> torch.Tensor:
+    # Covers both shards of a TP=2 split, including each shard's first and last id.
+    return torch.tensor([[0, 11], [5, 6], [6, 5], [11, 0], [3, 9]])
+
+
+def _reference(full: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    logits = full.detach().clone().requires_grad_()
+    loss = -torch.log_softmax(logits.float(), dim=-1).gather(-1, target.unsqueeze(-1)).squeeze(-1)
+    loss.backward(_grad_output())
+    return loss.detach(), logits.grad
+
+
+def _grad_output() -> torch.Tensor:
+    return torch.randn(SEQ, BATCH, generator=torch.Generator().manual_seed(11))
+
+
+def _check_against_full_softmax(rank: int, world_size: int, port: int) -> None:
+    init_gloo(rank, world_size, port=port)
+    try:
+        group = dist.new_group(list(range(world_size)))
+        shard = VOCAB // world_size
+        for dtype, tol in ((torch.float32, 1e-6), (torch.bfloat16, 1e-2)):
+            full = _full_logits(dtype)
+            expected_loss, expected_grad = _reference(full, _target())
+            local = full[..., rank * shard : (rank + 1) * shard].clone().requires_grad_()
+            local_before = local.detach().clone()
+
+            loss = vocab_parallel_cross_entropy(local, _target(), group)
+            loss.backward(_grad_output())
+
+            torch.testing.assert_close(loss, expected_loss, atol=tol, rtol=tol)
+            assert local.grad.dtype == dtype
+            torch.testing.assert_close(
+                local.grad, expected_grad[..., rank * shard : (rank + 1) * shard], atol=tol, rtol=tol
+            )
+            torch.testing.assert_close(local.detach(), local_before, atol=0, rtol=0)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("world_size", [1, 2])
+def test_matches_the_full_vocabulary_log_softmax_on_every_shard(world_size: int) -> None:
+    """Loss and gradient equal a plain log-softmax over the whole vocabulary, however it is split."""
+    run_multiprocess(_check_against_full_softmax, world_size=world_size)
+
+
+def _check_against_megatron(rank: int, world_size: int, port: int) -> None:
+    from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy
+
+    init_gloo(rank, world_size, port=port)
+    try:
+        group = dist.new_group(list(range(world_size)))
+        shard = VOCAB // world_size
+        full = _full_logits(torch.bfloat16)
+        results = []
+        for fn in (vocab_parallel_cross_entropy, fused_vocab_parallel_cross_entropy):
+            local = full[..., rank * shard : (rank + 1) * shard].clone().requires_grad_()
+            loss = fn(local, _target(), group)
+            loss.backward(_grad_output())
+            results.append((loss.detach(), local.grad))
+        (loss, grad), (megatron_loss, megatron_grad) = results
+        torch.testing.assert_close(loss, megatron_loss, atol=1e-6, rtol=1e-6)
+        torch.testing.assert_close(grad, megatron_grad, atol=0, rtol=0)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("world_size", [1, 2])
+def test_matches_megatron_fused_cross_entropy(world_size: int) -> None:
+    """Where Megatron is installed, the stand-in computes what its fused kernel computes."""
+    pytest.importorskip("megatron.core.fusions.fused_cross_entropy")
+    run_multiprocess(_check_against_megatron, world_size=world_size)
+
+
+def _check_compute_log_probs_without_megatron(rank: int, world_size: int, port: int) -> None:
+    sys.modules["megatron"] = None  # importing any of it now raises, as on an install without it
+    math_utils._MEGATRON_INSTALLED = False
+    init_gloo(rank, world_size, port=port)
+    try:
+        group = dist.new_group([0])
+        logits = _full_logits(torch.float32)[:, 0]
+        tokens = _target()[:, 0]
+        expected = torch.log_softmax(logits, dim=-1).gather(-1, tokens.unsqueeze(-1)).squeeze(-1)
+
+        torch.testing.assert_close(compute_log_probs(logits, tokens, group), expected, atol=1e-6, rtol=1e-6)
+    finally:
+        dist.destroy_process_group()
+
+
+def test_compute_log_probs_runs_without_megatron() -> None:
+    """The FSDP and torchtitan log-prob path imports nothing from Megatron when it is not installed."""
+    run_multiprocess(_check_compute_log_probs_without_megatron, world_size=1)

From 137c7871d62716cd2f7f4d4c63619a35280a3415 Mon Sep 17 00:00:00 2001
From: Hejian Sang <sanghj0923@gmail.com>
Date: Thu, 22 Oct 2026 13:29:51 -0700
Subject: [PATCH 2/9] test: compute_log_probs returns [seq, 1]

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../loss_hub/test_vocab_parallel_cross_entropy.py              | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
index 10b6f9b1160..6cd31625125 100644
--- a/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
+++ b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
@@ -99,7 +99,8 @@ def _check_compute_log_probs_without_megatron(rank: int, world_size: int, port:
         group = dist.new_group([0])
         logits = _full_logits(torch.float32)[:, 0]
         tokens = _target()[:, 0]
-        expected = torch.log_softmax(logits, dim=-1).gather(-1, tokens.unsqueeze(-1)).squeeze(-1)
+        # [seq, 1], the shape compute_log_probs has always returned
+        expected = torch.log_softmax(logits, dim=-1).gather(-1, tokens.unsqueeze(-1))
 
         torch.testing.assert_close(compute_log_probs(logits, tokens, group), expected, atol=1e-6, rtol=1e-6)
     finally:

From 804331e9647d783e47e525cb82eae07b9c165f1a Mon Sep 17 00:00:00 2001
From: Hejian Sang <sanghj0923@gmail.com>
Date: Thu, 22 Oct 2026 13:31:09 -0700
Subject: [PATCH 3/9] test: share one process group per tensor-parallel size

Spawning dominates the runtime; two spawns keep the file inside the fast CPU budget.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../test_vocab_parallel_cross_entropy.py      | 111 ++++++++----------
 1 file changed, 50 insertions(+), 61 deletions(-)

diff --git a/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
index 6cd31625125..4664180b651 100644
--- a/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
+++ b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
@@ -33,80 +33,69 @@ def _grad_output() -> torch.Tensor:
     return torch.randn(SEQ, BATCH, generator=torch.Generator().manual_seed(11))
 
 
-def _check_against_full_softmax(rank: int, world_size: int, port: int) -> None:
-    init_gloo(rank, world_size, port=port)
-    try:
-        group = dist.new_group(list(range(world_size)))
-        shard = VOCAB // world_size
-        for dtype, tol in ((torch.float32, 1e-6), (torch.bfloat16, 1e-2)):
-            full = _full_logits(dtype)
-            expected_loss, expected_grad = _reference(full, _target())
-            local = full[..., rank * shard : (rank + 1) * shard].clone().requires_grad_()
-            local_before = local.detach().clone()
-
-            loss = vocab_parallel_cross_entropy(local, _target(), group)
-            loss.backward(_grad_output())
-
-            torch.testing.assert_close(loss, expected_loss, atol=tol, rtol=tol)
-            assert local.grad.dtype == dtype
-            torch.testing.assert_close(
-                local.grad, expected_grad[..., rank * shard : (rank + 1) * shard], atol=tol, rtol=tol
-            )
-            torch.testing.assert_close(local.detach(), local_before, atol=0, rtol=0)
-    finally:
-        dist.destroy_process_group()
-
-
-@pytest.mark.parametrize("world_size", [1, 2])
-def test_matches_the_full_vocabulary_log_softmax_on_every_shard(world_size: int) -> None:
+def _check_against_full_softmax(rank: int, world_size: int, group: dist.ProcessGroup) -> None:
     """Loss and gradient equal a plain log-softmax over the whole vocabulary, however it is split."""
-    run_multiprocess(_check_against_full_softmax, world_size=world_size)
-
+    shard = VOCAB // world_size
+    for dtype, tol in ((torch.float32, 1e-6), (torch.bfloat16, 1e-2)):
+        full = _full_logits(dtype)
+        expected_loss, expected_grad = _reference(full, _target())
+        local = full[..., rank * shard : (rank + 1) * shard].clone().requires_grad_()
+        local_before = local.detach().clone()
 
-def _check_against_megatron(rank: int, world_size: int, port: int) -> None:
-    from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy
+        loss = vocab_parallel_cross_entropy(local, _target(), group)
+        loss.backward(_grad_output())
 
-    init_gloo(rank, world_size, port=port)
-    try:
-        group = dist.new_group(list(range(world_size)))
-        shard = VOCAB // world_size
-        full = _full_logits(torch.bfloat16)
-        results = []
-        for fn in (vocab_parallel_cross_entropy, fused_vocab_parallel_cross_entropy):
-            local = full[..., rank * shard : (rank + 1) * shard].clone().requires_grad_()
-            loss = fn(local, _target(), group)
-            loss.backward(_grad_output())
-            results.append((loss.detach(), local.grad))
-        (loss, grad), (megatron_loss, megatron_grad) = results
-        torch.testing.assert_close(loss, megatron_loss, atol=1e-6, rtol=1e-6)
-        torch.testing.assert_close(grad, megatron_grad, atol=0, rtol=0)
-    finally:
-        dist.destroy_process_group()
+        torch.testing.assert_close(loss, expected_loss, atol=tol, rtol=tol)
+        assert local.grad.dtype == dtype
+        torch.testing.assert_close(
+            local.grad, expected_grad[..., rank * shard : (rank + 1) * shard], atol=tol, rtol=tol
+        )
+        torch.testing.assert_close(local.detach(), local_before, atol=0, rtol=0)
 
 
-@pytest.mark.parametrize("world_size", [1, 2])
-def test_matches_megatron_fused_cross_entropy(world_size: int) -> None:
+def _check_against_megatron(rank: int, world_size: int, group: dist.ProcessGroup) -> None:
     """Where Megatron is installed, the stand-in computes what its fused kernel computes."""
-    pytest.importorskip("megatron.core.fusions.fused_cross_entropy")
-    run_multiprocess(_check_against_megatron, world_size=world_size)
+    from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy
 
+    shard = VOCAB // world_size
+    full = _full_logits(torch.bfloat16)
+    results = []
+    for fn in (vocab_parallel_cross_entropy, fused_vocab_parallel_cross_entropy):
+        local = full[..., rank * shard : (rank + 1) * shard].clone().requires_grad_()
+        loss = fn(local, _target(), group)
+        loss.backward(_grad_output())
+        results.append((loss.detach(), local.grad))
+    (loss, grad), (megatron_loss, megatron_grad) = results
+    torch.testing.assert_close(loss, megatron_loss, atol=1e-6, rtol=1e-6)
+    torch.testing.assert_close(grad, megatron_grad, atol=0, rtol=0)
 
-def _check_compute_log_probs_without_megatron(rank: int, world_size: int, port: int) -> None:
+
+def _check_compute_log_probs_without_megatron(group: dist.ProcessGroup) -> None:
+    """The FSDP and torchtitan log-prob path imports nothing from Megatron when it is not installed."""
     sys.modules["megatron"] = None  # importing any of it now raises, as on an install without it
     math_utils._MEGATRON_INSTALLED = False
+    logits = _full_logits(torch.float32)[:, 0]
+    tokens = _target()[:, 0]
+    # [seq, 1], the shape compute_log_probs has always returned
+    expected = torch.log_softmax(logits, dim=-1).gather(-1, tokens.unsqueeze(-1))
+
+    torch.testing.assert_close(compute_log_probs(logits, tokens, group), expected, atol=1e-6, rtol=1e-6)
+
+
+def _run_checks(rank: int, world_size: int, port: int) -> None:
     init_gloo(rank, world_size, port=port)
     try:
-        group = dist.new_group([0])
-        logits = _full_logits(torch.float32)[:, 0]
-        tokens = _target()[:, 0]
-        # [seq, 1], the shape compute_log_probs has always returned
-        expected = torch.log_softmax(logits, dim=-1).gather(-1, tokens.unsqueeze(-1))
-
-        torch.testing.assert_close(compute_log_probs(logits, tokens, group), expected, atol=1e-6, rtol=1e-6)
+        group = dist.new_group(list(range(world_size)))
+        _check_against_full_softmax(rank, world_size, group)
+        _check_against_megatron(rank, world_size, group)
+        if world_size == 1:
+            # last: it hides Megatron from the rest of this process
+            _check_compute_log_probs_without_megatron(group)
     finally:
         dist.destroy_process_group()
 
 
-def test_compute_log_probs_runs_without_megatron() -> None:
-    """The FSDP and torchtitan log-prob path imports nothing from Megatron when it is not installed."""
-    run_multiprocess(_check_compute_log_probs_without_megatron, world_size=1)
+@pytest.mark.parametrize("world_size", [1, 2])
+def test_vocab_parallel_cross_entropy(world_size: int) -> None:
+    """One process group per tensor-parallel size: spawning is what costs, so every check shares it."""
+    run_multiprocess(_run_checks, world_size=world_size)

From c5951460d175571c862d3e04de70b28749f3706b Mon Sep 17 00:00:00 2001
From: Hejian Sang <sanghj0923@gmail.com>
Date: Thu, 22 Oct 2026 14:07:17 -0700
Subject: [PATCH 4/9] Check megatron.core, and let the fallback take the
 default group as None

megatron-energon installs a bare megatron namespace without megatron.core,
which find_spec("megatron") counted as Megatron. The fallback also took
group.rank() on a None group, which compute_log_probs callers pass for the
default group; dist.get_rank(None) gives torch.distributed's meaning.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../training_utils/loss/hub/math_utils.py     | 25 +++++++++++++------
 .../test_vocab_parallel_cross_entropy.py      | 24 ++++++++++++++++++
 2 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/miles/backends/training_utils/loss/hub/math_utils.py b/miles/backends/training_utils/loss/hub/math_utils.py
index 9da07ebb98c..1a9372b9f2f 100644
--- a/miles/backends/training_utils/loss/hub/math_utils.py
+++ b/miles/backends/training_utils/loss/hub/math_utils.py
@@ -18,9 +18,19 @@
 from miles.backends.training_utils.parallel import get_parallel_state
 
 _LOG_RATIO_EXP_CLAMP = 20.0
-# A pip install for the FSDP or torchtitan backend has no Megatron; a broken Megatron install
-# still fails loudly at its import instead of silently switching implementations.
-_MEGATRON_INSTALLED = importlib.util.find_spec("megatron") is not None
+
+
+def _megatron_core_installed() -> bool:
+    """A pip install for the FSDP or torchtitan backend has no Megatron. ``megatron.core`` holds the fused
+    kernel: megatron-energon alone installs a bare ``megatron`` namespace. A broken Megatron install still
+    fails loudly at its import instead of silently switching implementations."""
+    try:
+        return importlib.util.find_spec("megatron.core") is not None
+    except ModuleNotFoundError:  # no megatron package at all
+        return False
+
+
+_MEGATRON_INSTALLED = _megatron_core_installed()
 
 
 def _safe_clamp_log_ratio(log_ratio: torch.Tensor) -> torch.Tensor:
@@ -320,12 +330,13 @@ def compute_log_probs(
 class _VocabParallelCrossEntropy(torch.autograd.Function):
     """Megatron's ``fused_vocab_parallel_cross_entropy`` math in plain torch, for installs without Megatron.
 
-    Logits are ``[seq, batch, vocab / tp]``, split over ``group`` along the vocabulary. Unlike
+    Logits are ``[seq, batch, vocab / tp]``, split over ``group`` (None: the default group, as in
+    ``torch.distributed``) along the vocabulary. Unlike
     Megatron's, it leaves its input untouched and returns the gradient in the input's dtype.
     """
 
     @staticmethod
-    def forward(ctx, logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:
+    def forward(ctx, logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup | None) -> torch.Tensor:
         ctx.input_dtype = logits.dtype
         logits = logits.float()
         logits_max = logits.max(dim=-1)[0]
@@ -333,7 +344,7 @@ def forward(ctx, logits: torch.Tensor, target: torch.Tensor, group: dist.Process
         logits = logits - logits_max.unsqueeze(-1)
 
         partition_vocab_size = logits.size(-1)
-        vocab_start = group.rank() * partition_vocab_size
+        vocab_start = dist.get_rank(group) * partition_vocab_size
         target_mask = (target < vocab_start) | (target >= vocab_start + partition_vocab_size)
         masked_target = (target - vocab_start).masked_fill(target_mask, 0)
         predicted_logits = logits.gather(-1, masked_target.unsqueeze(-1)).squeeze(-1).masked_fill(target_mask, 0.0)
@@ -357,7 +368,7 @@ def backward(ctx, grad_output: torch.Tensor):
         return (grad_input * grad_output.unsqueeze(-1)).to(ctx.input_dtype), None, None
 
 
-def vocab_parallel_cross_entropy(logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup):
+def vocab_parallel_cross_entropy(logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup | None):
     return _VocabParallelCrossEntropy.apply(logits, target, group)
 
 
diff --git a/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
index 4664180b651..cadbe4f40ad 100644
--- a/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
+++ b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
@@ -80,6 +80,8 @@ def _check_compute_log_probs_without_megatron(group: dist.ProcessGroup) -> None:
     expected = torch.log_softmax(logits, dim=-1).gather(-1, tokens.unsqueeze(-1))
 
     torch.testing.assert_close(compute_log_probs(logits, tokens, group), expected, atol=1e-6, rtol=1e-6)
+    # None is the default group, as in torch.distributed and the callers that pass it
+    torch.testing.assert_close(compute_log_probs(logits, tokens, None), expected, atol=1e-6, rtol=1e-6)
 
 
 def _run_checks(rank: int, world_size: int, port: int) -> None:
@@ -99,3 +101,25 @@ def _run_checks(rank: int, world_size: int, port: int) -> None:
 def test_vocab_parallel_cross_entropy(world_size: int) -> None:
     """One process group per tensor-parallel size: spawning is what costs, so every check shares it."""
     run_multiprocess(_run_checks, world_size=world_size)
+
+
+def _find_spec_seeing(found: dict[str, bool]):
+    def find_spec(name: str):
+        if name.split(".")[0] not in found:  # importing the parent package fails, as find_spec does
+            raise ModuleNotFoundError(name)
+        return object() if found.get(name, False) else None
+
+    return find_spec
+
+
+@pytest.mark.parametrize(
+    "found,installed",
+    [
+        ({"megatron": True, "megatron.core": True}, True),
+        ({"megatron": True}, False),  # megatron-energon's bare namespace, no megatron.core
+        ({}, False),  # no megatron package at all
+    ],
+)
+def test_megatron_counts_as_installed_only_with_megatron_core(monkeypatch, found, installed) -> None:
+    monkeypatch.setattr(math_utils.importlib.util, "find_spec", _find_spec_seeing(found))
+    assert math_utils._megatron_core_installed() is installed

From cd5e45ab837856a0e98f9a73e371565dd0299246 Mon Sep 17 00:00:00 2001
From: Hejian Sang <sanghj0923@gmail.com>
Date: Thu, 22 Oct 2026 14:20:39 -0700
Subject: [PATCH 5/9] Fallback cross entropy: one copy, then in place as
 Megatron's kernel

The shift, exp and gradient each allocated another full-size fp32 logits
buffer, so a run without Megatron could OOM where the fused kernel fits.
One copy keeps the caller's logits untouched; the rest works in place.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 miles/backends/training_utils/loss/hub/math_utils.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/miles/backends/training_utils/loss/hub/math_utils.py b/miles/backends/training_utils/loss/hub/math_utils.py
index 1a9372b9f2f..6d905e25ea5 100644
--- a/miles/backends/training_utils/loss/hub/math_utils.py
+++ b/miles/backends/training_utils/loss/hub/math_utils.py
@@ -338,17 +338,19 @@ class _VocabParallelCrossEntropy(torch.autograd.Function):
     @staticmethod
     def forward(ctx, logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup | None) -> torch.Tensor:
         ctx.input_dtype = logits.dtype
-        logits = logits.float()
+        # The one full-size copy: it keeps the caller's logits untouched, and the rest works on it in
+        # place as Megatron's kernel does, so the fallback needs no more memory than that kernel.
+        logits = logits.to(torch.float32, copy=True)
         logits_max = logits.max(dim=-1)[0]
         dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=group)
-        logits = logits - logits_max.unsqueeze(-1)
+        logits.sub_(logits_max.unsqueeze(-1))
 
         partition_vocab_size = logits.size(-1)
         vocab_start = dist.get_rank(group) * partition_vocab_size
         target_mask = (target < vocab_start) | (target >= vocab_start + partition_vocab_size)
         masked_target = (target - vocab_start).masked_fill(target_mask, 0)
         predicted_logits = logits.gather(-1, masked_target.unsqueeze(-1)).squeeze(-1).masked_fill(target_mask, 0.0)
-        exp_logits = logits.exp()
+        exp_logits = logits.exp_()
         sum_exp_logits = exp_logits.sum(dim=-1)
 
         # One all-reduce for both, as Megatron batches them.
@@ -364,8 +366,8 @@ def forward(ctx, logits: torch.Tensor, target: torch.Tensor, group: dist.Process
     def backward(ctx, grad_output: torch.Tensor):
         softmax, target_mask, masked_target = ctx.saved_tensors
         # d(loss)/d(logits) = softmax - one_hot(target), the one-hot only on the rank owning the target.
-        grad_input = softmax.scatter_add(-1, masked_target.unsqueeze(-1), -(~target_mask).unsqueeze(-1).float())
-        return (grad_input * grad_output.unsqueeze(-1)).to(ctx.input_dtype), None, None
+        grad_input = softmax.scatter_add_(-1, masked_target.unsqueeze(-1), -(~target_mask).unsqueeze(-1).float())
+        return grad_input.mul_(grad_output.unsqueeze(-1)).to(ctx.input_dtype), None, None
 
 
 def vocab_parallel_cross_entropy(logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup | None):

From e305343ba5630a6b0742b51ace93806f478c1bad Mon Sep 17 00:00:00 2001
From: Hejian Sang <sanghj0923@gmail.com>
Date: Thu, 22 Oct 2026 15:20:15 -0700
Subject: [PATCH 6/9] Stop calling the plain-torch cross entropy "fused" in
 compute_log_probs

Without Megatron the local name fused_vocab_parallel_cross_entropy held the
plain-torch fallback; both branches now bind cross_entropy.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 miles/backends/training_utils/loss/hub/math_utils.py | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/miles/backends/training_utils/loss/hub/math_utils.py b/miles/backends/training_utils/loss/hub/math_utils.py
index 6d905e25ea5..1f1bf61fbe3 100644
--- a/miles/backends/training_utils/loss/hub/math_utils.py
+++ b/miles/backends/training_utils/loss/hub/math_utils.py
@@ -309,22 +309,20 @@ def compute_log_probs(
         return log_probs.gather(dim=-1, index=tokens.unsqueeze(-1)).squeeze(-1)
 
     if _MEGATRON_INSTALLED:
-        from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy
+        from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy as cross_entropy
     else:
-        fused_vocab_parallel_cross_entropy = vocab_parallel_cross_entropy
+        cross_entropy = vocab_parallel_cross_entropy
 
-    # convert to [seq_len, batch_size, vocab_size] as expected by fused_vocab_parallel_cross_entropy
+    # convert to [seq_len, batch_size, vocab_size] as cross_entropy expects
     logits = logits.unsqueeze(1)
     tokens = tokens.unsqueeze(1)
 
     if debug_unified_grad_fused_logprob and not torch.is_grad_enabled():
         with torch.enable_grad():
-            log_probs = -fused_vocab_parallel_cross_entropy(
-                logits.detach().requires_grad_(True), tokens, process_group
-            )
+            log_probs = -cross_entropy(logits.detach().requires_grad_(True), tokens, process_group)
         return log_probs.detach()
 
-    return -fused_vocab_parallel_cross_entropy(logits, tokens, process_group)
+    return -cross_entropy(logits, tokens, process_group)
 
 
 class _VocabParallelCrossEntropy(torch.autograd.Function):

From 752a065842cf879df1af2de909978d06952f667c Mon Sep 17 00:00:00 2001
From: Hejian Sang <sanghj0923@gmail.com>
Date: Thu, 22 Oct 2026 15:20:17 -0700
Subject: [PATCH 7/9] Trim the comments in the Megatron-free cross entropy

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../training_utils/loss/hub/math_utils.py      | 18 +++++++-----------
 .../test_vocab_parallel_cross_entropy.py       | 10 +++++-----
 2 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/miles/backends/training_utils/loss/hub/math_utils.py b/miles/backends/training_utils/loss/hub/math_utils.py
index 1f1bf61fbe3..cb3b4c3f5ca 100644
--- a/miles/backends/training_utils/loss/hub/math_utils.py
+++ b/miles/backends/training_utils/loss/hub/math_utils.py
@@ -21,9 +21,7 @@
 
 
 def _megatron_core_installed() -> bool:
-    """A pip install for the FSDP or torchtitan backend has no Megatron. ``megatron.core`` holds the fused
-    kernel: megatron-energon alone installs a bare ``megatron`` namespace. A broken Megatron install still
-    fails loudly at its import instead of silently switching implementations."""
+    """``megatron.core`` holds the fused kernel; megatron-energon alone installs a bare ``megatron``."""
     try:
         return importlib.util.find_spec("megatron.core") is not None
     except ModuleNotFoundError:  # no megatron package at all
@@ -326,18 +324,16 @@ def compute_log_probs(
 
 
 class _VocabParallelCrossEntropy(torch.autograd.Function):
-    """Megatron's ``fused_vocab_parallel_cross_entropy`` math in plain torch, for installs without Megatron.
+    """Megatron's ``fused_vocab_parallel_cross_entropy`` in plain torch, for installs without Megatron.
 
-    Logits are ``[seq, batch, vocab / tp]``, split over ``group`` (None: the default group, as in
-    ``torch.distributed``) along the vocabulary. Unlike
-    Megatron's, it leaves its input untouched and returns the gradient in the input's dtype.
+    Logits are ``[seq, batch, vocab / tp]``, split over ``group`` along the vocabulary. Unlike Megatron's,
+    it leaves its input untouched and returns the gradient in the input's dtype.
     """
 
     @staticmethod
     def forward(ctx, logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup | None) -> torch.Tensor:
         ctx.input_dtype = logits.dtype
-        # The one full-size copy: it keeps the caller's logits untouched, and the rest works on it in
-        # place as Megatron's kernel does, so the fallback needs no more memory than that kernel.
+        # the one full-size copy; the rest is in place, as in Megatron
         logits = logits.to(torch.float32, copy=True)
         logits_max = logits.max(dim=-1)[0]
         dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=group)
@@ -351,7 +347,7 @@ def forward(ctx, logits: torch.Tensor, target: torch.Tensor, group: dist.Process
         exp_logits = logits.exp_()
         sum_exp_logits = exp_logits.sum(dim=-1)
 
-        # One all-reduce for both, as Megatron batches them.
+        # one all-reduce for both, as in Megatron
         predicted_and_sum = torch.cat((predicted_logits, sum_exp_logits))
         dist.all_reduce(predicted_and_sum, op=dist.ReduceOp.SUM, group=group)
         predicted_logits, sum_exp_logits = predicted_and_sum.split(predicted_logits.size(0))
@@ -363,7 +359,7 @@ def forward(ctx, logits: torch.Tensor, target: torch.Tensor, group: dist.Process
     @staticmethod
     def backward(ctx, grad_output: torch.Tensor):
         softmax, target_mask, masked_target = ctx.saved_tensors
-        # d(loss)/d(logits) = softmax - one_hot(target), the one-hot only on the rank owning the target.
+        # softmax - one_hot(target); the one-hot lives on the rank owning the target
         grad_input = softmax.scatter_add_(-1, masked_target.unsqueeze(-1), -(~target_mask).unsqueeze(-1).float())
         return grad_input.mul_(grad_output.unsqueeze(-1)).to(ctx.input_dtype), None, None
 
diff --git a/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
index cadbe4f40ad..f41a9dc2853 100644
--- a/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
+++ b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
@@ -1,4 +1,4 @@
-"""The torch-only vocab-parallel cross entropy that stands in for Megatron's fused one when it is absent."""
+"""The plain-torch vocab-parallel cross entropy must compute what Megatron's fused kernel does."""
 
 import sys
 
@@ -18,7 +18,7 @@ def _full_logits(dtype: torch.dtype) -> torch.Tensor:
 
 
 def _target() -> torch.Tensor:
-    # Covers both shards of a TP=2 split, including each shard's first and last id.
+    # each TP=2 shard's first and last id
     return torch.tensor([[0, 11], [5, 6], [6, 5], [11, 0], [3, 9]])
 
 
@@ -76,11 +76,11 @@ def _check_compute_log_probs_without_megatron(group: dist.ProcessGroup) -> None:
     math_utils._MEGATRON_INSTALLED = False
     logits = _full_logits(torch.float32)[:, 0]
     tokens = _target()[:, 0]
-    # [seq, 1], the shape compute_log_probs has always returned
+    # compute_log_probs returns [seq, 1]
     expected = torch.log_softmax(logits, dim=-1).gather(-1, tokens.unsqueeze(-1))
 
     torch.testing.assert_close(compute_log_probs(logits, tokens, group), expected, atol=1e-6, rtol=1e-6)
-    # None is the default group, as in torch.distributed and the callers that pass it
+    # None: the default group
     torch.testing.assert_close(compute_log_probs(logits, tokens, None), expected, atol=1e-6, rtol=1e-6)
 
 
@@ -99,7 +99,7 @@ def _run_checks(rank: int, world_size: int, port: int) -> None:
 
 @pytest.mark.parametrize("world_size", [1, 2])
 def test_vocab_parallel_cross_entropy(world_size: int) -> None:
-    """One process group per tensor-parallel size: spawning is what costs, so every check shares it."""
+    """Loss and gradient match a full-vocabulary log-softmax and Megatron's kernel on every TP shard."""
     run_multiprocess(_run_checks, world_size=world_size)
 
 

From 46c51b21f0ae685f48796b02f95e49b55527409e Mon Sep 17 00:00:00 2001
From: Hejian Sang <sanghj0923@gmail.com>
Date: Fri, 23 Oct 2026 13:57:37 -0700
Subject: [PATCH 8/9] Use the torch vocab-parallel cross entropy on every
 install

One implementation instead of a switch on whether megatron.core is installed,
so the log-probs no longer depend on what happens to be importable: with
Megatron installed, the FSDP and torchtitan backends ran Megatron's fused kernel,
and on a pip install the torch one. The torch version is bit-identical to
Megatron's unfused vocab-parallel cross entropy; its fused kernel differs in the
last bits (torch.compile's exp and summation order, up to 1.9e-6 in the loss on
H200) and rounds the gradient to bf16. On H200 at 16k x 151936 the torch
version is about 9% faster than the fused kernel with the same peak memory.

The forward works in place on its fp32 input, as Megatron's does: every caller
passes a fresh copy (_upcast_chunk_to_fp32), so a second copy would add an fp32
[T, V] buffer to every Megatron run. The fake-kernel tests patch the miles
function, and the Megatron comparison is its own test, skipped without Megatron.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../training_utils/loss/hub/math_utils.py     | 33 ++------
 .../training_utils/loss_hub/conftest.py       |  4 +-
 .../loss_hub/test_math_utils_log_probs.py     |  8 +-
 .../test_vocab_parallel_cross_entropy.py      | 82 ++++++++-----------
 4 files changed, 46 insertions(+), 81 deletions(-)

diff --git a/miles/backends/training_utils/loss/hub/math_utils.py b/miles/backends/training_utils/loss/hub/math_utils.py
index cb3b4c3f5ca..2ca448ec521 100644
--- a/miles/backends/training_utils/loss/hub/math_utils.py
+++ b/miles/backends/training_utils/loss/hub/math_utils.py
@@ -1,7 +1,6 @@
 # Adapt from https://github.com/OpenRLHF/OpenRLHF/blob/10c733694ed9fbb78a0a2ff6a05efc7401584d46/openrlhf/models/utils.py
 # and https://github.com/OpenRLHF/OpenRLHF/blob/10c733694ed9fbb78a0a2ff6a05efc7401584d46/openrlhf/trainer/ppo_utils/experience_maker.py
 
-import importlib.util
 from argparse import Namespace
 from pathlib import Path
 
@@ -20,17 +19,6 @@
 _LOG_RATIO_EXP_CLAMP = 20.0
 
 
-def _megatron_core_installed() -> bool:
-    """``megatron.core`` holds the fused kernel; megatron-energon alone installs a bare ``megatron``."""
-    try:
-        return importlib.util.find_spec("megatron.core") is not None
-    except ModuleNotFoundError:  # no megatron package at all
-        return False
-
-
-_MEGATRON_INSTALLED = _megatron_core_installed()
-
-
 def _safe_clamp_log_ratio(log_ratio: torch.Tensor) -> torch.Tensor:
     log_ratio = torch.nan_to_num(
         log_ratio.float(),
@@ -306,35 +294,30 @@ def compute_log_probs(
         log_probs = torch.log_softmax(full_logits, dim=-1)
         return log_probs.gather(dim=-1, index=tokens.unsqueeze(-1)).squeeze(-1)
 
-    if _MEGATRON_INSTALLED:
-        from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy as cross_entropy
-    else:
-        cross_entropy = vocab_parallel_cross_entropy
-
-    # convert to [seq_len, batch_size, vocab_size] as cross_entropy expects
+    # [seq, 1]: the shape callers have always received
     logits = logits.unsqueeze(1)
     tokens = tokens.unsqueeze(1)
 
     if debug_unified_grad_fused_logprob and not torch.is_grad_enabled():
         with torch.enable_grad():
-            log_probs = -cross_entropy(logits.detach().requires_grad_(True), tokens, process_group)
+            log_probs = -vocab_parallel_cross_entropy(logits.detach().requires_grad_(True), tokens, process_group)
         return log_probs.detach()
 
-    return -cross_entropy(logits, tokens, process_group)
+    return -vocab_parallel_cross_entropy(logits, tokens, process_group)
 
 
 class _VocabParallelCrossEntropy(torch.autograd.Function):
-    """Megatron's ``fused_vocab_parallel_cross_entropy`` in plain torch, for installs without Megatron.
+    """Cross entropy over logits split along the vocabulary over ``group``, in plain torch.
 
-    Logits are ``[seq, batch, vocab / tp]``, split over ``group`` along the vocabulary. Unlike Megatron's,
-    it leaves its input untouched and returns the gradient in the input's dtype.
+    Megatron's vocab-parallel cross entropy, bit for bit, so every backend and install runs one
+    implementation. Like it, the forward works in place on fp32 ``logits`` (callers pass a copy).
+    The gradient comes back in the input's dtype.
     """
 
     @staticmethod
     def forward(ctx, logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup | None) -> torch.Tensor:
         ctx.input_dtype = logits.dtype
-        # the one full-size copy; the rest is in place, as in Megatron
-        logits = logits.to(torch.float32, copy=True)
+        logits = logits.float()
         logits_max = logits.max(dim=-1)[0]
         dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=group)
         logits.sub_(logits_max.unsqueeze(-1))
diff --git a/tests/fast/backends/training_utils/loss_hub/conftest.py b/tests/fast/backends/training_utils/loss_hub/conftest.py
index 760eccd6572..883fd63cb7c 100644
--- a/tests/fast/backends/training_utils/loss_hub/conftest.py
+++ b/tests/fast/backends/training_utils/loss_hub/conftest.py
@@ -2,7 +2,7 @@
 
 import pytest
 import torch
-from megatron.core.fusions import fused_cross_entropy
+from miles.backends.training_utils.loss.hub import math_utils
 
 
 @pytest.fixture
@@ -21,5 +21,5 @@ def fake_fused_cross_entropy(
         return logits.sum(dim=-1) * 0
 
     fake: Callable[[torch.Tensor, torch.Tensor, object], torch.Tensor] = fake_fused_cross_entropy
-    monkeypatch.setattr(fused_cross_entropy, "fused_vocab_parallel_cross_entropy", fake)
+    monkeypatch.setattr(math_utils, "vocab_parallel_cross_entropy", fake)
     return calls
diff --git a/tests/fast/backends/training_utils/loss_hub/test_math_utils_log_probs.py b/tests/fast/backends/training_utils/loss_hub/test_math_utils_log_probs.py
index 351d7d21fb2..1b71e643630 100644
--- a/tests/fast/backends/training_utils/loss_hub/test_math_utils_log_probs.py
+++ b/tests/fast/backends/training_utils/loss_hub/test_math_utils_log_probs.py
@@ -2,7 +2,7 @@
 
 import pytest
 import torch
-from megatron.core.fusions import fused_cross_entropy
+from miles.backends.training_utils.loss.hub import math_utils
 
 from miles.backends.training_utils.loss.hub.math_utils import compute_log_probs
 
@@ -18,11 +18,7 @@ def fake_fused_cross_entropy(
         calls.append((torch.is_grad_enabled(), logits.requires_grad))
         return logits.sum(dim=-1) * 0
 
-    monkeypatch.setattr(
-        fused_cross_entropy,
-        "fused_vocab_parallel_cross_entropy",
-        fake_fused_cross_entropy,
-    )
+    monkeypatch.setattr(math_utils, "vocab_parallel_cross_entropy", fake_fused_cross_entropy)
     return calls
 
 
diff --git a/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
index f41a9dc2853..6deb93302f6 100644
--- a/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
+++ b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
@@ -1,13 +1,10 @@
-"""The plain-torch vocab-parallel cross entropy must compute what Megatron's fused kernel does."""
-
-import sys
+"""The vocab-parallel cross entropy behind every backend's torch log-probs."""
 
 import pytest
 import torch
 import torch.distributed as dist
 from tests.fast.dist_utils import init_gloo, run_multiprocess
 
-from miles.backends.training_utils.loss.hub import math_utils
 from miles.backends.training_utils.loss.hub.math_utils import compute_log_probs, vocab_parallel_cross_entropy
 
 SEQ, BATCH, VOCAB = 5, 2, 12
@@ -40,7 +37,6 @@ def _check_against_full_softmax(rank: int, world_size: int, group: dist.ProcessG
         full = _full_logits(dtype)
         expected_loss, expected_grad = _reference(full, _target())
         local = full[..., rank * shard : (rank + 1) * shard].clone().requires_grad_()
-        local_before = local.detach().clone()
 
         loss = vocab_parallel_cross_entropy(local, _target(), group)
         loss.backward(_grad_output())
@@ -50,38 +46,36 @@ def _check_against_full_softmax(rank: int, world_size: int, group: dist.ProcessG
         torch.testing.assert_close(
             local.grad, expected_grad[..., rank * shard : (rank + 1) * shard], atol=tol, rtol=tol
         )
-        torch.testing.assert_close(local.detach(), local_before, atol=0, rtol=0)
 
 
 def _check_against_megatron(rank: int, world_size: int, group: dist.ProcessGroup) -> None:
-    """Where Megatron is installed, the stand-in computes what its fused kernel computes."""
-    from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy
+    from megatron.core.tensor_parallel.cross_entropy import vocab_parallel_cross_entropy as megatron_cross_entropy
+
+    def megatron(logits, target, group):
+        return megatron_cross_entropy(logits, target, tp_group=group)
 
     shard = VOCAB // world_size
-    full = _full_logits(torch.bfloat16)
-    results = []
-    for fn in (vocab_parallel_cross_entropy, fused_vocab_parallel_cross_entropy):
-        local = full[..., rank * shard : (rank + 1) * shard].clone().requires_grad_()
-        loss = fn(local, _target(), group)
-        loss.backward(_grad_output())
-        results.append((loss.detach(), local.grad))
-    (loss, grad), (megatron_loss, megatron_grad) = results
-    torch.testing.assert_close(loss, megatron_loss, atol=1e-6, rtol=1e-6)
-    torch.testing.assert_close(grad, megatron_grad, atol=0, rtol=0)
+    for dtype in (torch.float32, torch.bfloat16):
+        full = _full_logits(dtype)
+        results = []
+        for fn in (vocab_parallel_cross_entropy, megatron):
+            local = full[..., rank * shard : (rank + 1) * shard].clone().requires_grad_()
+            loss = fn(local, _target(), group)
+            loss.backward(_grad_output())
+            results.append((loss.detach(), local.grad))
+        (loss, grad), (megatron_loss, megatron_grad) = results
+        assert torch.equal(loss, megatron_loss) and torch.equal(grad, megatron_grad), dtype
 
 
-def _check_compute_log_probs_without_megatron(group: dist.ProcessGroup) -> None:
-    """The FSDP and torchtitan log-prob path imports nothing from Megatron when it is not installed."""
-    sys.modules["megatron"] = None  # importing any of it now raises, as on an install without it
-    math_utils._MEGATRON_INSTALLED = False
+def _check_compute_log_probs(group: dist.ProcessGroup) -> None:
     logits = _full_logits(torch.float32)[:, 0]
     tokens = _target()[:, 0]
     # compute_log_probs returns [seq, 1]
     expected = torch.log_softmax(logits, dim=-1).gather(-1, tokens.unsqueeze(-1))
 
-    torch.testing.assert_close(compute_log_probs(logits, tokens, group), expected, atol=1e-6, rtol=1e-6)
+    torch.testing.assert_close(compute_log_probs(logits.clone(), tokens, group), expected, atol=1e-6, rtol=1e-6)
     # None: the default group
-    torch.testing.assert_close(compute_log_probs(logits, tokens, None), expected, atol=1e-6, rtol=1e-6)
+    torch.testing.assert_close(compute_log_probs(logits.clone(), tokens, None), expected, atol=1e-6, rtol=1e-6)
 
 
 def _run_checks(rank: int, world_size: int, port: int) -> None:
@@ -89,37 +83,29 @@ def _run_checks(rank: int, world_size: int, port: int) -> None:
     try:
         group = dist.new_group(list(range(world_size)))
         _check_against_full_softmax(rank, world_size, group)
-        _check_against_megatron(rank, world_size, group)
         if world_size == 1:
-            # last: it hides Megatron from the rest of this process
-            _check_compute_log_probs_without_megatron(group)
+            _check_compute_log_probs(group)
+    finally:
+        dist.destroy_process_group()
+
+
+def _run_megatron_check(rank: int, world_size: int, port: int) -> None:
+    init_gloo(rank, world_size, port=port)
+    try:
+        _check_against_megatron(rank, world_size, dist.new_group(list(range(world_size))))
     finally:
         dist.destroy_process_group()
 
 
 @pytest.mark.parametrize("world_size", [1, 2])
 def test_vocab_parallel_cross_entropy(world_size: int) -> None:
-    """Loss and gradient match a full-vocabulary log-softmax and Megatron's kernel on every TP shard."""
+    """Loss and gradient equal a plain log-softmax over the whole vocabulary, however it is split."""
     run_multiprocess(_run_checks, world_size=world_size)
 
 
-def _find_spec_seeing(found: dict[str, bool]):
-    def find_spec(name: str):
-        if name.split(".")[0] not in found:  # importing the parent package fails, as find_spec does
-            raise ModuleNotFoundError(name)
-        return object() if found.get(name, False) else None
-
-    return find_spec
-
-
-@pytest.mark.parametrize(
-    "found,installed",
-    [
-        ({"megatron": True, "megatron.core": True}, True),
-        ({"megatron": True}, False),  # megatron-energon's bare namespace, no megatron.core
-        ({}, False),  # no megatron package at all
-    ],
-)
-def test_megatron_counts_as_installed_only_with_megatron_core(monkeypatch, found, installed) -> None:
-    monkeypatch.setattr(math_utils.importlib.util, "find_spec", _find_spec_seeing(found))
-    assert math_utils._megatron_core_installed() is installed
+@pytest.mark.parametrize("world_size", [1, 2])
+def test_matches_megatron_vocab_parallel_cross_entropy(world_size: int) -> None:
+    """Bit for bit Megatron's math. Its fused kernel differs in the last bits, from torch.compile's
+    exp and summation order, and rounds the gradient to bf16."""
+    pytest.importorskip("megatron.core.tensor_parallel.cross_entropy")
+    run_multiprocess(_run_megatron_check, world_size=world_size)

From 9022cb2ebaa761f8faef397a14dc455d9c96b2b2 Mon Sep 17 00:00:00 2001
From: Hejian Sang <sanghj0923@gmail.com>
Date: Tue, 27 Oct 2026 10:42:46 -0700
Subject: [PATCH 9/9] Fail loudly on a None group; drop the last Megatron
 import from the log-prob path

- The torch kernel took dist.get_rank(None), i.e. the world group, where
  callers elsewhere mean no tensor parallelism by None: with 8 ranks and no
  TP every rank but 0 masked its targets and the SUM mixed data-parallel
  ranks. It now takes group.rank(), as Megatron's kernel did, so None or a
  non-member group raises. Every caller passes a group.
- The tensor-parallel true-on-policy gather imported Megatron's
  dist_all_gather_func, which is torch.distributed.all_gather_into_tensor;
  call that directly, so a wheel install can use it.
- --debug-unified-grad-fused-logprob's help no longer mentions a fused kernel.

Found by an adversarial review of this PR.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 miles/backends/training_utils/loss/hub/math_utils.py  | 11 +++++------
 miles/utils/args/configs/debug.py                     |  5 ++---
 .../loss_hub/test_vocab_parallel_cross_entropy.py     |  2 --
 3 files changed, 7 insertions(+), 11 deletions(-)

diff --git a/miles/backends/training_utils/loss/hub/math_utils.py b/miles/backends/training_utils/loss/hub/math_utils.py
index 2ca448ec521..c97e83f1008 100644
--- a/miles/backends/training_utils/loss/hub/math_utils.py
+++ b/miles/backends/training_utils/loss/hub/math_utils.py
@@ -315,7 +315,7 @@ class _VocabParallelCrossEntropy(torch.autograd.Function):
     """
 
     @staticmethod
-    def forward(ctx, logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup | None) -> torch.Tensor:
+    def forward(ctx, logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:
         ctx.input_dtype = logits.dtype
         logits = logits.float()
         logits_max = logits.max(dim=-1)[0]
@@ -323,7 +323,8 @@ def forward(ctx, logits: torch.Tensor, target: torch.Tensor, group: dist.Process
         logits.sub_(logits_max.unsqueeze(-1))
 
         partition_vocab_size = logits.size(-1)
-        vocab_start = dist.get_rank(group) * partition_vocab_size
+        # group.rank(), as Megatron: None or a rank outside the group fails here instead of reducing over the world
+        vocab_start = group.rank() * partition_vocab_size
         target_mask = (target < vocab_start) | (target >= vocab_start + partition_vocab_size)
         masked_target = (target - vocab_start).masked_fill(target_mask, 0)
         predicted_logits = logits.gather(-1, masked_target.unsqueeze(-1)).squeeze(-1).masked_fill(target_mask, 0.0)
@@ -347,7 +348,7 @@ def backward(ctx, grad_output: torch.Tensor):
         return grad_input.mul_(grad_output.unsqueeze(-1)).to(ctx.input_dtype), None, None
 
 
-def vocab_parallel_cross_entropy(logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup | None):
+def vocab_parallel_cross_entropy(logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup):
     return _VocabParallelCrossEntropy.apply(logits, target, group)
 
 
@@ -453,12 +454,10 @@ def forward(ctx, input_: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor
         if world_size == 1:
             return input_.contiguous()
 
-        from megatron.core.tensor_parallel.mappings import dist_all_gather_func
-
         gather_shape = list(input_.shape)
         gather_shape[0] *= world_size
         gathered = torch.empty(gather_shape, dtype=input_.dtype, device=input_.device)
-        dist_all_gather_func(gathered, input_.contiguous(), group=group)
+        dist.all_gather_into_tensor(gathered, input_.contiguous(), group=group)
         return torch.cat(gathered.chunk(world_size, dim=0), dim=-1).contiguous()
 
     @staticmethod
diff --git a/miles/utils/args/configs/debug.py b/miles/utils/args/configs/debug.py
index 6319554edfa..c99de49dc8a 100644
--- a/miles/utils/args/configs/debug.py
+++ b/miles/utils/args/configs/debug.py
@@ -265,9 +265,8 @@ class DebugConfig(BaseConfig):
         bool,
         Arg(
             help=(
-                "Debug/test only: compute the stored log probabilities through the same grad-enabled fused "
-                "cross entropy the training step uses, then detach the result, so the two invocations of the "
-                "fused kernel take one execution path instead of two."
+                "Debug/test only: compute the stored log probabilities through the same grad-enabled "
+                "cross entropy the training step uses, then detach the result."
             )
         ),
     ] = False
diff --git a/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
index 6deb93302f6..e21dfe2a2a3 100644
--- a/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
+++ b/tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py
@@ -74,8 +74,6 @@ def _check_compute_log_probs(group: dist.ProcessGroup) -> None:
     expected = torch.log_softmax(logits, dim=-1).gather(-1, tokens.unsqueeze(-1))
 
     torch.testing.assert_close(compute_log_probs(logits.clone(), tokens, group), expected, atol=1e-6, rtol=1e-6)
-    # None: the default group
-    torch.testing.assert_close(compute_log_probs(logits.clone(), tokens, None), expected, atol=1e-6, rtol=1e-6)
 
 
 def _run_checks(rank: int, world_size: int, port: int) -> None: