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>

Author: HJSang

miles/backends/training_utils/loss/hub/math_utils.py

@@ -315,15 +315,16 @@ 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]
         dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=group)
         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)
 
 
@@ -439,12 +440,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

miles/utils/arguments.py

@@ -2456,9 +2456,8 @@ def add_debug_arguments(parser):
                 "--debug-unified-grad-fused-logprob",
                 action="store_true",
                 default=False,
-                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.",
+                help="Debug/test only: compute the stored log probabilities through the same grad-enabled "
+                "cross entropy the training step uses, then detach the result.",
             )
             parser.add_argument(
                 "--debug-deterministic-collective",

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: