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>

Author: HJSang

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,20 +330,21 @@ 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]
         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
+        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)
 
 

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