From 3b7b8565997649ce0ab9b90a75be0e24d6bf7079 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 16 Sep 2026 19:35:53 +0800 Subject: [PATCH 1/2] Add optional Liger vocabulary-parallel cross entropy backend Signed-off-by: 0z5a (cherry picked from commit 5453799d020cbdf081f2e2e9d95eda4330e72075) --- deepspeed/runtime/engine.py | 6 +- deepspeed/runtime/tensor_parallel/config.py | 5 +- deepspeed/sequence/cross_entropy.py | 79 +++++++- docs/_tutorials/autotp-training.md | 10 + ...test_liger_vocab_parallel_cross_entropy.py | 180 ++++++++++++++++++ 5 files changed, 269 insertions(+), 11 deletions(-) create mode 100644 tests/unit/v1/sequence_parallelism/test_liger_vocab_parallel_cross_entropy.py diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index e7f10d1dac1b..039020a043fa 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -1057,12 +1057,16 @@ def finalize_autotp(autotp=None, attach_uc_metadata=False): raise ValueError("Unable to choose a loss for multiple no-gather vocab-parallel LM heads") if vocab_parallel_heads: from deepspeed.sequence.cross_entropy import configure_vocab_parallel_loss - configure_vocab_parallel_loss(model, vocab_parallel_heads[0]) + configure_vocab_parallel_loss(model, + vocab_parallel_heads[0], + backend=tp_config.vocab_parallel_ce_backend) elif tp_config.vocab_parallel_lm_head: # Every partitioning path must agree; otherwise the request degrades into ordinary # AutoTP with a gathered head and no distributed loss, which is easy to miss. raise ValueError( "vocab_parallel_lm_head requires a supported nn.Linear named 'lm_head' or 'embed_out'") + elif tp_config.vocab_parallel_ce_backend != "torch": + raise ValueError("vocab_parallel_ce_backend='liger' requires a no-gather vocabulary-parallel LM head") if attach_uc_metadata: setattr(model, UNIVERSAL_CHECKPOINT_INFO, collect_autotp_universal_checkpoint_info(model)) diff --git a/deepspeed/runtime/tensor_parallel/config.py b/deepspeed/runtime/tensor_parallel/config.py index 9d196eca63df..2448ee14d28f 100644 --- a/deepspeed/runtime/tensor_parallel/config.py +++ b/deepspeed/runtime/tensor_parallel/config.py @@ -7,7 +7,7 @@ from deepspeed.runtime.config_utils import DeepSpeedConfigModel import torch from pydantic import Field -from typing import Optional, Dict, Any +from typing import Optional, Dict, Any, Literal class AUTOTP_MODE(Enum): @@ -53,6 +53,9 @@ class TPTrainingConfig(DeepSpeedConfigModel): vocab_parallel_lm_head: bool = False """Keep an untied LM head vocabulary-sharded and install a compatible distributed loss.""" + vocab_parallel_ce_backend: Literal["torch", "liger"] = "torch" + """Optional Liger CE acceleration; the PyTorch reference remains the default.""" + tensor_parallel: TPConfig = Field({}, alias="tp") """ Configuration for tensor parallelism used to split the model across several diff --git a/deepspeed/sequence/cross_entropy.py b/deepspeed/sequence/cross_entropy.py index 81699f950add..d47a8c5026d6 100644 --- a/deepspeed/sequence/cross_entropy.py +++ b/deepspeed/sequence/cross_entropy.py @@ -5,11 +5,53 @@ import torch from torch import nn +from functools import lru_cache import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator from deepspeed.utils.logging import logger +@lru_cache(maxsize=1) +def _load_liger_vocab_ce(): + try: + from liger_kernel.ops.vocab_parallel_cross_entropy import LigerVocabParallelCEFunction + except ImportError as error: + raise ImportError("Liger vocabulary-parallel CE requires liger-kernel>=0.8.1; " + "install it or select backend='torch'.") from error + return LigerVocabParallelCEFunction + + +def _liger_vocab_cross_entropy(logits, target, tp_group, ignore_index): + kernel = _load_liger_vocab_ce() + loss = kernel.apply(logits.reshape(-1, 1, logits.shape[-1]), target.reshape(-1, 1), tp_group, ignore_index, 0.0, + "pytorch", "global").reshape_as(target) + if loss.requires_grad: + consumed = False + + def guard_backward(gradient): + nonlocal consumed + # Liger overwrites its saved exp buffer with the first gradient, so a retained-graph + # second backward must fail rather than silently reuse that buffer. + if consumed or torch.is_grad_enabled(): + raise RuntimeError("Liger CE supports one first-order backward per forward; " + "use backend='torch' for retained-graph or higher-order gradients.") + consumed = True + return gradient + + loss.register_hook(guard_backward) + return loss + + +def _liger_can_run(vocab_parallel_logits, target, local_vocab_size, global_vocab_size, tp_group): + # Liger's kernel all-gathers fixed-width shards, while the PyTorch path trims uneven + # ones, so the kernel stays restricted to equal shards. + tp_world_size = dist.get_world_size(tp_group) if tp_group is not None else 1 + return (get_accelerator().device_name() == "cuda" and get_accelerator().on_accelerator(vocab_parallel_logits) + and vocab_parallel_logits.dtype in (torch.float32, torch.float16, torch.bfloat16) and target.numel() > 0 + and local_vocab_size * tp_world_size == global_vocab_size) + + class _VocabParallelCrossEntropy(torch.autograd.Function): @staticmethod @@ -156,7 +198,8 @@ def vocab_parallel_cross_entropy(vocab_parallel_logits, global_vocab_size=None, ignore_index=-100, reduction="mean", - gather_sequence_loss=False): + gather_sequence_loss=False, + backend="torch"): """Compute cross entropy over vocabulary-sharded logits. Tensor parallel ranks collectively own the last (vocabulary) dimension. Sequence @@ -175,6 +218,8 @@ def vocab_parallel_cross_entropy(vocab_parallel_logits, raise ValueError("vocab_parallel_logits must contain at least one local vocabulary entry") if reduction not in ("none", "sum", "mean"): raise ValueError(f"Unsupported reduction: {reduction}") + if backend not in ("torch", "liger"): + raise ValueError(f"Unsupported vocabulary-parallel CE backend: {backend}") if gather_sequence_loss and reduction != "none": raise ValueError("gather_sequence_loss is only supported with reduction='none'") @@ -194,8 +239,16 @@ def vocab_parallel_cross_entropy(vocab_parallel_logits, if invalid_target.any().item(): raise ValueError(f"Target is out of range for vocabulary size {global_vocab_size}") - loss = _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target, tp_group, vocab_start_index, - vocab_end_index, ignore_index) + use_liger = backend == "liger" and _liger_can_run(vocab_parallel_logits, target, local_vocab_size, + global_vocab_size, tp_group) + if use_liger: + loss = _liger_vocab_cross_entropy(vocab_parallel_logits, target.to(dtype=torch.long), tp_group, ignore_index) + else: + if backend == "liger": + logger.warning_once("Liger CE requires nonempty CUDA tokens, a supported dtype, and equal vocabulary " + "shards; using the PyTorch reference backend for this layout.") + loss = _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target, tp_group, vocab_start_index, + vocab_end_index, ignore_index) if reduction == "none": if gather_sequence_loss: if sp_group is None: @@ -221,8 +274,11 @@ def vocab_sequence_parallel_cross_entropy(vocab_parallel_logits, vocab_end_index=None, ignore_index=-100, reduction="none", - gather_sequence_loss=True): + gather_sequence_loss=True, + backend="torch"): """Sequence-parallel wrapper preserving the legacy local-slice gradient.""" + if backend not in ("torch", "liger"): + raise ValueError(f"Unsupported vocabulary-parallel CE backend: {backend}") if gather_sequence_loss and reduction != "none": raise ValueError("gather_sequence_loss is only supported with reduction='none'") if gather_sequence_loss and sp_group is None: @@ -235,7 +291,8 @@ def vocab_sequence_parallel_cross_entropy(vocab_parallel_logits, vocab_start_index=vocab_start_index, vocab_end_index=vocab_end_index, ignore_index=ignore_index, - reduction=reduction) + reduction=reduction, + backend=backend) if not gather_sequence_loss: return loss return _GatherSequenceLoss.apply(loss, sp_group, False) @@ -250,7 +307,8 @@ def __init__(self, vocab_end_index=None, ignore_index=-100, reduction="mean", - gather_sequence_loss=False): + gather_sequence_loss=False, + backend="torch"): super().__init__() self.tp_group = tp_group self.sp_group = sp_group @@ -259,6 +317,7 @@ def __init__(self, self.ignore_index = ignore_index self.reduction = reduction self.gather_sequence_loss = gather_sequence_loss + self.backend = backend def forward(self, vocab_parallel_logits, target): return vocab_parallel_cross_entropy(vocab_parallel_logits, @@ -269,7 +328,8 @@ def forward(self, vocab_parallel_logits, target): vocab_end_index=self.vocab_end_index, ignore_index=self.ignore_index, reduction=self.reduction, - gather_sequence_loss=self.gather_sequence_loss) + gather_sequence_loss=self.gather_sequence_loss, + backend=self.backend) class VocabParallelCausalLMLoss: @@ -350,7 +410,7 @@ def __call__(self, logits, labels=None, vocab_size=None, shift_labels=None, num_ return loss -def configure_vocab_parallel_loss(model, vocab_parallel_head, sp_group=None, ignore_index=-100): +def configure_vocab_parallel_loss(model, vocab_parallel_head, sp_group=None, ignore_index=-100, backend="torch"): """Install the causal-LM loss required by a no-gather vocabulary projection. Leave ``sp_group`` as ``None`` when running under DeepSpeed's Ulysses @@ -365,7 +425,8 @@ def configure_vocab_parallel_loss(model, vocab_parallel_head, sp_group=None, ign sp_group=sp_group, vocab_start_index=vocab_parallel_head.vocab_start_index, vocab_end_index=vocab_parallel_head.vocab_end_index, - ignore_index=ignore_index) + ignore_index=ignore_index, + backend=backend) original_loss_function = model.loss_function registered_loss_module = getattr(model, "_modules", {}).pop("loss_function", None) diff --git a/docs/_tutorials/autotp-training.md b/docs/_tutorials/autotp-training.md index b997735b681f..4a7560ab3d8d 100644 --- a/docs/_tutorials/autotp-training.md +++ b/docs/_tutorials/autotp-training.md @@ -166,6 +166,16 @@ installs a pure-PyTorch distributed causal-LM loss through the model's log-sum-exp and target lookup without gathering vocabulary logits. Uneven vocabulary shards are supported. +For optional fused CE, install `liger-kernel>=0.8.1` and set +`tensor_parallel.vocab_parallel_ce_backend` to `"liger"` alongside +`vocab_parallel_lm_head`. The default remains `"torch"`, with no Liger dependency. +The Liger backend accepts nonempty CUDA logits in FP32, FP16, or BF16 with equal +vocabulary shards. Uneven shards and unsupported layouts use the PyTorch backend +on every TP rank. This fuses CE, not the linear projection: local logits are still +materialized. Liger supports one first-order backward per forward; retained-graph +second backward and higher-order gradients raise an error rather than reusing +its overwritten gradient buffer. Select `"torch"` for those workloads. + Only an `nn.Linear` whose final name segment is `lm_head` or `embed_out` is supported. If no such head exists, initialization fails instead of quietly falling back to an ordinary gathered head. When `partition_config` or a diff --git a/tests/unit/v1/sequence_parallelism/test_liger_vocab_parallel_cross_entropy.py b/tests/unit/v1/sequence_parallelism/test_liger_vocab_parallel_cross_entropy.py new file mode 100644 index 000000000000..c4c5a80ce284 --- /dev/null +++ b/tests/unit/v1/sequence_parallelism/test_liger_vocab_parallel_cross_entropy.py @@ -0,0 +1,180 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from copy import deepcopy + +import pytest +import torch +from torch import nn +from torch.nn import functional as F + +import deepspeed +import deepspeed.comm as dist +import deepspeed.sequence.cross_entropy as cross_entropy +from deepspeed.accelerator import get_accelerator +from deepspeed.module_inject.layers import VocabParallelLinear +from deepspeed.module_inject.tp_shard import get_shard_size_list +from deepspeed.utils import groups +from unit.common import DistributedTest +from unit.v1.autotp.test_autotp_training import reset_tp_model_init_state + + +def test_liger_cpu_fallback_matches_torch(): + logits = torch.randn(3, 18, requires_grad=True) + reference = logits.detach().clone().requires_grad_(True) + target = torch.tensor([0, -100, 17]) + actual = cross_entropy.vocab_parallel_cross_entropy(logits, target, backend="liger") + expected = F.cross_entropy(reference, target) + torch.testing.assert_close(actual, expected) + actual.backward() + expected.backward() + torch.testing.assert_close(logits.grad, reference.grad) + + +class LigerOutputTrainingModel(nn.Module): + + def __init__(self, vocabulary): + super().__init__() + self.projection = nn.Linear(32, 32) + self.lm_head = nn.Linear(32, vocabulary) + self.loss_function = cross_entropy.VocabParallelCausalLMLoss() + + def forward(self, inputs, labels): + logits = self.lm_head(self.projection(inputs).tanh()) + return logits, self.loss_function(logits, labels=labels) + + +class TestLigerVocabParallelCE(DistributedTest): + world_size = 2 + reuse_dist_env = False + + def test_losses_and_gradients(self): + pytest.importorskip("liger_kernel.ops.vocab_parallel_cross_entropy") + device = get_accelerator().current_device_name() + if not str(device).startswith("cuda"): + pytest.skip("Liger CE requires CUDA") + rank = dist.get_rank() + for vocabulary in (34, 35): + # Each layer layout owns a distinct process group; metadata is cached + # for a fixed layout throughout a training run. + group = dist.new_group(ranks=list(range(self.world_size))) + sizes = [vocabulary // 2, vocabulary - vocabulary // 2] + start = sum(sizes[:rank]) + end = start + sizes[rank] + for dtype in (torch.float32, torch.float16, torch.bfloat16): + for reduction in ("none", "sum", "mean"): + torch.manual_seed(8173) + full = (torch.randn(2, 3, vocabulary, device=device) * 20).to(dtype) + reference = full.detach().float().clone().requires_grad_(True) + local = full[..., start:end].detach().contiguous().requires_grad_(True) + before = local.detach().clone() + target = torch.tensor([[0, -100, vocabulary - 1], [1, 17, 8]], device=device) + expected = F.cross_entropy(reference.reshape(-1, vocabulary), + target.reshape(-1), + reduction=reduction) + if reduction == "none": + expected = expected.reshape_as(target) + actual = cross_entropy.vocab_parallel_cross_entropy(local, + target, + group, + vocab_start_index=start, + vocab_end_index=end, + reduction=reduction, + backend="liger") + torch.testing.assert_close(actual, expected, atol=2e-5, rtol=2e-5) + actual.sum().backward() + expected.sum().backward() + torch.testing.assert_close(local.grad, + reference.grad[..., start:end].to(dtype), + atol=2e-5, + rtol=2e-3) + torch.testing.assert_close(local.detach(), before, atol=0, rtol=0) + group = dist.new_group(ranks=list(range(self.world_size))) + logits = torch.randn(2, 3, 17, device=device, requires_grad=True) + ignored = torch.full((2, 3), -100, device=device) + loss = cross_entropy.vocab_parallel_cross_entropy(logits, ignored, group, backend="liger") + torch.testing.assert_close(loss, torch.zeros_like(loss)) + loss.backward() + torch.testing.assert_close(logits.grad, torch.zeros_like(logits)) + logits = torch.randn(2, 3, 17, device=device, requires_grad=True) + loss = cross_entropy.vocab_parallel_cross_entropy(logits, torch.zeros_like(ignored), group, backend="liger") + loss.backward(retain_graph=True) + with pytest.raises(RuntimeError, match="one first-order backward"): + loss.backward() + logits = torch.randn(2, 3, 17, device=device, requires_grad=True) + loss = cross_entropy.vocab_parallel_cross_entropy(logits, torch.zeros_like(ignored), group, backend="liger") + with pytest.raises(RuntimeError, match="one first-order backward"): + torch.autograd.grad(loss, logits, create_graph=True) + + @pytest.mark.parametrize("vocabulary", [34, 35]) + def test_five_optimizer_steps(self, vocabulary): + pytest.importorskip("liger_kernel.ops.vocab_parallel_cross_entropy") + reset_tp_model_init_state() + device = get_accelerator().current_device_name() + if not str(device).startswith("cuda"): + pytest.skip("Liger CE requires CUDA") + torch.manual_seed(8173) + model = LigerOutputTrainingModel(vocabulary).to(device) + reference = deepcopy(model) + optimizer = torch.optim.SGD(model.parameters(), lr=0.05) + reference_optimizer = torch.optim.SGD(reference.parameters(), lr=0.05) + config = { + "train_micro_batch_size_per_gpu": 1, + "gradient_accumulation_steps": 2, + "gradient_clipping": 0.0, + "zero_optimization": { + "stage": 0 + }, + "tensor_parallel": { + "autotp_size": 2, + "vocab_parallel_lm_head": True, + "vocab_parallel_ce_backend": "liger", + "partition_config": { + "use_default_specs": False, + "layer_specs": [] + } + } + } + try: + engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config) + head = engine.module.lm_head + assert isinstance(head, VocabParallelLinear) + assert engine.module.loss_function.backend == "liger" + rank = dist.get_rank(groups.get_tensor_model_parallel_group()) + sizes = get_shard_size_list(vocabulary, 2, head.tp_meta, "lm_head") + start = sum(sizes[:rank]) + for step in range(5): + reference_optimizer.zero_grad() + for micro in range(2): + torch.manual_seed(100 + step * 2 + micro) + inputs = torch.randn(2, 6, 32, device=device) + labels = torch.randint(vocabulary, (2, 6), device=device) + labels[0, 2] = -100 + reference_logits, expected = reference(inputs, labels) + logits, actual = engine(inputs, labels) + torch.testing.assert_close(logits, reference_logits[..., start:start + sizes[rank]]) + torch.testing.assert_close(actual, expected) + (expected / 2).backward() + engine.backward(actual) + torch.testing.assert_close(head.weight.grad, + reference.lm_head.weight.grad[start:start + sizes[rank]], + atol=1e-6, + rtol=1e-5) + torch.testing.assert_close(engine.module.projection.weight.grad, + reference.projection.weight.grad, + atol=1e-6, + rtol=1e-5) + engine.step() + reference_optimizer.step() + assert engine.global_steps == step + 1 + torch.testing.assert_close(head.weight, reference.lm_head.weight[start:start + sizes[rank]]) + torch.testing.assert_close(head.bias, reference.lm_head.bias[start:start + sizes[rank]]) + torch.testing.assert_close(engine.module.projection.weight, reference.projection.weight) + torch.testing.assert_close(engine.module.projection.bias, reference.projection.bias) + if rank == 0: + print(f"PR-D vocab={vocabulary} step={step + 1} loss={actual.item():.6f} reference-aligned", + flush=True) + finally: + reset_tp_model_init_state() From 11525798c87a82c8bae03a5a2713f118c43b036c Mon Sep 17 00:00:00 2001 From: 0z5a Date: Mon, 21 Sep 2026 10:09:38 +0800 Subject: [PATCH 2/2] Address review feedback on the Liger vocabulary-parallel CE backend - Select Liger only behind an explicit tensor-parallel group. A None group is the default world to some Liger versions, which would read data-parallel replicas as vocabulary shards and produce a wrong global loss. - Gate the kernel on the accelerator's Triton support instead of CUDA only, so Triton-capable XPU and other accelerators take the same path. - Accept backend on VocabParallelCausalLMLoss so configure_vocab_parallel_loss no longer raises TypeError for vocab_parallel_ce_backend. - Install liger-kernel in the development requirements so the v1 test executes in CI instead of being skipped by importorskip. - Cover the new dispatch contract with tests and keep the documented boundaries current. Signed-off-by: 0z5a --- deepspeed/sequence/cross_entropy.py | 30 +++++-- docs/_tutorials/autotp-training.md | 8 +- requirements/requirements-dev.txt | 1 + ...test_liger_vocab_parallel_cross_entropy.py | 82 +++++++++++++++++-- 4 files changed, 101 insertions(+), 20 deletions(-) diff --git a/deepspeed/sequence/cross_entropy.py b/deepspeed/sequence/cross_entropy.py index d47a8c5026d6..a8699d9afb45 100644 --- a/deepspeed/sequence/cross_entropy.py +++ b/deepspeed/sequence/cross_entropy.py @@ -44,10 +44,15 @@ def guard_backward(gradient): def _liger_can_run(vocab_parallel_logits, target, local_vocab_size, global_vocab_size, tp_group): - # Liger's kernel all-gathers fixed-width shards, while the PyTorch path trims uneven - # ones, so the kernel stays restricted to equal shards. - tp_world_size = dist.get_world_size(tp_group) if tp_group is not None else 1 - return (get_accelerator().device_name() == "cuda" and get_accelerator().on_accelerator(vocab_parallel_logits) + # A ``None`` group is ambiguous to Liger: depending on its version it may resolve to the + # default world, which would fold data-parallel replicas into the vocabulary dimension, so + # the kernel only runs behind a real tensor-parallel group. Liger's kernel all-gathers + # fixed-width shards while the PyTorch path trims uneven ones, so the kernel stays + # restricted to equal shards. + if tp_group is None: + return False + tp_world_size = dist.get_world_size(tp_group) + return (get_accelerator().is_triton_supported() and get_accelerator().on_accelerator(vocab_parallel_logits) and vocab_parallel_logits.dtype in (torch.float32, torch.float16, torch.bfloat16) and target.numel() > 0 and local_vocab_size * tp_world_size == global_vocab_size) @@ -240,12 +245,13 @@ def vocab_parallel_cross_entropy(vocab_parallel_logits, raise ValueError(f"Target is out of range for vocabulary size {global_vocab_size}") use_liger = backend == "liger" and _liger_can_run(vocab_parallel_logits, target, local_vocab_size, - global_vocab_size, tp_group) + global_vocab_size, tp_group) if use_liger: loss = _liger_vocab_cross_entropy(vocab_parallel_logits, target.to(dtype=torch.long), tp_group, ignore_index) else: if backend == "liger": - logger.warning_once("Liger CE requires nonempty CUDA tokens, a supported dtype, and equal vocabulary " + logger.warning_once("Liger CE requires an explicit tensor-parallel group, a Triton-supported " + "accelerator, nonempty accelerator tokens, a supported dtype, and equal vocabulary " "shards; using the PyTorch reference backend for this layout.") loss = _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target, tp_group, vocab_start_index, vocab_end_index, ignore_index) @@ -352,12 +358,19 @@ class VocabParallelCausalLMLoss: metadata. """ - def __init__(self, tp_group=None, sp_group=None, vocab_start_index=None, vocab_end_index=None, ignore_index=-100): + def __init__(self, + tp_group=None, + sp_group=None, + vocab_start_index=None, + vocab_end_index=None, + ignore_index=-100, + backend="torch"): self.tp_group = tp_group self.sp_group = sp_group self.vocab_start_index = vocab_start_index self.vocab_end_index = vocab_end_index self.ignore_index = ignore_index + self.backend = backend self._resolved_vocab_key = None self._resolved_global_vocab_size = None @@ -403,7 +416,8 @@ def __call__(self, logits, labels=None, vocab_size=None, shift_labels=None, num_ vocab_end_index=self.vocab_end_index, global_vocab_size=global_vocab_size, ignore_index=self.ignore_index, - reduction=reduction) + reduction=reduction, + backend=self.backend) if num_items_in_batch is not None: denominator = torch.as_tensor(num_items_in_batch, device=loss.device, dtype=loss.dtype) loss = loss / denominator.clamp_min(1) diff --git a/docs/_tutorials/autotp-training.md b/docs/_tutorials/autotp-training.md index 4a7560ab3d8d..9ae232b567ae 100644 --- a/docs/_tutorials/autotp-training.md +++ b/docs/_tutorials/autotp-training.md @@ -169,9 +169,11 @@ vocabulary shards are supported. For optional fused CE, install `liger-kernel>=0.8.1` and set `tensor_parallel.vocab_parallel_ce_backend` to `"liger"` alongside `vocab_parallel_lm_head`. The default remains `"torch"`, with no Liger dependency. -The Liger backend accepts nonempty CUDA logits in FP32, FP16, or BF16 with equal -vocabulary shards. Uneven shards and unsupported layouts use the PyTorch backend -on every TP rank. This fuses CE, not the linear projection: local logits are still +The Liger backend accepts nonempty accelerator logits in FP32, FP16, or BF16 with +equal vocabulary shards behind an explicit tensor-parallel group, on any accelerator +whose Triton support the DeepSpeed accelerator reports. Uneven shards, a missing +tensor-parallel group, and unsupported layouts use the PyTorch backend on every TP +rank. This fuses CE, not the linear projection: local logits are still materialized. Liger supports one first-order backward per forward; retained-graph second backward and higher-order gradients raise an error rather than reusing its overwritten gradient buffer. Select `"torch"` for those workloads. diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 10a24fc697e9..ee1b0502ae9a 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -5,6 +5,7 @@ deepspeed-kernels ; sys_platform == 'linux' docutils<0.18 future importlib-metadata>=4 +liger-kernel>=0.8.1 mup pre-commit>=3.2.0 pytest>=7.2.0,<8.4.0 diff --git a/tests/unit/v1/sequence_parallelism/test_liger_vocab_parallel_cross_entropy.py b/tests/unit/v1/sequence_parallelism/test_liger_vocab_parallel_cross_entropy.py index c4c5a80ce284..c5ee8645863c 100644 --- a/tests/unit/v1/sequence_parallelism/test_liger_vocab_parallel_cross_entropy.py +++ b/tests/unit/v1/sequence_parallelism/test_liger_vocab_parallel_cross_entropy.py @@ -21,7 +21,23 @@ from unit.v1.autotp.test_autotp_training import reset_tp_model_init_state -def test_liger_cpu_fallback_matches_torch(): +def _require_liger(): + pytest.importorskip("liger_kernel.ops.vocab_parallel_cross_entropy") + if not get_accelerator().is_triton_supported(): + pytest.skip("Liger CE requires a Triton-supported accelerator") + + +def _record_liger_calls(monkeypatch, calls): + original = cross_entropy._liger_vocab_cross_entropy + + def record(*args, **kwargs): + calls.append(args) + return original(*args, **kwargs) + + monkeypatch.setattr(cross_entropy, "_liger_vocab_cross_entropy", record) + + +def test_liger_without_tp_group_uses_pytorch_reference(): logits = torch.randn(3, 18, requires_grad=True) reference = logits.detach().clone().requires_grad_(True) target = torch.tensor([0, -100, 17]) @@ -33,6 +49,11 @@ def test_liger_cpu_fallback_matches_torch(): torch.testing.assert_close(logits.grad, reference.grad) +def test_vocab_parallel_causal_lm_loss_forwards_backend(): + assert cross_entropy.VocabParallelCausalLMLoss().backend == "torch" + assert cross_entropy.VocabParallelCausalLMLoss(backend="liger").backend == "liger" + + class LigerOutputTrainingModel(nn.Module): def __init__(self, vocabulary): @@ -46,15 +67,61 @@ def forward(self, inputs, labels): return logits, self.loss_function(logits, labels=labels) +class TestLigerBackendSelection(DistributedTest): + world_size = 2 + reuse_dist_env = False + + def test_missing_tp_group_uses_pytorch_reference(self, monkeypatch): + """A None group is the default world to some Liger versions, so it must not select Liger.""" + device = get_accelerator().current_device_name() + calls = [] + _record_liger_calls(monkeypatch, calls) + torch.manual_seed(8173) + logits = torch.randn(2, 3, 17, device=device, requires_grad=True) + reference = logits.detach().clone().requires_grad_(True) + target = torch.tensor([[0, 4, 16], [9, -100, 3]], device=device) + actual = cross_entropy.vocab_parallel_cross_entropy(logits, target, backend="liger") + expected = F.cross_entropy(reference.reshape(-1, reference.shape[-1]), target.reshape(-1)) + torch.testing.assert_close(actual, expected) + actual.backward() + expected.backward() + torch.testing.assert_close(logits.grad, reference.grad) + assert calls == [] + + def test_tp_group_selects_liger_only_for_equal_shards(self, monkeypatch): + _require_liger() + device = get_accelerator().current_device_name() + rank = dist.get_rank() + group = dist.new_group(ranks=list(range(self.world_size))) + target = torch.tensor([[0, 4, 16], [9, -100, 3]], device=device) + for vocabulary in (18, 17): + calls = [] + _record_liger_calls(monkeypatch, calls) + torch.manual_seed(8173) + full = torch.randn(2, 3, vocabulary, device=device) * 20 + sizes = [vocabulary // 2, vocabulary - vocabulary // 2] + start = sum(sizes[:rank]) + end = start + sizes[rank] + local = full[..., start:end].detach().contiguous().requires_grad_(True) + reference = full.detach().float().clone().requires_grad_(True) + expected = F.cross_entropy(reference.reshape(-1, vocabulary), target.reshape(-1)) + actual = cross_entropy.vocab_parallel_cross_entropy(local, + target, + group, + vocab_start_index=start, + vocab_end_index=end, + backend="liger") + torch.testing.assert_close(actual, expected, atol=2e-5, rtol=2e-5) + assert len(calls) == (1 if vocabulary % 2 == 0 else 0) + + class TestLigerVocabParallelCE(DistributedTest): world_size = 2 reuse_dist_env = False def test_losses_and_gradients(self): - pytest.importorskip("liger_kernel.ops.vocab_parallel_cross_entropy") + _require_liger() device = get_accelerator().current_device_name() - if not str(device).startswith("cuda"): - pytest.skip("Liger CE requires CUDA") rank = dist.get_rank() for vocabulary in (34, 35): # Each layer layout owns a distinct process group; metadata is cached @@ -110,11 +177,9 @@ def test_losses_and_gradients(self): @pytest.mark.parametrize("vocabulary", [34, 35]) def test_five_optimizer_steps(self, vocabulary): - pytest.importorskip("liger_kernel.ops.vocab_parallel_cross_entropy") + _require_liger() reset_tp_model_init_state() device = get_accelerator().current_device_name() - if not str(device).startswith("cuda"): - pytest.skip("Liger CE requires CUDA") torch.manual_seed(8173) model = LigerOutputTrainingModel(vocabulary).to(device) reference = deepcopy(model) @@ -174,7 +239,6 @@ def test_five_optimizer_steps(self, vocabulary): torch.testing.assert_close(engine.module.projection.weight, reference.projection.weight) torch.testing.assert_close(engine.module.projection.bias, reference.projection.bias) if rank == 0: - print(f"PR-D vocab={vocabulary} step={step + 1} loss={actual.item():.6f} reference-aligned", - flush=True) + print(f"vocab={vocabulary} step={step + 1} loss={actual.item():.6f} reference-aligned", flush=True) finally: reset_tp_model_init_state()