Compute log probs with one plain-torch vocab-parallel cross entropy, without Megatron

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

@@ -294,21 +294,62 @@ 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
-
-    # convert to [seq_len, batch_size, vocab_size] as expected by fused_vocab_parallel_cross_entropy
+    # [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 = -fused_vocab_parallel_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 -fused_vocab_parallel_cross_entropy(logits, tokens, process_group)
+    return -vocab_parallel_cross_entropy(logits, tokens, process_group)
+
+
+class _VocabParallelCrossEntropy(torch.autograd.Function):
+    """Cross entropy over logits split along the vocabulary over ``group``, in plain torch.
+
+    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) -> 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)
+        # 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)
+        exp_logits = logits.exp_()
+        sum_exp_logits = exp_logits.sum(dim=-1)
+
+        # 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))
+
+        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
+        # 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
+
+
+def vocab_parallel_cross_entropy(logits: torch.Tensor, target: torch.Tensor, group: dist.ProcessGroup):
+    return _VocabParallelCrossEntropy.apply(logits, target, group)
 
 
 def _apply_sampling_mask(
@@ -413,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

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

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

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
 
 

tests/fast/backends/training_utils/loss_hub/test_vocab_parallel_cross_entropy.py

@@ -0,0 +1,109 @@
+"""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.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:
+    # each TP=2 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, group: dist.ProcessGroup) -> None:
+    """Loss and gradient equal a plain log-softmax over the whole vocabulary, however it is split."""
+    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_()
+
+        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
+        )
+
+
+def _check_against_megatron(rank: int, world_size: int, group: dist.ProcessGroup) -> None:
+    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
+    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(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.clone(), 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(list(range(world_size)))
+        _check_against_full_softmax(rank, world_size, group)
+        if world_size == 1:
+            _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 equal a plain log-softmax over the whole vocabulary, however it is split."""
+    run_multiprocess(_run_checks, world_size=world_size)
+
+
+@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)