diff --git a/miles/backends/training_utils/loss/hub/math_utils.py b/miles/backends/training_utils/loss/hub/math_utils.py index 8a15144044..c97e83f100 100644 --- a/miles/backends/training_utils/loss/hub/math_utils.py +++ b/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 diff --git a/miles/utils/args/configs/debug.py b/miles/utils/args/configs/debug.py index 6319554edf..c99de49dc8 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/conftest.py b/tests/fast/backends/training_utils/loss_hub/conftest.py index 760eccd657..883fd63cb7 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 351d7d21fb..1b71e64363 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 new file mode 100644 index 0000000000..e21dfe2a2a --- /dev/null +++ b/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)