From a003eee14b8ed3beec3e3e2584cb46b142bafcb6 Mon Sep 17 00:00:00 2001 From: Justin Hu Date: Tue, 11 Aug 2026 17:06:18 +0000 Subject: [PATCH 1/7] feat: add Megatron fused linear cross entropy backends Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...ark_megatron_fused_linear_cross_entropy.py | 529 +++++++++++++++ src/liger_kernel/megatron/__init__.py | 4 + .../megatron/fused_linear_cross_entropy.py | 43 ++ src/liger_kernel/ops/__init__.py | 6 + src/liger_kernel/ops/cutile/ops/__init__.py | 4 + .../megatron_fused_linear_cross_entropy.py | 597 +++++++++++++++++ .../megatron_fused_linear_cross_entropy.py | 200 ++++++ src/liger_kernel/ops/triton/__init__.py | 12 + src/liger_kernel/ops/triton/ops/__init__.py | 9 + .../megatron_fused_linear_cross_entropy.py | 600 ++++++++++++++++++ .../test_cutile_fused_linear_cross_entropy.py | 152 +++++ .../test_fused_linear_cross_entropy.py | 213 +++++++ 12 files changed, 2369 insertions(+) create mode 100644 benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py create mode 100644 src/liger_kernel/megatron/fused_linear_cross_entropy.py create mode 100644 src/liger_kernel/ops/cutile/ops/megatron_fused_linear_cross_entropy.py create mode 100644 src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py create mode 100644 src/liger_kernel/ops/triton/__init__.py create mode 100644 src/liger_kernel/ops/triton/ops/__init__.py create mode 100644 src/liger_kernel/ops/triton/ops/megatron_fused_linear_cross_entropy.py create mode 100644 test/megatron/test_cutile_fused_linear_cross_entropy.py create mode 100644 test/megatron/test_fused_linear_cross_entropy.py diff --git a/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py b/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py new file mode 100644 index 000000000..8b012c5ed --- /dev/null +++ b/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py @@ -0,0 +1,529 @@ +"""Benchmark hidden-to-loss FLCE against Megatron's materialized output-loss stack. + +Megatron-Core does not provide a fused linear cross-entropy kernel. Its comparable +training path is: + + vocab-parallel linear -> materialized local logits -> fused vocab-parallel CE + +This script compares that path with ``LigerMegatronFusedLinearCrossEntropy``, +which saves low-precision CE state to avoid projection recomputation. The +``liger-triton`` provider additionally replaces all three local GEMMs with +portable Triton kernels. When Megatron-Core is installed, the ``megatron-core`` +provider uses its fused CE. The always-available ``megatron-compatible`` +provider uses Liger's drop-in Megatron CE. + +Backward timing creates a fresh graph outside each timed event pair, so only +backward execution is measured while respecting Megatron's single-use fused CE +graph. Fixed iteration counts keep all tensor-parallel ranks in collective +lockstep. + +Examples: + + python benchmark_megatron_fused_linear_cross_entropy.py --tp-size 1 + torchrun --help # not needed; the script spawns TP ranks itself + python benchmark_megatron_fused_linear_cross_entropy.py --tp-size 4 \ + --token-counts 512 2048 --vocab-sizes 32000 128256 +""" + +from __future__ import annotations + +import argparse +import os +import tempfile + +from dataclasses import dataclass +from pathlib import Path + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F + +from utils import BenchmarkData +from utils import get_formatted_time +from utils import get_gpu_name +from utils import update_benchmark_data_csv + +from liger_kernel.megatron import LigerMegatronCrossEntropy +from liger_kernel.megatron import LigerMegatronFusedLinearCrossEntropy +from liger_kernel.ops.triton.ops.megatron_fused_linear_cross_entropy import ( + liger_megatron_fused_linear_cross_entropy as triton_megatron_fused_linear_cross_entropy, +) + +try: + from liger_kernel.ops.cutile.ops.megatron_fused_linear_cross_entropy import ( + liger_megatron_fused_linear_cross_entropy as cutile_megatron_fused_linear_cross_entropy, + ) + + _CUTILE_AVAILABLE = True +except ImportError: + cutile_megatron_fused_linear_cross_entropy = None + _CUTILE_AVAILABLE = False + +try: + from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy + + _MEGATRON_CORE_AVAILABLE = True +except ImportError: + fused_vocab_parallel_cross_entropy = None + _MEGATRON_CORE_AVAILABLE = False + + +_SPEED_SAMPLES = 5 +_MEMORY_SAMPLES = 3 +_DTYPES = { + "bf16": torch.bfloat16, + "fp16": torch.float16, +} + + +@dataclass +class _ProviderState: + hidden: torch.Tensor + weight: torch.Tensor + bias: torch.Tensor | None + target: torch.Tensor + forward: object + + def clear_grads(self): + self.hidden.grad = None + self.weight.grad = None + if self.bias is not None: + self.bias.grad = None + + +def _all_reduce_hidden_grad(grad: torch.Tensor, tp_group): + dist.all_reduce(grad, op=dist.ReduceOp.SUM, group=tp_group) + return grad + + +def _make_state( + provider: str, + hidden_master: torch.Tensor, + weight_master: torch.Tensor, + bias_master: torch.Tensor | None, + target: torch.Tensor, + tp_group, + tp_size: int, +) -> _ProviderState: + hidden = hidden_master.clone().requires_grad_(True) + weight = weight_master.clone().requires_grad_(True) + bias = bias_master.clone().requires_grad_(True) if bias_master is not None else None + + if provider == "liger": + loss = LigerMegatronFusedLinearCrossEntropy() + forward = lambda: loss(hidden, weight, target, bias=bias, tp_group=tp_group) + elif provider == "liger-triton": + forward = lambda: triton_megatron_fused_linear_cross_entropy( + hidden, + weight, + target, + bias=bias, + tp_group=tp_group, + ) + elif provider == "liger-cutile": + if not _CUTILE_AVAILABLE: + raise RuntimeError("provider 'liger-cutile' requires the cuda-tile package.") + forward = lambda: cutile_megatron_fused_linear_cross_entropy( + hidden, + weight, + target, + bias=bias, + tp_group=tp_group, + ) + else: + if tp_size > 1: + hidden.register_hook(lambda grad: _all_reduce_hidden_grad(grad, tp_group)) + + if provider == "megatron-core": + if not _MEGATRON_CORE_AVAILABLE: + raise RuntimeError("provider 'megatron-core' requires the megatron-core package.") + ce_forward = lambda logits: fused_vocab_parallel_cross_entropy(logits, target, tp_group) + elif provider == "megatron-compatible": + ce = LigerMegatronCrossEntropy() + ce_forward = lambda logits: ce(logits, target, tp_group=tp_group) + else: + raise ValueError(f"unknown provider: {provider!r}") + + forward = lambda: ce_forward(F.linear(hidden, weight, bias)) + + return _ProviderState(hidden, weight, bias, target, forward) + + +def _synchronized_elapsed_ms(step, tp_group, iterations: int) -> float: + dist.barrier(group=tp_group) + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + step() + end.record() + torch.cuda.synchronize() + elapsed = torch.tensor(start.elapsed_time(end) / iterations, device="cuda") + dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=tp_group) + return float(elapsed) + + +def _synchronized_backward_ms(state: _ProviderState, tp_group, iterations: int) -> float: + dist.barrier(group=tp_group) + torch.cuda.synchronize() + event_pairs = [] + for _ in range(iterations): + state.clear_grads() + loss = state.forward() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + loss.backward(torch.ones_like(loss)) + end.record() + event_pairs.append((start, end)) + torch.cuda.synchronize() + elapsed = torch.tensor( + sum(start.elapsed_time(end) for start, end in event_pairs) / iterations, + device="cuda", + ) + dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=tp_group) + return float(elapsed) + + +def _quantiles(samples: list[float]) -> tuple[float, float, float]: + values = torch.tensor(samples) + return ( + float(values.quantile(0.5)), + float(values.quantile(0.2)), + float(values.quantile(0.8)), + ) + + +def _speed( + state: _ProviderState, + tp_group, + warmup_iterations: int, + measure_iterations: int, +) -> dict[str, tuple[float, float, float]]: + def forward_step(): + state.forward() + + def full_step(): + state.clear_grads() + loss = state.forward() + loss.backward(torch.ones_like(loss)) + + for _ in range(warmup_iterations): + full_step() + torch.cuda.synchronize() + + forward_samples = [ + _synchronized_elapsed_ms(forward_step, tp_group, measure_iterations) for _ in range(_SPEED_SAMPLES) + ] + backward_samples = [_synchronized_backward_ms(state, tp_group, measure_iterations) for _ in range(_SPEED_SAMPLES)] + full_samples = [_synchronized_elapsed_ms(full_step, tp_group, measure_iterations) for _ in range(_SPEED_SAMPLES)] + return { + "forward": _quantiles(forward_samples), + "backward": _quantiles(backward_samples), + "full": _quantiles(full_samples), + } + + +def _memory(state: _ProviderState, tp_group) -> tuple[float, float, float]: + def full_step(): + state.clear_grads() + loss = state.forward() + loss.backward(torch.ones_like(loss)) + + full_step() + torch.cuda.synchronize() + samples = [] + for _ in range(_MEMORY_SAMPLES): + state.clear_grads() + torch.cuda.reset_peak_memory_stats() + full_step() + torch.cuda.synchronize() + peak = torch.tensor(torch.cuda.max_memory_allocated() / 2**20, device="cuda") + dist.all_reduce(peak, op=dist.ReduceOp.MAX, group=tp_group) + samples.append(float(peak)) + return _quantiles(samples) + + +def _make_masters( + rank: int, + tp_size: int, + num_tokens: int, + hidden_size: int, + vocab_global: int, + dtype: torch.dtype, + with_bias: bool, + device: torch.device, +): + if vocab_global % tp_size: + raise ValueError(f"vocab size {vocab_global} must be divisible by TP={tp_size}.") + vocab_local = vocab_global // tp_size + generator = torch.Generator(device=device) + generator.manual_seed(17) + hidden = torch.randn( + num_tokens, + 1, + hidden_size, + device=device, + dtype=dtype, + generator=generator, + ) + target = torch.randint( + vocab_global, + (num_tokens, 1), + device=device, + dtype=torch.long, + generator=generator, + ) + dist.broadcast(hidden, src=0) + dist.broadcast(target, src=0) + + generator.manual_seed(1000 + rank) + weight = torch.randn( + vocab_local, + hidden_size, + device=device, + dtype=dtype, + generator=generator, + ) + bias = torch.randn(vocab_local, device=device, dtype=dtype, generator=generator) if with_bias else None + return hidden, weight, bias, target + + +def _check_correctness( + rank: int, + tp_size: int, + tp_group, + dtype: torch.dtype, + device: torch.device, + providers, +): + hidden, weight, bias, target = _make_masters( + rank, + tp_size, + num_tokens=32, + hidden_size=256, + vocab_global=1024, + dtype=dtype, + with_bias=True, + device=device, + ) + weight.mul_(0.02) + bias.mul_(0.02) + upstream = torch.randn_like(target, dtype=torch.float32) + dist.broadcast(upstream, src=0) + outputs = {} + correctness_providers = ["megatron-compatible"] + correctness_providers.extend( + provider for provider in ("liger", "liger-triton", "liger-cutile") if provider in providers + ) + for provider in correctness_providers: + state = _make_state(provider, hidden, weight, bias, target, tp_group, tp_size) + loss = state.forward() + loss.backward(upstream) + outputs[provider] = ( + loss.detach().float(), + state.hidden.grad.detach().float(), + state.weight.grad.detach().float(), + state.bias.grad.detach().float(), + ) + + reference = outputs["megatron-compatible"] + names = ("loss", "grad_hidden", "grad_weight", "grad_bias") + for provider in correctness_providers[1:]: + actual = outputs[provider] + for name, actual_tensor, reference_tensor in zip(names, actual, reference): + torch.testing.assert_close( + actual_tensor, + reference_tensor, + atol=5e-3, + rtol=5e-2, + msg=f"{provider}: {name}", + ) + + +def _worker( + rank, + tp_size, + providers, + token_counts, + vocab_sizes, + hidden_size, + dtype_name, + with_bias, + warmup_iterations, + measure_iterations, + rendezvous, + result_path, + overwrite, +): + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29500") + dist.init_process_group( + backend="nccl", + init_method=f"file://{rendezvous}", + rank=rank, + world_size=tp_size, + ) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + tp_group = dist.group.WORLD + dtype = _DTYPES[dtype_name] + + _check_correctness(rank, tp_size, tp_group, dtype, device, providers) + if rank == 0: + print("Correctness: loss and gradients match the materialized reference.", flush=True) + + grouped_speed = {} if rank == 0 else None + grouped_memory = {} if rank == 0 else None + for num_tokens in token_counts: + for vocab_global in vocab_sizes: + masters = _make_masters( + rank, + tp_size, + num_tokens, + hidden_size, + vocab_global, + dtype, + with_bias, + device, + ) + for provider in providers: + state = _make_state(provider, *masters, tp_group, tp_size) + speed = _speed(state, tp_group, warmup_iterations, measure_iterations) + memory = _memory(state, tp_group) + if rank == 0: + for mode, (p50, p20, p80) in speed.items(): + grouped_speed.setdefault((provider, mode, num_tokens), []).append((vocab_global, p50, p20, p80)) + print( + f"[speed] {provider:>21s} TP={tp_size} BT={num_tokens:>5d} " + f"V={vocab_global:>6d} {mode:>8s}: {p50:.4f} ms", + flush=True, + ) + grouped_memory.setdefault((provider, "full", num_tokens), []).append((vocab_global, *memory)) + print( + f"[memory] {provider:>21s} TP={tp_size} BT={num_tokens:>5d} " + f"V={vocab_global:>6d}: {memory[0]:.1f} MB", + flush=True, + ) + del state + torch.cuda.empty_cache() + dist.barrier(group=tp_group) + del masters + + if rank == 0: + timestamp = get_formatted_time() + gpu_name = get_gpu_name() + rows = [] + common = { + "kernel_name": "megatron_fused_linear_cross_entropy", + "gpu_name": gpu_name, + "x_name": "V", + "x_label": "global vocab size", + "timestamp": timestamp, + } + for (provider, mode, num_tokens), samples in grouped_speed.items(): + samples.sort() + rows.append( + BenchmarkData( + kernel_provider=provider, + metric_name="speed", + metric_unit="ms", + x_values=[row[0] for row in samples], + y_values_50=[row[1] for row in samples], + y_values_20=[row[2] for row in samples], + y_values_80=[row[3] for row in samples], + kernel_operation_mode=mode, + extra_benchmark_config_str=( + f'{{"BT": {num_tokens}, "H": {hidden_size}, "TP": {tp_size}, ' + f'"dtype": "{dtype_name}", "bias": {str(with_bias).lower()}}}' + ), + **common, + ) + ) + for (provider, mode, num_tokens), samples in grouped_memory.items(): + samples.sort() + rows.append( + BenchmarkData( + kernel_provider=provider, + metric_name="memory", + metric_unit="MB", + x_values=[row[0] for row in samples], + y_values_50=[row[1] for row in samples], + y_values_20=[row[2] for row in samples], + y_values_80=[row[3] for row in samples], + kernel_operation_mode=mode, + extra_benchmark_config_str=( + f'{{"BT": {num_tokens}, "H": {hidden_size}, "TP": {tp_size}, ' + f'"dtype": "{dtype_name}", "bias": {str(with_bias).lower()}}}' + ), + **common, + ) + ) + update_benchmark_data_csv(rows, filename=result_path, overwrite=overwrite) + + dist.destroy_process_group() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tp-size", type=int, default=1) + parser.add_argument("--token-counts", type=int, nargs="+", default=[512, 2048]) + parser.add_argument("--hidden-size", type=int, default=4096) + parser.add_argument("--vocab-sizes", type=int, nargs="+", default=[32000, 128256]) + parser.add_argument("--dtype", choices=sorted(_DTYPES), default="bf16") + parser.add_argument("--with-bias", action="store_true") + parser.add_argument("--warmup-iterations", type=int, default=3) + parser.add_argument("--measure-iterations", type=int, default=10) + parser.add_argument( + "--providers", + nargs="+", + choices=["liger", "liger-triton", "liger-cutile", "megatron-compatible", "megatron-core"], + ) + parser.add_argument( + "--output", + type=Path, + default=Path(__file__).resolve().parents[1] / "data" / "all_benchmark_data_megatron_flce.csv", + ) + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + + if args.tp_size > torch.cuda.device_count(): + raise RuntimeError(f"--tp-size={args.tp_size} requires {args.tp_size} GPUs; found {torch.cuda.device_count()}.") + providers = args.providers or ["megatron-compatible", "liger"] + if _MEGATRON_CORE_AVAILABLE and args.providers is None: + providers.insert(1, "megatron-core") + if "megatron-core" in providers and not _MEGATRON_CORE_AVAILABLE: + raise RuntimeError("provider 'megatron-core' requested, but megatron-core is not installed.") + if "liger-cutile" in providers and not _CUTILE_AVAILABLE: + raise RuntimeError("provider 'liger-cutile' requested, but cuda-tile is not installed.") + if min(args.token_counts) <= 0 or args.hidden_size <= 0 or min(args.vocab_sizes) <= 0: + raise ValueError("token counts, hidden size, and vocabulary sizes must be positive.") + + output = args.output.resolve() + output.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile() as rendezvous: + mp.spawn( + _worker, + args=( + args.tp_size, + providers, + args.token_counts, + args.vocab_sizes, + args.hidden_size, + args.dtype, + args.with_bias, + args.warmup_iterations, + args.measure_iterations, + rendezvous.name, + str(output), + args.overwrite, + ), + nprocs=args.tp_size, + join=True, + ) + + +if __name__ == "__main__": + main() diff --git a/src/liger_kernel/megatron/__init__.py b/src/liger_kernel/megatron/__init__.py index 94f3da552..4cb498d13 100644 --- a/src/liger_kernel/megatron/__init__.py +++ b/src/liger_kernel/megatron/__init__.py @@ -10,6 +10,8 @@ experts. Mode 1 patches ``fused_bias_swiglu.SwiGLUFunction`` instead; both fall back to Megatron for FP8 input store and CPU activation offload, and neither touches the bias or MoE-routed variants. + LigerMegatronFusedLinearCrossEntropy — hidden-state-to-loss fused output + projection for tensor-parallel vocabulary shards. apply_liger_kernel_to_megatron — patches Megatron-Core so existing training scripts pick up Liger kernels with one line. Currently supports RMSNorm (via BackendSpecProvider), both the fused and unfused @@ -23,12 +25,14 @@ """ from liger_kernel.megatron.cross_entropy import LigerMegatronCrossEntropy +from liger_kernel.megatron.fused_linear_cross_entropy import LigerMegatronFusedLinearCrossEntropy from liger_kernel.megatron.monkey_patch import apply_liger_kernel_to_megatron from liger_kernel.megatron.rms_norm import LigerMegatronRMSNorm from liger_kernel.megatron.swiglu import LigerMegatronSwiGLU __all__ = [ "LigerMegatronCrossEntropy", + "LigerMegatronFusedLinearCrossEntropy", "LigerMegatronRMSNorm", "LigerMegatronSwiGLU", "apply_liger_kernel_to_megatron", diff --git a/src/liger_kernel/megatron/fused_linear_cross_entropy.py b/src/liger_kernel/megatron/fused_linear_cross_entropy.py new file mode 100644 index 000000000..ea3fd6ed7 --- /dev/null +++ b/src/liger_kernel/megatron/fused_linear_cross_entropy.py @@ -0,0 +1,43 @@ +"""Megatron-facing module for tensor-parallel fused linear cross entropy.""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from liger_kernel.ops import liger_megatron_fused_linear_cross_entropy + + +class LigerMegatronFusedLinearCrossEntropy(nn.Module): + """Fuse a vocab-sharded output projection with per-token cross entropy. + + ``hidden`` is replicated across TP ranks and ``weight`` contains the local + contiguous vocabulary shard. The output shape matches ``target``. + """ + + def __init__( + self, + ignore_index: int = -100, + ): + super().__init__() + self.ignore_index = ignore_index + + def forward( + self, + hidden: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + bias: torch.Tensor | None = None, + tp_group=None, + ) -> torch.Tensor: + return liger_megatron_fused_linear_cross_entropy( + hidden, + weight, + target, + bias=bias, + tp_group=tp_group, + ignore_index=self.ignore_index, + ) + + def extra_repr(self) -> str: + return f"ignore_index={self.ignore_index}" diff --git a/src/liger_kernel/ops/__init__.py b/src/liger_kernel/ops/__init__.py index 6bed68643..19b1f7676 100644 --- a/src/liger_kernel/ops/__init__.py +++ b/src/liger_kernel/ops/__init__.py @@ -67,6 +67,12 @@ from liger_kernel.ops.layer_norm import layer_norm_backward # noqa: F401 from liger_kernel.ops.layer_norm import layer_norm_forward # noqa: F401 from liger_kernel.ops.llama4_rope import LigerLlama4RopeFunction # noqa: F401 +from liger_kernel.ops.megatron_fused_linear_cross_entropy import ( # noqa: F401 + LigerMegatronFusedLinearCrossEntropyFunction as LigerMegatronFusedLinearCrossEntropyFunction, +) +from liger_kernel.ops.megatron_fused_linear_cross_entropy import ( # noqa: F401 + liger_megatron_fused_linear_cross_entropy as liger_megatron_fused_linear_cross_entropy, +) from liger_kernel.ops.mhc import LigerMHCCoeffsFunction # noqa: F401 from liger_kernel.ops.mhc import LigerMHCPostResFunction # noqa: F401 from liger_kernel.ops.mhc import LigerMHCPreFunction # noqa: F401 diff --git a/src/liger_kernel/ops/cutile/ops/__init__.py b/src/liger_kernel/ops/cutile/ops/__init__.py index fcbd1cc84..492a7f97e 100644 --- a/src/liger_kernel/ops/cutile/ops/__init__.py +++ b/src/liger_kernel/ops/cutile/ops/__init__.py @@ -33,6 +33,8 @@ from liger_kernel.ops.cutile.ops.layer_norm import layer_norm_backward from liger_kernel.ops.cutile.ops.layer_norm import layer_norm_forward from liger_kernel.ops.cutile.ops.llama4_rope import LigerLlama4RopeFunction +from liger_kernel.ops.cutile.ops.megatron_fused_linear_cross_entropy import LigerMegatronFusedLinearCrossEntropyFunction +from liger_kernel.ops.cutile.ops.megatron_fused_linear_cross_entropy import liger_megatron_fused_linear_cross_entropy from liger_kernel.ops.cutile.ops.multi_token_attention import LigerMultiTokenAttentionFunction from liger_kernel.ops.cutile.ops.qwen2vl_mrope import LigerQwen2VLMRopeFunction from liger_kernel.ops.cutile.ops.rope import LigerRopeFunction @@ -64,6 +66,8 @@ "layer_norm_backward", "layer_norm_forward", "LigerLlama4RopeFunction", + "LigerMegatronFusedLinearCrossEntropyFunction", + "liger_megatron_fused_linear_cross_entropy", "LigerMultiTokenAttentionFunction", "LigerQwen2VLMRopeFunction", "LigerRopeFunction", diff --git a/src/liger_kernel/ops/cutile/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/cutile/ops/megatron_fused_linear_cross_entropy.py new file mode 100644 index 000000000..2f7284a11 --- /dev/null +++ b/src/liger_kernel/ops/cutile/ops/megatron_fused_linear_cross_entropy.py @@ -0,0 +1,597 @@ +"""CuTile tensor-parallel fused linear cross entropy for Megatron.""" + +from __future__ import annotations + +import math + +import cuda.tile as ct +import torch +import torch.distributed as dist + +from liger_kernel.ops.cutile.ops.utils import _next_power_of_2 +from liger_kernel.ops.megatron_fused_linear_cross_entropy import _tp_rank_and_world + +ConstBool = ct.Constant[bool] +ConstInt = ct.Constant[int] +LOG2E = 1.4426950408889634 +MAX_ROW_BLOCK_SIZE = 4096 + + +@ct.function +def _matmul_body( + a, + b, + bias, + output, + TILE_M: ConstInt, + TILE_N: ConstInt, + TILE_K: ConstInt, + HAS_BIAS: ConstBool, + SWIZZLE: ConstBool, +): + num_m_tiles = ct.num_tiles(a, axis=0, shape=(TILE_M, TILE_K)) + num_n_tiles = ct.num_tiles(b, axis=1, shape=(TILE_K, TILE_N)) + num_k_tiles = ct.num_tiles(a, axis=1, shape=(TILE_M, TILE_K)) + block = ct.bid(0) + if SWIZZLE: + group_size_m = 8 + blocks_per_group = group_size_m * num_n_tiles + group = block // blocks_per_group + first_tile_m = group * group_size_m + active_group_size_m = min(num_m_tiles - first_tile_m, group_size_m) + tile_m = first_tile_m + (block % active_group_size_m) + tile_n = (block % blocks_per_group) // active_group_size_m + else: + tile_m = block // num_n_tiles + tile_n = block % num_n_tiles + + accumulator = ct.full((TILE_M, TILE_N), 0.0, dtype=ct.float32) + for tile_k in range(num_k_tiles): + a_tile = ct.load( + a, + index=(tile_m, tile_k), + shape=(TILE_M, TILE_K), + padding_mode=ct.PaddingMode.ZERO, + ) + b_tile = ct.load( + b, + index=(tile_k, tile_n), + shape=(TILE_K, TILE_N), + padding_mode=ct.PaddingMode.ZERO, + ) + accumulator = ct.mma(a_tile, b_tile, accumulator) + + if HAS_BIAS: + bias_tile = ct.load( + bias, + index=(tile_n,), + shape=(TILE_N,), + padding_mode=ct.PaddingMode.ZERO, + ) + accumulator = accumulator + ct.astype(bias_tile, ct.float32) + + ct.store(output, index=(tile_m, tile_n), tile=ct.astype(accumulator, output.dtype)) + + +@ct.kernel(num_ctas=1) +def _matmul_1cta_kernel( + a, + b, + bias, + output, + TILE_M: ConstInt, + TILE_N: ConstInt, + TILE_K: ConstInt, + HAS_BIAS: ConstBool, + SWIZZLE: ConstBool, +): + _matmul_body(a, b, bias, output, TILE_M, TILE_N, TILE_K, HAS_BIAS, SWIZZLE) + + +@ct.kernel(num_ctas=2) +def _matmul_2cta_kernel( + a, + b, + bias, + output, + TILE_M: ConstInt, + TILE_N: ConstInt, + TILE_K: ConstInt, + HAS_BIAS: ConstBool, + SWIZZLE: ConstBool, +): + _matmul_body(a, b, bias, output, TILE_M, TILE_N, TILE_K, HAS_BIAS, SWIZZLE) + + +@ct.kernel(occupancy=4) +def _row_max_kernel( + input, + output, + n_cols, + BLOCK_SIZE: ConstInt, +): + row = ct.bid(0) + num_chunks = (n_cols + BLOCK_SIZE - 1) // BLOCK_SIZE + row_max_tile = ct.full((1,), -math.inf, dtype=ct.float32) + + for chunk in range(num_chunks): + columns = ct.arange(BLOCK_SIZE, dtype=ct.int32) + chunk * BLOCK_SIZE + values = ct.astype( + ct.gather( + input, + (row, columns), + check_bounds=True, + padding_value=-math.inf, + latency=3, + ), + ct.float32, + ) + row_max = ct.maximum( + ct.sum(row_max_tile, 0, keepdims=False), + ct.max(values, 0, keepdims=False), + ) + row_max_tile = ct.full((1,), row_max, dtype=ct.float32) + + ct.scatter(output, row, ct.sum(row_max_tile, 0, keepdims=False)) + + +@ct.kernel(occupancy=4) +def _vocab_parallel_ce_forward_kernel( + logits, + logits_max, + target, + predicted_logit, + sum_exp, + vocab_start, + n_cols, + ignore_index, + BLOCK_SIZE: ConstInt, +): + row = ct.bid(0) + y_global = ct.load(target, row, shape=()) + maximum = ct.astype(ct.load(logits_max, row, shape=()), ct.float32) + is_ignored = y_global == ignore_index + target_off_rank = (y_global < vocab_start) or (y_global >= vocab_start + n_cols) + y_local = ct.astype(y_global - vocab_start, ct.int32) + + if is_ignored or target_off_rank: + predicted = 0.0 + else: + target_index = ct.add(ct.arange(1, dtype=ct.int32), y_local) + target_tile = ct.gather(logits, (row, target_index), check_bounds=False) + predicted = ct.sum(ct.astype(target_tile, ct.float32), 0, keepdims=False) - maximum + + num_chunks = (n_cols + BLOCK_SIZE - 1) // BLOCK_SIZE + sum_exp_tile = ct.full((1,), 0.0, dtype=ct.float32) + for chunk in range(num_chunks): + columns = ct.arange(BLOCK_SIZE, dtype=ct.int32) + chunk * BLOCK_SIZE + in_bounds = columns < n_cols + values = ct.astype( + ct.gather( + logits, + (row, columns), + check_bounds=True, + padding_value=-math.inf, + latency=3, + ), + ct.float32, + ) + exponentials = ct.exp2((values - maximum) * LOG2E, flush_to_zero=True) + exponentials = ct.where(in_bounds, exponentials, 0.0) + running_sum = ct.sum(sum_exp_tile, 0, keepdims=False) + sum_exp_tile = ct.full( + (1,), + running_sum + ct.sum(exponentials, 0, keepdims=False), + dtype=ct.float32, + ) + ct.scatter( + logits, + (row, columns), + ct.astype(exponentials, logits.dtype), + check_bounds=True, + ) + + ct.scatter(predicted_logit, row, predicted) + ct.scatter(sum_exp, row, ct.sum(sum_exp_tile, 0, keepdims=False)) + + +@ct.kernel(occupancy=4) +def _vocab_parallel_ce_backward_kernel( + exp_buffer, + sum_exp, + target, + grad_output, + vocab_start, + n_cols, + ignore_index, + BLOCK_SIZE: ConstInt, +): + row = ct.bid(0) + y_global = ct.load(target, row, shape=()) + num_chunks = (n_cols + BLOCK_SIZE - 1) // BLOCK_SIZE + + if y_global == ignore_index: + for chunk in range(num_chunks): + columns = ct.arange(BLOCK_SIZE, dtype=ct.int32) + chunk * BLOCK_SIZE + zeros = ct.full((BLOCK_SIZE,), 0.0, dtype=exp_buffer.dtype) + ct.scatter(exp_buffer, (row, columns), zeros, check_bounds=True) + return + + target_off_rank = (y_global < vocab_start) or (y_global >= vocab_start + n_cols) + y_local = ct.astype(y_global - vocab_start, ct.int32) + global_sum = ct.astype(ct.load(sum_exp, row, shape=()), ct.float32) + upstream = ct.astype(ct.load(grad_output, row, shape=()), ct.float32) + + for chunk in range(num_chunks): + columns = ct.arange(BLOCK_SIZE, dtype=ct.int32) + chunk * BLOCK_SIZE + exponentials = ct.astype( + ct.gather(exp_buffer, (row, columns), check_bounds=True, padding_value=0.0), + ct.float32, + ) + gradient = exponentials / global_sum + if not target_off_rank: + gradient = ct.where(columns == y_local, gradient - 1.0, gradient) + gradient = gradient * upstream + ct.scatter( + exp_buffer, + (row, columns), + ct.astype(gradient, exp_buffer.dtype), + check_bounds=True, + ) + + +@ct.kernel(occupancy=4) +def _loss_kernel( + sum_exp, + predicted_logit, + target, + output, + ignore_index, +): + row = ct.bid(0) + y_global = ct.load(target, row, shape=()) + if y_global == ignore_index: + loss = 0.0 + else: + denominator = ct.astype(ct.load(sum_exp, row, shape=()), ct.float32) + predicted = ct.astype(ct.load(predicted_logit, row, shape=()), ct.float32) + loss = ct.log(denominator) - predicted + ct.scatter(output, row, loss) + + +@ct.kernel(occupancy=4) +def _column_sum_kernel( + input, + output, + n_rows, + BLOCK_SIZE: ConstInt, +): + column = ct.bid(0) + num_chunks = (n_rows + BLOCK_SIZE - 1) // BLOCK_SIZE + total_tile = ct.full((1,), 0.0, dtype=ct.float32) + + for chunk in range(num_chunks): + rows = ct.arange(BLOCK_SIZE, dtype=ct.int32) + chunk * BLOCK_SIZE + values = ct.astype( + ct.gather(input, (rows, column), check_bounds=True, padding_value=0.0), + ct.float32, + ) + running_total = ct.sum(total_tile, 0, keepdims=False) + total_tile = ct.full( + (1,), + running_total + ct.sum(values, 0, keepdims=False), + dtype=ct.float32, + ) + + ct.scatter(output, column, ct.astype(ct.sum(total_tile, 0, keepdims=False), output.dtype)) + + +def _select_row_block_size(size: int) -> int: + return min(MAX_ROW_BLOCK_SIZE, _next_power_of_2(size)) + + +def _cutile_matmul( + a: torch.Tensor, + b: torch.Tensor, + *, + operation: str, + bias: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, +) -> torch.Tensor: + if a.ndim != 2 or b.ndim != 2 or a.shape[1] != b.shape[0]: + raise ValueError(f"matmul expects [M, K] @ [K, N], got {tuple(a.shape)} and {tuple(b.shape)}.") + + output = torch.empty( + (a.shape[0], b.shape[1]), + device=a.device, + dtype=output_dtype or a.dtype, + ) + if operation == "projection": + kernel, tile = ( + (_matmul_2cta_kernel, (256, 256, 128)) if a.shape[0] <= 1024 else (_matmul_2cta_kernel, (512, 256, 64)) + ) + elif operation == "dx": + if a.shape[0] <= 1024: + kernel, tile = _matmul_1cta_kernel, (128, 128, 64) + elif a.shape[1] > 16000 or a.shape[0] > 16384: + kernel, tile = _matmul_2cta_kernel, (512, 256, 64) + else: + kernel, tile = _matmul_1cta_kernel, (256, 256, 64) + elif operation == "dw": + if a.shape[1] >= 16384 and (a.shape[0] > 16000 or a.shape[1] == 16384): + kernel, tile = _matmul_2cta_kernel, (512, 256, 64) + else: + use_single_cta = a.shape[1] <= 1024 or a.shape[0] > 16000 or a.shape[1] > 16384 + kernel = _matmul_1cta_kernel if use_single_cta else _matmul_2cta_kernel + tile = (256, 256, 64) + else: + raise ValueError(f"unknown FLCE GEMM operation: {operation!r}.") + + tile_m, tile_n, tile_k = tile + swizzle = operation == "projection" and b.shape[1] > 16000 + grid = ( + ct.cdiv(a.shape[0], tile_m) * ct.cdiv(b.shape[1], tile_n), + 1, + 1, + ) + ct.launch( + torch.cuda.current_stream(), + grid, + kernel, + ( + a, + b, + bias if bias is not None else output, + output, + tile_m, + tile_n, + tile_k, + bias is not None, + swizzle, + ), + ) + return output + + +def _cutile_row_max(input: torch.Tensor) -> torch.Tensor: + output = torch.empty(input.shape[0], device=input.device, dtype=torch.float32) + block_size = 16384 if input.shape[1] > 16384 else _select_row_block_size(input.shape[1]) + ct.launch( + torch.cuda.current_stream(), + (input.shape[0], 1, 1), + _row_max_kernel, + (input, output, int(input.shape[1]), int(block_size)), + ) + return output + + +def _cutile_ce_forward( + logits: torch.Tensor, + logits_max: torch.Tensor, + target: torch.Tensor, + vocab_start: int, + ignore_index: int, +) -> tuple[torch.Tensor, torch.Tensor]: + rows, vocab_local = logits.shape + stats = torch.empty((2, rows), device=logits.device, dtype=torch.float32) + predicted_logit = stats[0] + sum_exp = stats[1] + block_size = 16384 if vocab_local > 16384 else _select_row_block_size(vocab_local) + ct.launch( + torch.cuda.current_stream(), + (rows, 1, 1), + _vocab_parallel_ce_forward_kernel, + ( + logits, + logits_max, + target, + predicted_logit, + sum_exp, + int(vocab_start), + int(vocab_local), + int(ignore_index), + int(block_size), + ), + ) + return logits, stats + + +def _cutile_ce_backward( + exp_buffer: torch.Tensor, + sum_exp: torch.Tensor, + target: torch.Tensor, + grad_output: torch.Tensor, + vocab_start: int, + ignore_index: int, +) -> None: + block_size = min(2048, _select_row_block_size(exp_buffer.shape[1])) + ct.launch( + torch.cuda.current_stream(), + (exp_buffer.shape[0], 1, 1), + _vocab_parallel_ce_backward_kernel, + ( + exp_buffer, + sum_exp, + target, + grad_output, + int(vocab_start), + int(exp_buffer.shape[1]), + int(ignore_index), + int(block_size), + ), + ) + + +def _cutile_loss( + sum_exp: torch.Tensor, + predicted_logit: torch.Tensor, + target: torch.Tensor, + ignore_index: int, +) -> torch.Tensor: + output = torch.empty_like(sum_exp) + ct.launch( + torch.cuda.current_stream(), + (target.numel(), 1, 1), + _loss_kernel, + (sum_exp, predicted_logit, target, output, int(ignore_index)), + ) + return output + + +def _cutile_column_sum(input: torch.Tensor, output_dtype: torch.dtype) -> torch.Tensor: + output = torch.empty(input.shape[1], device=input.device, dtype=output_dtype) + block_size = _select_row_block_size(input.shape[0]) + ct.launch( + torch.cuda.current_stream(), + (input.shape[1], 1, 1), + _column_sum_kernel, + (input, output, int(input.shape[0]), int(block_size)), + ) + return output + + +def _materialized_backward(ctx, grad_output: torch.Tensor): + hidden, weight, exp_buffer, sum_exp, target = ctx.saved_tensors + grad_output_1d = grad_output.contiguous().reshape(-1).float() + _cutile_ce_backward( + exp_buffer, + sum_exp, + target, + grad_output_1d, + ctx.vocab_start, + ctx.ignore_index, + ) + + grad_hidden = _cutile_matmul( + exp_buffer, + weight, + operation="dx", + output_dtype=torch.float32 if exp_buffer.shape[0] <= 1024 else None, + ) + reduce_work = ( + dist.all_reduce( + grad_hidden, + op=dist.ReduceOp.SUM, + group=ctx.tp_group, + async_op=True, + ) + if ctx.tp_world > 1 + else None + ) + grad_weight = _cutile_matmul(exp_buffer.t(), hidden, operation="dw") + grad_bias = _cutile_column_sum(exp_buffer, ctx.bias_dtype) if ctx.has_bias else None + + if reduce_work is not None: + reduce_work.wait() + grad_hidden = grad_hidden.to(ctx.hidden_dtype).reshape(ctx.original_hidden_shape) + return grad_hidden, grad_weight, grad_bias + + +class LigerMegatronFusedLinearCrossEntropyFunction(torch.autograd.Function): + """Hidden-to-loss tensor-parallel FLCE using CuTile local kernels.""" + + @staticmethod + def forward( + ctx, + hidden: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + bias: torch.Tensor | None, + tp_group, + ignore_index: int, + ) -> torch.Tensor: + if hidden.ndim < 2: + raise ValueError(f"hidden must have at least 2 dimensions, got shape {tuple(hidden.shape)}.") + if weight.ndim != 2: + raise ValueError(f"weight must be 2-D [V_local, H], got shape {tuple(weight.shape)}.") + if tuple(target.shape) != tuple(hidden.shape[:-1]): + raise ValueError( + f"target shape must equal hidden.shape[:-1]; got target={tuple(target.shape)}, " + f"hidden={tuple(hidden.shape)}." + ) + if hidden.shape[-1] != weight.shape[1]: + raise ValueError(f"hidden size mismatch: hidden has H={hidden.shape[-1]}, weight has H={weight.shape[1]}.") + if hidden.dtype != weight.dtype: + raise TypeError(f"hidden and weight must have the same dtype, got {hidden.dtype} and {weight.dtype}.") + if hidden.device != weight.device or hidden.device != target.device: + raise ValueError("hidden, weight, and target must be on the same device.") + if bias is not None: + if bias.ndim != 1 or bias.shape[0] != weight.shape[0]: + raise ValueError(f"bias must have shape ({weight.shape[0]},), got {tuple(bias.shape)}.") + if bias.device != hidden.device or bias.dtype != hidden.dtype: + raise TypeError("bias must have the same device and dtype as hidden.") + if hidden.device.type != "cuda" or hidden.dtype not in (torch.bfloat16, torch.float16): + raise RuntimeError("CuTile Megatron FLCE requires a CUDA GPU and float16 or bfloat16 inputs.") + + tp_rank, tp_world = _tp_rank_and_world(tp_group) + vocab_local = weight.shape[0] + vocab_global = vocab_local * tp_world + vocab_start = tp_rank * vocab_local + + flat_target = target.reshape(-1).to(torch.int64).contiguous() + valid = flat_target != ignore_index + invalid = valid & ((flat_target < 0) | (flat_target >= vocab_global)) + valid_targets = ~torch.any(invalid) + if hasattr(torch, "_assert_async"): + torch._assert_async(valid_targets, f"non-ignored targets must be in [0, {vocab_global}).") + elif not valid_targets.item(): + raise ValueError(f"non-ignored targets must be in [0, {vocab_global}).") + + original_hidden_shape = hidden.shape + hidden_2d = hidden.reshape(-1, hidden.shape[-1]).contiguous() + weight_2d = weight.contiguous() + bias_1d = bias.contiguous() if bias is not None else None + + logits = _cutile_matmul(hidden_2d, weight_2d.t(), operation="projection", bias=bias_1d) + logits_max = _cutile_row_max(logits) + if tp_world > 1: + dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group) + + exp_buffer, stats = _cutile_ce_forward( + logits, + logits_max, + flat_target, + vocab_start, + ignore_index, + ) + if tp_world > 1: + dist.all_reduce(stats, op=dist.ReduceOp.SUM, group=tp_group) + predicted_logit = stats[0] + sum_exp = stats[1] + + loss = _cutile_loss(sum_exp, predicted_logit, flat_target, ignore_index) + + ctx.save_for_backward(hidden_2d, weight_2d, exp_buffer, sum_exp, flat_target) + ctx.has_bias = bias is not None + ctx.bias_dtype = bias.dtype if bias is not None else None + ctx.tp_group = tp_group + ctx.tp_world = tp_world + ctx.vocab_start = vocab_start + ctx.ignore_index = ignore_index + ctx.original_hidden_shape = original_hidden_shape + ctx.hidden_dtype = hidden.dtype + return loss.reshape(target.shape) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + grad_hidden, grad_weight, grad_bias = _materialized_backward(ctx, grad_output) + return grad_hidden, grad_weight, None, grad_bias, None, None + + +def liger_megatron_fused_linear_cross_entropy( + hidden: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + bias: torch.Tensor | None = None, + tp_group=None, + ignore_index: int = -100, +) -> torch.Tensor: + """Compute Megatron FLCE with CuTile local kernels and NCCL TP collectives.""" + return LigerMegatronFusedLinearCrossEntropyFunction.apply( + hidden, + weight, + target, + bias, + tp_group, + ignore_index, + ) diff --git a/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py new file mode 100644 index 000000000..c28e2cad4 --- /dev/null +++ b/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py @@ -0,0 +1,200 @@ +"""Materialized tensor-parallel fused linear cross entropy for Megatron. + +Each tensor-parallel rank owns a contiguous vocabulary shard. Forward performs +one local projection GEMM, computes globally normalized cross entropy, and saves +shifted exponentials in the projection dtype. Backward converts that buffer to +dlogits in-place, avoiding projection recomputation before forming dX and dW. +""" + +from __future__ import annotations + +import operator + +import torch +import torch.distributed as dist + +from liger_kernel.ops.utils import compare_version + +_SUPPORTS_OUT_DTYPE = compare_version("torch", operator.ge, "2.8.0") + + +def _tp_rank_and_world(tp_group) -> tuple[int, int]: + if tp_group is None: + return 0, 1 + world = dist.get_world_size(tp_group) + if world == 1: + return 0, 1 + return dist.get_rank(tp_group), world + + +def _materialized_backward(ctx, grad_output: torch.Tensor): + """Convert saved CE state to dlogits and form projection gradients.""" + from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps + from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_backward_kernel + + hidden, weight, exp_buf, sum_exp_global, target = ctx.saved_tensors + grad_out = grad_output.contiguous().reshape(-1).float() + num_warps = _get_num_warps(ctx.ce_block_size) + liger_vocab_parallel_ce_backward_kernel[(hidden.shape[0],)]( + EXP_ptr=exp_buf, + EXP_stride=exp_buf.stride(0), + sum_exp_ptr=sum_exp_global, + Y_ptr=target, + grad_out_ptr=grad_out, + vocab_start=ctx.vocab_start, + n_cols=weight.shape[0], + ignore_index=ctx.ignore_index, + alpha_eff=0.0, + eps_eff=0.0, + HAS_LABEL_SMOOTHING=False, + BLOCK_SIZE=ctx.ce_block_size, + num_warps=num_warps, + ) + + if _SUPPORTS_OUT_DTYPE: + grad_hidden = torch.mm(exp_buf, weight, out_dtype=torch.float32) + else: + grad_hidden = exp_buf.float() @ weight.float() + grad_weight = exp_buf.t() @ hidden + grad_bias = exp_buf.sum(dim=0, dtype=torch.float32).to(ctx.bias_dtype) if ctx.has_bias else None + + if ctx.tp_world > 1: + dist.all_reduce(grad_hidden, op=dist.ReduceOp.SUM, group=ctx.tp_group) + grad_hidden = grad_hidden.to(ctx.hidden_dtype).reshape(ctx.original_hidden_shape) + return grad_hidden, grad_weight, grad_bias + + +class LigerMegatronFusedLinearCrossEntropyFunction(torch.autograd.Function): + """Hidden-to-loss tensor-parallel FLCE with saved low-precision CE state.""" + + @staticmethod + def forward( + ctx, + hidden: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + bias: torch.Tensor | None, + tp_group, + ignore_index: int, + ) -> torch.Tensor: + if hidden.ndim < 2: + raise ValueError(f"hidden must have at least 2 dimensions, got shape {tuple(hidden.shape)}.") + if weight.ndim != 2: + raise ValueError(f"weight must be 2-D [V_local, H], got shape {tuple(weight.shape)}.") + if tuple(target.shape) != tuple(hidden.shape[:-1]): + raise ValueError( + f"target shape must equal hidden.shape[:-1]; got target={tuple(target.shape)}, " + f"hidden={tuple(hidden.shape)}." + ) + if hidden.shape[-1] != weight.shape[1]: + raise ValueError(f"hidden size mismatch: hidden has H={hidden.shape[-1]}, weight has H={weight.shape[1]}.") + if hidden.dtype != weight.dtype: + raise TypeError(f"hidden and weight must have the same dtype, got {hidden.dtype} and {weight.dtype}.") + if hidden.device != weight.device or hidden.device != target.device: + raise ValueError("hidden, weight, and target must be on the same device.") + if bias is not None: + if bias.ndim != 1 or bias.shape[0] != weight.shape[0]: + raise ValueError(f"bias must have shape ({weight.shape[0]},), got {tuple(bias.shape)}.") + if bias.device != hidden.device or bias.dtype != hidden.dtype: + raise TypeError("bias must have the same device and dtype as hidden.") + if hidden.device.type != "cuda" or hidden.dtype not in (torch.bfloat16, torch.float16): + raise RuntimeError("Megatron FLCE requires a CUDA GPU and float16 or bfloat16 inputs.") + + tp_rank, tp_world = _tp_rank_and_world(tp_group) + vocab_local = weight.shape[0] + vocab_global = vocab_local * tp_world + vocab_start = tp_rank * vocab_local + + flat_target = target.reshape(-1).to(torch.int64).contiguous() + valid = flat_target != ignore_index + invalid = valid & ((flat_target < 0) | (flat_target >= vocab_global)) + valid_targets = ~torch.any(invalid) + if hasattr(torch, "_assert_async"): + torch._assert_async(valid_targets, f"non-ignored targets must be in [0, {vocab_global}).") + elif not valid_targets.item(): + raise ValueError(f"non-ignored targets must be in [0, {vocab_global}).") + + original_hidden_shape = hidden.shape + hidden_2d = hidden.reshape(-1, hidden.shape[-1]).contiguous() + weight_2d = weight.contiguous() + bias_1d = bias.contiguous() if bias is not None else None + + logits = torch.mm(hidden_2d, weight_2d.t()) + if bias_1d is not None: + logits.add_(bias_1d) + + logits_max = logits.amax(dim=-1).float() + if tp_world > 1: + dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group) + + from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps + from liger_kernel.ops.vocab_parallel_cross_entropy import _select_block_size + from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_forward_kernel + + exp_buf = torch.empty( + hidden_2d.shape[0], + vocab_local, + device=hidden.device, + dtype=hidden.dtype, + ) + predicted_logit = torch.empty(hidden_2d.shape[0], device=hidden.device, dtype=torch.float32) + sum_exp = torch.empty_like(predicted_logit) + ce_block_size = _select_block_size(vocab_local) + num_warps = _get_num_warps(ce_block_size) + liger_vocab_parallel_ce_forward_kernel[(hidden_2d.shape[0],)]( + X_ptr=logits, + X_stride=logits.stride(0), + EXP_ptr=exp_buf, + EXP_stride=exp_buf.stride(0), + logits_max_ptr=logits_max, + Y_ptr=flat_target, + pred_ptr=predicted_logit, + sum_exp_ptr=sum_exp, + vocab_start=vocab_start, + n_cols=vocab_local, + ignore_index=ignore_index, + BLOCK_SIZE=ce_block_size, + num_warps=num_warps, + ) + if tp_world > 1: + dist.all_reduce(predicted_logit, op=dist.ReduceOp.SUM, group=tp_group) + dist.all_reduce(sum_exp, op=dist.ReduceOp.SUM, group=tp_group) + + loss = torch.log(sum_exp) - predicted_logit + loss = torch.where(valid, loss, torch.zeros_like(loss)) + + ctx.save_for_backward(hidden_2d, weight_2d, exp_buf, sum_exp, flat_target) + ctx.has_bias = bias is not None + ctx.bias_dtype = bias.dtype if bias is not None else None + ctx.tp_group = tp_group + ctx.tp_world = tp_world + ctx.vocab_start = vocab_start + ctx.ignore_index = ignore_index + ctx.ce_block_size = ce_block_size + ctx.original_hidden_shape = original_hidden_shape + ctx.hidden_dtype = hidden.dtype + return loss.reshape(target.shape) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + grad_hidden, grad_weight, grad_bias = _materialized_backward(ctx, grad_output) + return grad_hidden, grad_weight, None, grad_bias, None, None + + +def liger_megatron_fused_linear_cross_entropy( + hidden: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + bias: torch.Tensor | None = None, + tp_group=None, + ignore_index: int = -100, +) -> torch.Tensor: + """Compute per-token loss from replicated hidden states and a local vocab shard.""" + return LigerMegatronFusedLinearCrossEntropyFunction.apply( + hidden, + weight, + target, + bias, + tp_group, + ignore_index, + ) diff --git a/src/liger_kernel/ops/triton/__init__.py b/src/liger_kernel/ops/triton/__init__.py new file mode 100644 index 000000000..0de2a07af --- /dev/null +++ b/src/liger_kernel/ops/triton/__init__.py @@ -0,0 +1,12 @@ +"""Opt-in portable Triton implementations.""" + +from liger_kernel.ops.backends.registry import ImplInfo +from liger_kernel.ops.backends.registry import register_impl + +register_impl( + ImplInfo( + name="triton", + devices=("cuda",), + module_path=f"{__name__}.ops", + ) +) diff --git a/src/liger_kernel/ops/triton/ops/__init__.py b/src/liger_kernel/ops/triton/ops/__init__.py new file mode 100644 index 000000000..f757b8cf8 --- /dev/null +++ b/src/liger_kernel/ops/triton/ops/__init__.py @@ -0,0 +1,9 @@ +"""Operators replaced by the opt-in all-Triton implementation.""" + +from liger_kernel.ops.triton.ops.megatron_fused_linear_cross_entropy import LigerMegatronFusedLinearCrossEntropyFunction +from liger_kernel.ops.triton.ops.megatron_fused_linear_cross_entropy import liger_megatron_fused_linear_cross_entropy + +__all__ = [ + "LigerMegatronFusedLinearCrossEntropyFunction", + "liger_megatron_fused_linear_cross_entropy", +] diff --git a/src/liger_kernel/ops/triton/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/triton/ops/megatron_fused_linear_cross_entropy.py new file mode 100644 index 000000000..5ffa8947e --- /dev/null +++ b/src/liger_kernel/ops/triton/ops/megatron_fused_linear_cross_entropy.py @@ -0,0 +1,600 @@ +"""Portable all-Triton tensor-parallel fused linear cross entropy for Megatron. + +Each tensor-parallel rank owns a contiguous vocabulary shard. Forward performs +one Triton projection GEMM, computes globally normalized cross entropy, and +saves shifted exponentials in the projection dtype. Backward converts that +buffer to dlogits in-place before Triton dX and dW GEMMs. Tensor-parallel +collectives remain NCCL/RCCL calls between architecture-independent kernels. +""" + +from __future__ import annotations + +import torch +import torch.distributed as dist +import triton +import triton.language as tl + + +def _matmul_autotune_configs(): + return [ + triton.Config( + {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 8}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, + num_stages=3, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 4}, + num_stages=3, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 4}, + num_stages=3, + num_warps=4, + ), + ] + + +def _split_k_matmul_autotune_configs(): + return [ + triton.Config( + { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 2, + }, + num_stages=3, + num_warps=4, + ), + triton.Config( + { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 4, + }, + num_stages=3, + num_warps=4, + ), + triton.Config( + { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 8, + }, + num_stages=3, + num_warps=4, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 4, + }, + num_stages=3, + num_warps=8, + ), + triton.Config( + { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 4, + }, + num_stages=3, + num_warps=8, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 4, + }, + num_stages=3, + num_warps=8, + ), + ] + + +@triton.autotune(configs=_matmul_autotune_configs(), key=["M", "N", "K"]) +@triton.jit +def _matmul_kernel( + a_ptr, + b_ptr, + bias_ptr, + output_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_om, + stride_on, + HAS_BIAS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % num_pid_in_group) % group_size_m + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k_start in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + k_remaining = K - k_start * BLOCK_SIZE_K + a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < k_remaining), other=0.0) + b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_remaining) & (offs_n[None, :] < N), other=0.0) + accumulator = tl.dot(a, b, accumulator) + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if HAS_BIAS: + bias = tl.load(bias_ptr + offs_n, mask=offs_n < N, other=0.0).to(tl.float32) + accumulator += bias[None, :] + + output_ptrs = output_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on + tl.store(output_ptrs, accumulator, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N)) + + +@triton.autotune( + configs=_split_k_matmul_autotune_configs(), + key=["M", "N", "K"], + reset_to_zero=["output_ptr"], +) +@triton.jit +def _split_k_matmul_kernel( + a_ptr, + b_ptr, + output_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_om, + stride_on, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + SPLIT_K: tl.constexpr, +): + pid = tl.program_id(0) + split_k_id = tl.program_id(1) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % num_pid_in_group) % group_size_m + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = split_k_id * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k_start in range(0, tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)): + k_remaining = K - k_start * BLOCK_SIZE_K * SPLIT_K + a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < k_remaining), other=0.0) + b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_remaining) & (offs_n[None, :] < N), other=0.0) + accumulator = tl.dot(a, b, accumulator) + a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak + b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk + + output_ptrs = output_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on + tl.atomic_add(output_ptrs, accumulator, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N)) + + +@triton.jit +def _row_max_kernel( + input_ptr, + output_ptr, + n_cols, + input_row_stride, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + row_ptr = input_ptr + row * input_row_stride + row_max = -float("inf") + for start in range(0, n_cols, BLOCK_SIZE): + offsets = start + tl.arange(0, BLOCK_SIZE) + values = tl.load(row_ptr + offsets, mask=offsets < n_cols, other=-float("inf")).to(tl.float32) + row_max = tl.maximum(row_max, tl.max(values)) + tl.store(output_ptr + row, row_max) + + +@triton.jit +def _loss_kernel( + sum_exp_ptr, + predicted_logit_ptr, + target_ptr, + loss_ptr, + n_rows, + ignore_index, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0).to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_rows + sum_exp = tl.load(sum_exp_ptr + offsets, mask=mask, other=1.0) + predicted_logit = tl.load(predicted_logit_ptr + offsets, mask=mask, other=0.0) + target = tl.load(target_ptr + offsets, mask=mask, other=ignore_index) + loss = tl.log(sum_exp) - predicted_logit + loss = tl.where(target == ignore_index, 0.0, loss) + tl.store(loss_ptr + offsets, loss, mask=mask) + + +@triton.jit +def _column_sum_kernel( + input_ptr, + output_ptr, + n_rows, + input_row_stride, + BLOCK_SIZE: tl.constexpr, +): + col = tl.program_id(0).to(tl.int64) + offsets = tl.arange(0, BLOCK_SIZE) + total = 0.0 + for start in range(0, n_rows, BLOCK_SIZE): + rows = start + offsets + values = tl.load(input_ptr + rows * input_row_stride + col, mask=rows < n_rows, other=0.0) + total += tl.sum(values.to(tl.float32)) + tl.store(output_ptr + col, total) + + +def _triton_matmul( + a: torch.Tensor, + b: torch.Tensor, + *, + bias: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, +) -> torch.Tensor: + if a.ndim != 2 or b.ndim != 2 or a.shape[1] != b.shape[0]: + raise ValueError(f"matmul expects [M, K] @ [K, N], got {tuple(a.shape)} and {tuple(b.shape)}.") + m, k = a.shape + n = b.shape[1] + output = torch.empty((m, n), device=a.device, dtype=output_dtype or a.dtype) + grid = lambda meta: (triton.cdiv(m, meta["BLOCK_SIZE_M"]) * triton.cdiv(n, meta["BLOCK_SIZE_N"]),) + _matmul_kernel[grid]( + a, + b, + bias if bias is not None else output, + output, + m, + n, + k, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + output.stride(0), + output.stride(1), + HAS_BIAS=bias is not None, + ) + return output + + +def _triton_dx_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + m, k = a.shape + n = b.shape[1] + if m > 1024 or k < 4096: + return _triton_matmul(a, b, output_dtype=torch.float32) + + output = torch.zeros((m, n), device=a.device, dtype=torch.float32) + grid = lambda meta: ( + triton.cdiv(m, meta["BLOCK_SIZE_M"]) * triton.cdiv(n, meta["BLOCK_SIZE_N"]), + meta["SPLIT_K"], + ) + _split_k_matmul_kernel[grid]( + a, + b, + output, + m, + n, + k, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + output.stride(0), + output.stride(1), + ) + return output + + +def _triton_row_max(input: torch.Tensor, block_size: int) -> torch.Tensor: + from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps + + output = torch.empty(input.shape[0], device=input.device, dtype=torch.float32) + _row_max_kernel[(input.shape[0],)]( + input, + output, + input.shape[1], + input.stride(0), + BLOCK_SIZE=block_size, + num_warps=_get_num_warps(block_size), + ) + return output + + +def _triton_loss( + sum_exp: torch.Tensor, + predicted_logit: torch.Tensor, + target: torch.Tensor, + ignore_index: int, +) -> torch.Tensor: + output = torch.empty_like(sum_exp) + block_size = 256 + _loss_kernel[(triton.cdiv(target.numel(), block_size),)]( + sum_exp, + predicted_logit, + target, + output, + target.numel(), + ignore_index, + BLOCK_SIZE=block_size, + num_warps=4, + ) + return output + + +def _triton_column_sum(input: torch.Tensor) -> torch.Tensor: + from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps + from liger_kernel.ops.vocab_parallel_cross_entropy import _select_block_size + + block_size = _select_block_size(input.shape[0]) + output = torch.empty(input.shape[1], device=input.device, dtype=torch.float32) + _column_sum_kernel[(input.shape[1],)]( + input, + output, + input.shape[0], + input.stride(0), + BLOCK_SIZE=block_size, + num_warps=_get_num_warps(block_size), + ) + return output + + +def _tp_rank_and_world(tp_group) -> tuple[int, int]: + if tp_group is None: + return 0, 1 + world = dist.get_world_size(tp_group) + if world == 1: + return 0, 1 + return dist.get_rank(tp_group), world + + +def _materialized_backward(ctx, grad_output: torch.Tensor): + """Convert saved CE state to dlogits and form projection gradients.""" + from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps + from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_backward_kernel + + hidden, weight, exp_buf, sum_exp_global, target = ctx.saved_tensors + grad_out = grad_output.contiguous().reshape(-1).float() + num_warps = _get_num_warps(ctx.ce_block_size) + liger_vocab_parallel_ce_backward_kernel[(hidden.shape[0],)]( + EXP_ptr=exp_buf, + EXP_stride=exp_buf.stride(0), + sum_exp_ptr=sum_exp_global, + Y_ptr=target, + grad_out_ptr=grad_out, + vocab_start=ctx.vocab_start, + n_cols=weight.shape[0], + ignore_index=ctx.ignore_index, + alpha_eff=0.0, + eps_eff=0.0, + HAS_LABEL_SMOOTHING=False, + BLOCK_SIZE=ctx.ce_block_size, + num_warps=num_warps, + ) + + grad_hidden = _triton_dx_matmul(exp_buf, weight) + grad_weight = _triton_matmul(exp_buf.t(), hidden) + grad_bias = _triton_column_sum(exp_buf).to(ctx.bias_dtype) if ctx.has_bias else None + + if ctx.tp_world > 1: + dist.all_reduce(grad_hidden, op=dist.ReduceOp.SUM, group=ctx.tp_group) + grad_hidden = grad_hidden.to(ctx.hidden_dtype).reshape(ctx.original_hidden_shape) + return grad_hidden, grad_weight, grad_bias + + +class LigerMegatronFusedLinearCrossEntropyFunction(torch.autograd.Function): + """Hidden-to-loss tensor-parallel FLCE with saved low-precision CE state.""" + + @staticmethod + def forward( + ctx, + hidden: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + bias: torch.Tensor | None, + tp_group, + ignore_index: int, + ) -> torch.Tensor: + if hidden.ndim < 2: + raise ValueError(f"hidden must have at least 2 dimensions, got shape {tuple(hidden.shape)}.") + if weight.ndim != 2: + raise ValueError(f"weight must be 2-D [V_local, H], got shape {tuple(weight.shape)}.") + if tuple(target.shape) != tuple(hidden.shape[:-1]): + raise ValueError( + f"target shape must equal hidden.shape[:-1]; got target={tuple(target.shape)}, " + f"hidden={tuple(hidden.shape)}." + ) + if hidden.shape[-1] != weight.shape[1]: + raise ValueError(f"hidden size mismatch: hidden has H={hidden.shape[-1]}, weight has H={weight.shape[1]}.") + if hidden.dtype != weight.dtype: + raise TypeError(f"hidden and weight must have the same dtype, got {hidden.dtype} and {weight.dtype}.") + if hidden.device != weight.device or hidden.device != target.device: + raise ValueError("hidden, weight, and target must be on the same device.") + if bias is not None: + if bias.ndim != 1 or bias.shape[0] != weight.shape[0]: + raise ValueError(f"bias must have shape ({weight.shape[0]},), got {tuple(bias.shape)}.") + if bias.device != hidden.device or bias.dtype != hidden.dtype: + raise TypeError("bias must have the same device and dtype as hidden.") + if hidden.device.type != "cuda" or hidden.dtype not in (torch.bfloat16, torch.float16): + raise RuntimeError("Megatron FLCE requires a CUDA GPU and float16 or bfloat16 inputs.") + + tp_rank, tp_world = _tp_rank_and_world(tp_group) + vocab_local = weight.shape[0] + vocab_global = vocab_local * tp_world + vocab_start = tp_rank * vocab_local + + flat_target = target.reshape(-1).to(torch.int64).contiguous() + valid = flat_target != ignore_index + invalid = valid & ((flat_target < 0) | (flat_target >= vocab_global)) + valid_targets = ~torch.any(invalid) + if hasattr(torch, "_assert_async"): + torch._assert_async(valid_targets, f"non-ignored targets must be in [0, {vocab_global}).") + elif not valid_targets.item(): + raise ValueError(f"non-ignored targets must be in [0, {vocab_global}).") + + original_hidden_shape = hidden.shape + hidden_2d = hidden.reshape(-1, hidden.shape[-1]).contiguous() + weight_2d = weight.contiguous() + bias_1d = bias.contiguous() if bias is not None else None + + from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps + from liger_kernel.ops.vocab_parallel_cross_entropy import _select_block_size + from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_forward_kernel + + logits = _triton_matmul(hidden_2d, weight_2d.t(), bias=bias_1d) + ce_block_size = _select_block_size(vocab_local) + logits_max = _triton_row_max(logits, ce_block_size) + if tp_world > 1: + dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group) + + exp_buf = torch.empty( + hidden_2d.shape[0], + vocab_local, + device=hidden.device, + dtype=hidden.dtype, + ) + predicted_logit = torch.empty(hidden_2d.shape[0], device=hidden.device, dtype=torch.float32) + sum_exp = torch.empty_like(predicted_logit) + num_warps = _get_num_warps(ce_block_size) + liger_vocab_parallel_ce_forward_kernel[(hidden_2d.shape[0],)]( + X_ptr=logits, + X_stride=logits.stride(0), + EXP_ptr=exp_buf, + EXP_stride=exp_buf.stride(0), + logits_max_ptr=logits_max, + Y_ptr=flat_target, + pred_ptr=predicted_logit, + sum_exp_ptr=sum_exp, + vocab_start=vocab_start, + n_cols=vocab_local, + ignore_index=ignore_index, + BLOCK_SIZE=ce_block_size, + num_warps=num_warps, + ) + if tp_world > 1: + dist.all_reduce(predicted_logit, op=dist.ReduceOp.SUM, group=tp_group) + dist.all_reduce(sum_exp, op=dist.ReduceOp.SUM, group=tp_group) + + loss = _triton_loss(sum_exp, predicted_logit, flat_target, ignore_index) + + ctx.save_for_backward(hidden_2d, weight_2d, exp_buf, sum_exp, flat_target) + ctx.has_bias = bias is not None + ctx.bias_dtype = bias.dtype if bias is not None else None + ctx.tp_group = tp_group + ctx.tp_world = tp_world + ctx.vocab_start = vocab_start + ctx.ignore_index = ignore_index + ctx.ce_block_size = ce_block_size + ctx.original_hidden_shape = original_hidden_shape + ctx.hidden_dtype = hidden.dtype + return loss.reshape(target.shape) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + grad_hidden, grad_weight, grad_bias = _materialized_backward(ctx, grad_output) + return grad_hidden, grad_weight, None, grad_bias, None, None + + +def liger_megatron_fused_linear_cross_entropy( + hidden: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + bias: torch.Tensor | None = None, + tp_group=None, + ignore_index: int = -100, +) -> torch.Tensor: + """Compute per-token loss from replicated hidden states and a local vocab shard.""" + return LigerMegatronFusedLinearCrossEntropyFunction.apply( + hidden, + weight, + target, + bias, + tp_group, + ignore_index, + ) diff --git a/test/megatron/test_cutile_fused_linear_cross_entropy.py b/test/megatron/test_cutile_fused_linear_cross_entropy.py new file mode 100644 index 000000000..ac32c5287 --- /dev/null +++ b/test/megatron/test_cutile_fused_linear_cross_entropy.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +import os +import subprocess +import sys +import tempfile + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F + +pytest.importorskip("cuda.tile") + +from liger_kernel.ops.cutile.ops.megatron_fused_linear_cross_entropy import ( # noqa: E402 + liger_megatron_fused_linear_cross_entropy, +) + + +def _reference_loss(hidden, weight, target, bias=None, ignore_index=-100): + logits = hidden.float() @ weight.float().t() + if bias is not None: + logits = logits + bias.float() + return F.cross_entropy( + logits.reshape(-1, logits.shape[-1]), + target.reshape(-1), + reduction="none", + ignore_index=ignore_index, + ).reshape(target.shape) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CuTile FLCE requires CUDA") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_cutile_megatron_flce_tp1_matches_pytorch(dtype, with_bias): + torch.manual_seed(42) + hidden_base = torch.randn(3, 2, 65, device="cuda", dtype=dtype) + weight_base = torch.randn(33, 65, device="cuda", dtype=dtype) * 0.02 + bias_base = torch.randn(33, device="cuda", dtype=dtype) * 0.02 if with_bias else None + target = torch.randint(0, 33, (3, 2), device="cuda") + target[0, 0] = -100 + upstream = torch.randn(3, 2, device="cuda") + + hidden_ref = hidden_base.clone().requires_grad_(True) + weight_ref = weight_base.clone().requires_grad_(True) + bias_ref = bias_base.clone().requires_grad_(True) if bias_base is not None else None + hidden_cutile = hidden_base.clone().requires_grad_(True) + weight_cutile = weight_base.clone().requires_grad_(True) + bias_cutile = bias_base.clone().requires_grad_(True) if bias_base is not None else None + + reference = _reference_loss(hidden_ref, weight_ref, target, bias_ref) + actual = liger_megatron_fused_linear_cross_entropy( + hidden_cutile, + weight_cutile, + target, + bias=bias_cutile, + ) + + torch.testing.assert_close(actual, reference, atol=5e-3, rtol=5e-2) + reference.backward(upstream) + actual.backward(upstream) + torch.testing.assert_close(hidden_cutile.grad, hidden_ref.grad, atol=5e-3, rtol=5e-2) + torch.testing.assert_close(weight_cutile.grad, weight_ref.grad, atol=5e-3, rtol=5e-2) + if with_bias: + torch.testing.assert_close(bias_cutile.grad, bias_ref.grad, atol=5e-3, rtol=5e-2) + + +def test_cutile_megatron_flce_backend_dispatch(): + env = os.environ.copy() + env["LIGER_KERNEL_IMPL"] = "cutile" + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "from liger_kernel.megatron.fused_linear_cross_entropy import " + "liger_megatron_fused_linear_cross_entropy as fn; print(fn.__module__)" + ), + ], + check=True, + capture_output=True, + text=True, + env=env, + ) + assert result.stdout.strip() == "liger_kernel.ops.cutile.ops.megatron_fused_linear_cross_entropy" + + +def _tp_worker(rank, world_size, file_name, dtype): + dist.init_process_group( + backend="nccl", + init_method=f"file://{file_name}", + rank=rank, + world_size=world_size, + ) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + tp_group = dist.group.WORLD + vocab_global = 64 + vocab_local = vocab_global // world_size + + torch.manual_seed(123) + hidden_base = torch.randn(3, 2, 65, device=device, dtype=dtype) + weight_global = torch.randn(vocab_global, 65, device=device, dtype=dtype) * 0.02 + bias_global = torch.randn(vocab_global, device=device, dtype=dtype) * 0.02 + target = torch.randint(0, vocab_global, (3, 2), device=device) + upstream = torch.randn(3, 2, device=device) + target[0, 0] = -100 + for tensor in (hidden_base, weight_global, bias_global, target, upstream): + dist.broadcast(tensor, src=0, group=tp_group) + + start = rank * vocab_local + end = start + vocab_local + hidden_cutile = hidden_base.clone().requires_grad_(True) + weight_local = weight_global[start:end].clone().requires_grad_(True) + bias_local = bias_global[start:end].clone().requires_grad_(True) + + actual = liger_megatron_fused_linear_cross_entropy( + hidden_cutile, + weight_local, + target, + bias=bias_local, + tp_group=tp_group, + ) + + hidden_ref = hidden_base.clone().requires_grad_(True) + weight_ref = weight_global.clone().requires_grad_(True) + bias_ref = bias_global.clone().requires_grad_(True) + reference = _reference_loss(hidden_ref, weight_ref, target, bias_ref) + + torch.testing.assert_close(actual, reference, atol=5e-3, rtol=5e-2) + actual.backward(upstream) + reference.backward(upstream) + torch.testing.assert_close(hidden_cutile.grad, hidden_ref.grad, atol=5e-3, rtol=5e-2) + torch.testing.assert_close(weight_local.grad, weight_ref.grad[start:end], atol=5e-3, rtol=5e-2) + torch.testing.assert_close(bias_local.grad, bias_ref.grad[start:end], atol=5e-3, rtol=5e-2) + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.device_count() < 2, + reason="requires at least two CUDA GPUs", +) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_cutile_megatron_flce_tp2_matches_global_reference(dtype): + with tempfile.NamedTemporaryFile() as rendezvous: + mp.spawn( + _tp_worker, + args=(2, rendezvous.name, dtype), + nprocs=2, + join=True, + ) diff --git a/test/megatron/test_fused_linear_cross_entropy.py b/test/megatron/test_fused_linear_cross_entropy.py new file mode 100644 index 000000000..39453870b --- /dev/null +++ b/test/megatron/test_fused_linear_cross_entropy.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +import os +import subprocess +import sys +import tempfile + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F + +from liger_kernel.megatron import LigerMegatronFusedLinearCrossEntropy +from liger_kernel.ops.megatron_fused_linear_cross_entropy import liger_megatron_fused_linear_cross_entropy +from liger_kernel.ops.triton.ops.megatron_fused_linear_cross_entropy import ( + liger_megatron_fused_linear_cross_entropy as triton_megatron_fused_linear_cross_entropy, +) + +_IMPLEMENTATIONS = [ + pytest.param(liger_megatron_fused_linear_cross_entropy, id="cublas"), + pytest.param(triton_megatron_fused_linear_cross_entropy, id="triton"), +] + + +def _reference_loss(hidden, weight, target, bias=None, ignore_index=-100): + logits = hidden.float() @ weight.float().t() + if bias is not None: + logits = logits + bias.float() + return F.cross_entropy( + logits.reshape(-1, logits.shape[-1]), + target.reshape(-1), + reduction="none", + ignore_index=ignore_index, + ).reshape(target.shape) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Megatron FLCE requires CUDA") +@pytest.mark.parametrize("shape", [(2, 3, 8, 16), (3, 2, 17, 32)]) +@pytest.mark.parametrize("with_bias", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("implementation", _IMPLEMENTATIONS) +def test_megatron_flce_tp1_matches_pytorch(shape, with_bias, dtype, implementation): + s, b, h, v = shape + torch.manual_seed(42) + hidden_base = torch.randn(s, b, h, device="cuda", dtype=dtype) + weight_base = torch.randn(v, h, device="cuda", dtype=dtype) * 0.02 + bias_base = torch.randn(v, device="cuda", dtype=dtype) * 0.02 if with_bias else None + target = torch.randint(0, v, (s, b), device="cuda") + target.reshape(-1)[0] = -100 + upstream = torch.randn(s, b, device="cuda") + + hidden_ref = hidden_base.clone().requires_grad_(True) + weight_ref = weight_base.clone().requires_grad_(True) + bias_ref = bias_base.clone().requires_grad_(True) if bias_base is not None else None + hidden_liger = hidden_base.clone().requires_grad_(True) + weight_liger = weight_base.clone().requires_grad_(True) + bias_liger = bias_base.clone().requires_grad_(True) if bias_base is not None else None + + reference = _reference_loss(hidden_ref, weight_ref, target, bias_ref) + actual = implementation( + hidden_liger, + weight_liger, + target, + bias=bias_liger, + ) + + torch.testing.assert_close(actual, reference, atol=5e-3, rtol=5e-2) + reference.backward(upstream) + actual.backward(upstream) + torch.testing.assert_close(hidden_liger.grad, hidden_ref.grad, atol=5e-3, rtol=5e-2) + torch.testing.assert_close(weight_liger.grad, weight_ref.grad, atol=5e-3, rtol=5e-2) + if with_bias: + torch.testing.assert_close(bias_liger.grad, bias_ref.grad, atol=5e-3, rtol=5e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Megatron FLCE requires CUDA") +def test_megatron_flce_module_contract(): + module = LigerMegatronFusedLinearCrossEntropy(ignore_index=-1) + hidden = torch.randn(2, 3, 8, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(16, 8, device="cuda", dtype=torch.bfloat16) * 0.02 + target = torch.randint(0, 16, (2, 3), device="cuda") + target[0, 0] = -1 + + actual = module(hidden, weight, target) + reference = _reference_loss(hidden, weight, target, ignore_index=-1) + torch.testing.assert_close(actual, reference, atol=5e-3, rtol=5e-2) + assert "ignore_index=-1" in repr(module) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Megatron FLCE requires CUDA") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_megatron_flce_triton_split_k_dx_matches_pytorch(dtype): + torch.manual_seed(44) + hidden_base = torch.randn(8, 4, 64, device="cuda", dtype=dtype) + weight_base = torch.randn(8192, 64, device="cuda", dtype=dtype) * 0.02 + target = torch.randint(0, weight_base.shape[0], hidden_base.shape[:-1], device="cuda") + upstream = torch.randn_like(target, dtype=torch.float32) + + hidden_ref = hidden_base.clone().requires_grad_(True) + weight_ref = weight_base.clone().requires_grad_(True) + hidden_triton = hidden_base.clone().requires_grad_(True) + weight_triton = weight_base.clone().requires_grad_(True) + + reference = _reference_loss(hidden_ref, weight_ref, target) + actual = triton_megatron_fused_linear_cross_entropy(hidden_triton, weight_triton, target) + torch.testing.assert_close(actual, reference, atol=5e-3, rtol=5e-2) + + reference.backward(upstream) + actual.backward(upstream) + torch.testing.assert_close(hidden_triton.grad, hidden_ref.grad, atol=5e-3, rtol=5e-2) + torch.testing.assert_close(weight_triton.grad, weight_ref.grad, atol=5e-3, rtol=5e-2) + + +def test_megatron_flce_triton_backend_dispatch(): + env = os.environ.copy() + env["LIGER_KERNEL_IMPL"] = "triton" + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "from liger_kernel.megatron.fused_linear_cross_entropy import " + "liger_megatron_fused_linear_cross_entropy as fn; print(fn.__module__)" + ), + ], + check=True, + capture_output=True, + text=True, + env=env, + ) + assert result.stdout.strip() == "liger_kernel.ops.triton.ops.megatron_fused_linear_cross_entropy" + + +def test_megatron_flce_rejects_cpu_inputs(): + hidden = torch.randn(2, 3, 8, dtype=torch.bfloat16) + weight = torch.randn(16, 8, dtype=torch.bfloat16) + target = torch.randint(0, 16, (2, 3)) + + with pytest.raises(RuntimeError, match="requires a CUDA GPU"): + liger_megatron_fused_linear_cross_entropy(hidden, weight, target) + + +def _tp_worker(rank, world_size, file_name, dtype, implementation_name): + dist.init_process_group( + backend="nccl", + init_method=f"file://{file_name}", + rank=rank, + world_size=world_size, + ) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + tp_group = dist.group.WORLD + implementation = ( + triton_megatron_fused_linear_cross_entropy + if implementation_name == "triton" + else liger_megatron_fused_linear_cross_entropy + ) + s, b, h, v_global = 3, 2, 17, 32 + v_local = v_global // world_size + + torch.manual_seed(123) + hidden_base = torch.randn(s, b, h, device=device, dtype=dtype) + weight_global = torch.randn(v_global, h, device=device, dtype=dtype) * 0.02 + bias_global = torch.randn(v_global, device=device, dtype=dtype) * 0.02 + target = torch.randint(0, v_global, (s, b), device=device) + upstream = torch.randn(s, b, device=device) + target.reshape(-1)[0] = -100 + for tensor in (hidden_base, weight_global, bias_global, target, upstream): + dist.broadcast(tensor, src=0, group=tp_group) + + start = rank * v_local + end = start + v_local + hidden_liger = hidden_base.clone().requires_grad_(True) + weight_local = weight_global[start:end].clone().requires_grad_(True) + bias_local = bias_global[start:end].clone().requires_grad_(True) + + actual = implementation( + hidden_liger, + weight_local, + target, + bias=bias_local, + tp_group=tp_group, + ) + + hidden_ref = hidden_base.clone().requires_grad_(True) + weight_ref = weight_global.clone().requires_grad_(True) + bias_ref = bias_global.clone().requires_grad_(True) + reference = _reference_loss(hidden_ref, weight_ref, target, bias_ref) + + torch.testing.assert_close(actual, reference, atol=5e-3, rtol=5e-2) + actual.backward(upstream) + reference.backward(upstream) + torch.testing.assert_close(hidden_liger.grad, hidden_ref.grad, atol=5e-3, rtol=5e-2) + torch.testing.assert_close(weight_local.grad, weight_ref.grad[start:end], atol=5e-3, rtol=5e-2) + torch.testing.assert_close(bias_local.grad, bias_ref.grad[start:end], atol=5e-3, rtol=5e-2) + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.device_count() < 2, + reason="requires at least two CUDA GPUs", +) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("implementation_name", ["cublas", "triton"]) +def test_megatron_flce_tp2_matches_global_reference(dtype, implementation_name): + with tempfile.NamedTemporaryFile() as rendezvous: + mp.spawn( + _tp_worker, + args=(2, rendezvous.name, dtype, implementation_name), + nprocs=2, + join=True, + ) From 52c080d08c375f40d3c78aaaa2e703bdf6cde4b6 Mon Sep 17 00:00:00 2001 From: Justin Hu Date: Tue, 11 Aug 2026 17:29:12 +0000 Subject: [PATCH 2/7] feat(cutedsl): add Megatron fused linear cross entropy Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- benchmark/README.md | 23 +- ...ark_megatron_fused_linear_cross_entropy.py | 39 ++- docs/High-Level-APIs.md | 29 ++- src/liger_kernel/ops/cutedsl/ops/__init__.py | 6 + .../megatron_fused_linear_cross_entropy.py | 244 ++++++++++++++++++ .../megatron_fused_linear_cross_entropy.py | 29 ++- .../megatron_fused_linear_cross_entropy.py | 29 ++- ...test_cutedsl_fused_linear_cross_entropy.py | 155 +++++++++++ 8 files changed, 519 insertions(+), 35 deletions(-) create mode 100644 src/liger_kernel/ops/cutedsl/ops/megatron_fused_linear_cross_entropy.py create mode 100644 test/megatron/test_cutedsl_fused_linear_cross_entropy.py diff --git a/benchmark/README.md b/benchmark/README.md index 1e33bcd2e..7afcaa95f 100644 --- a/benchmark/README.md +++ b/benchmark/README.md @@ -186,4 +186,25 @@ visualizer at it with `--data-file`: python ../benchmarks_visualizer.py \ --kernel-name cross_entropy --metric-name speed \ --data-file data/all_benchmark_data_cutedsl.csv -``` \ No newline at end of file +``` + +### Megatron fused linear cross-entropy + +The Megatron FLCE benchmark checks loss and gradient parity before measuring +the default cuBLAS/Triton implementation and the opt-in Triton, CuTile, and +CuTe DSL backends: + +```bash +cd benchmark/scripts +python benchmark_megatron_fused_linear_cross_entropy.py \ + --tp-size 4 --token-counts 512 2048 --vocab-sizes 32000 128256 \ + --providers megatron-compatible liger liger-triton liger-cutile + +# SM100 with nvidia-cutlass-dsl installed +python benchmark_megatron_fused_linear_cross_entropy.py \ + --tp-size 4 --providers megatron-compatible liger-cutedsl +``` + +Targets are global vocabulary indices; each rank owns a contiguous vocabulary +shard. Results include forward, backward, full-step latency, and peak Torch +memory. diff --git a/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py b/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py index 8b012c5ed..d13158afc 100644 --- a/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py +++ b/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py @@ -8,9 +8,11 @@ This script compares that path with ``LigerMegatronFusedLinearCrossEntropy``, which saves low-precision CE state to avoid projection recomputation. The ``liger-triton`` provider additionally replaces all three local GEMMs with -portable Triton kernels. When Megatron-Core is installed, the ``megatron-core`` -provider uses its fused CE. The always-available ``megatron-compatible`` -provider uses Liger's drop-in Megatron CE. +portable Triton kernels. ``liger-cutile`` replaces all local compute with +CuTile, while ``liger-cutedsl`` uses a persistent SM100 CuTe DSL projection. +When Megatron-Core is installed, the ``megatron-core`` provider uses its fused +CE. The always-available ``megatron-compatible`` provider uses Liger's drop-in +Megatron CE. Backward timing creates a fresh graph outside each timed event pair, so only backward execution is measured while respecting Megatron's single-use fused CE @@ -50,6 +52,16 @@ liger_megatron_fused_linear_cross_entropy as triton_megatron_fused_linear_cross_entropy, ) +try: + from liger_kernel.ops.cutedsl.ops.megatron_fused_linear_cross_entropy import ( + liger_megatron_fused_linear_cross_entropy as cutedsl_megatron_fused_linear_cross_entropy, + ) + + _CUTEDSL_AVAILABLE = True +except ImportError: + cutedsl_megatron_fused_linear_cross_entropy = None + _CUTEDSL_AVAILABLE = False + try: from liger_kernel.ops.cutile.ops.megatron_fused_linear_cross_entropy import ( liger_megatron_fused_linear_cross_entropy as cutile_megatron_fused_linear_cross_entropy, @@ -121,6 +133,16 @@ def _make_state( bias=bias, tp_group=tp_group, ) + elif provider == "liger-cutedsl": + if not _CUTEDSL_AVAILABLE: + raise RuntimeError("provider 'liger-cutedsl' requires nvidia-cutlass-dsl.") + forward = lambda: cutedsl_megatron_fused_linear_cross_entropy( + hidden, + weight, + target, + bias=bias, + tp_group=tp_group, + ) elif provider == "liger-cutile": if not _CUTILE_AVAILABLE: raise RuntimeError("provider 'liger-cutile' requires the cuda-tile package.") @@ -479,7 +501,14 @@ def main(): parser.add_argument( "--providers", nargs="+", - choices=["liger", "liger-triton", "liger-cutile", "megatron-compatible", "megatron-core"], + choices=[ + "liger", + "liger-triton", + "liger-cutedsl", + "liger-cutile", + "megatron-compatible", + "megatron-core", + ], ) parser.add_argument( "--output", @@ -496,6 +525,8 @@ def main(): providers.insert(1, "megatron-core") if "megatron-core" in providers and not _MEGATRON_CORE_AVAILABLE: raise RuntimeError("provider 'megatron-core' requested, but megatron-core is not installed.") + if "liger-cutedsl" in providers and not _CUTEDSL_AVAILABLE: + raise RuntimeError("provider 'liger-cutedsl' requested, but nvidia-cutlass-dsl is not installed.") if "liger-cutile" in providers and not _CUTILE_AVAILABLE: raise RuntimeError("provider 'liger-cutile' requested, but cuda-tile is not installed.") if min(args.token_counts) <= 0 or args.hidden_size <= 0 or min(args.vocab_sizes) <= 0: diff --git a/docs/High-Level-APIs.md b/docs/High-Level-APIs.md index 6bbe008a9..dae52662d 100644 --- a/docs/High-Level-APIs.md +++ b/docs/High-Level-APIs.md @@ -98,17 +98,22 @@ You can also use the Patching APIs to use the kernels for a specific model archi Liger also exposes a patch for the [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) training framework, replacing Megatron's native RMSNorm and both vocab-parallel -cross-entropy paths (fused and unfused) with Liger's Triton kernels. +cross-entropy paths (fused and unfused) with Liger kernels. Liger also exposes +a hidden-state-to-loss FLCE module for contiguous tensor-parallel vocabulary +shards. | **Framework** | **API** | **Supported Operations** | |---------------|--------------------------------------------------------|--------------------------| | Megatron-LM | `liger_kernel.megatron.apply_liger_kernel_to_megatron` | RMSNorm, CrossEntropyLoss | +| Megatron-LM | `liger_kernel.megatron.LigerMegatronFusedLinearCrossEntropy` | Fused output projection + CrossEntropyLoss | -**Scope**: Initial release supports `tensor_model_parallel_size=1` only for -cross-entropy. Vocab-parallel cross-entropy (TP>1) is follow-up work — with -TP>1, each rank holds a sharded `[N, V/tp]` logits slice and cross-entropy -requires cross-rank all-reduces that Liger's kernel does not perform. The -patch raises a `RuntimeError` at patch time or call time if TP>1 is detected. +`LigerMegatronFusedLinearCrossEntropy` accepts replicated hidden states, +the calling rank's contiguous `[V_local, H]` output-weight shard, and global +target indices. It supports TP1 and TP>1 through the supplied process group. +The default implementation uses cuBLAS, Triton CE, and NCCL. Set +`LIGER_KERNEL_IMPL=triton`, `cutile`, or `cutedsl` before importing Liger to +select an all-Triton local path, a CuTile local path, or the SM100 CuTe DSL +persistent projection. **Usage**: @@ -118,6 +123,12 @@ from liger_kernel.megatron import apply_liger_kernel_to_megatron # Call before Megatron's forward pass reaches compute_language_model_loss. # Defaults match Megatron's native CE behavior; no CE-specific config needed. apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True) + +# Or wire the hidden-state-to-loss operation into a vocab-sharded output layer. +from liger_kernel.megatron import LigerMegatronFusedLinearCrossEntropy + +loss_fn = LigerMegatronFusedLinearCrossEntropy(ignore_index=-100) +loss = loss_fn(hidden, local_output_weight, global_targets, tp_group=tp_group) ``` Both the fused (`config.cross_entropy_loss_fusion=True`, @@ -130,6 +141,12 @@ For training setups that need explicit kernel configuration (custom `LigerMegatronCrossEntropy` directly and wire it into your model — see `examples/megatron/run_mode2_hand_spec.py`. +::: liger_kernel.megatron.LigerMegatronFusedLinearCrossEntropy + options: + extra: + show_docstring: true + show_signature: true + ::: liger_kernel.megatron.apply_liger_kernel_to_megatron options: extra: diff --git a/src/liger_kernel/ops/cutedsl/ops/__init__.py b/src/liger_kernel/ops/cutedsl/ops/__init__.py index a5b051b1a..de25f51bc 100644 --- a/src/liger_kernel/ops/cutedsl/ops/__init__.py +++ b/src/liger_kernel/ops/cutedsl/ops/__init__.py @@ -18,6 +18,10 @@ from liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90 import LigerFusedScaledCrossEntropySM90Function from liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90 import fused_scaled_cross_entropy_backward from liger_kernel.ops.cutedsl.ops.fused_scaled_cross_entropy_sm90 import fused_scaled_cross_entropy_forward +from liger_kernel.ops.cutedsl.ops.megatron_fused_linear_cross_entropy import ( + LigerMegatronFusedLinearCrossEntropyFunction, +) +from liger_kernel.ops.cutedsl.ops.megatron_fused_linear_cross_entropy import liger_megatron_fused_linear_cross_entropy from liger_kernel.ops.cutedsl.ops.rms_norm import LigerRMSNormFunction from liger_kernel.ops.cutedsl.ops.rms_norm import rms_norm_backward from liger_kernel.ops.cutedsl.ops.rms_norm import rms_norm_forward @@ -41,6 +45,8 @@ "LigerFusedScaledCrossEntropySM90Function", "fused_scaled_cross_entropy_backward", "fused_scaled_cross_entropy_forward", + "LigerMegatronFusedLinearCrossEntropyFunction", + "liger_megatron_fused_linear_cross_entropy", "LigerRMSNormFunction", "rms_norm_backward", "rms_norm_forward", diff --git a/src/liger_kernel/ops/cutedsl/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/cutedsl/ops/megatron_fused_linear_cross_entropy.py new file mode 100644 index 000000000..c3a7a9026 --- /dev/null +++ b/src/liger_kernel/ops/cutedsl/ops/megatron_fused_linear_cross_entropy.py @@ -0,0 +1,244 @@ +"""CuTe DSL tensor-parallel fused linear cross entropy for Megatron. + +The SM100 path uses a persistent CuTe DSL GEMM for the local vocabulary +projection, Triton for vocabulary-parallel cross entropy, and NCCL for +tensor-parallel collectives. Shifted exponentials overwrite the projection +buffer and are reused by backward. +""" + +from __future__ import annotations + +import operator + +import cutlass +import cutlass.cute as cute +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from liger_kernel.ops.cutedsl.ops._sm100_gemm import K_ALIGNMENT +from liger_kernel.ops.cutedsl.ops._sm100_gemm import run_epilogue_gemm +from liger_kernel.ops.megatron_fused_linear_cross_entropy import _tp_rank_and_world +from liger_kernel.ops.megatron_fused_linear_cross_entropy import ( + liger_megatron_fused_linear_cross_entropy as default_megatron_fused_linear_cross_entropy, +) +from liger_kernel.ops.utils import compare_version + +_SUPPORTS_OUT_DTYPE = compare_version("torch", operator.ge, "2.8.0") + + +@cute.jit +def _identity_epilogue(accumulator, output): + output_dtype = output.element_type + for element in cutlass.range_constexpr(cute.size(accumulator)): + output[element] = accumulator[element].to(output_dtype) + + +def _native_cutedsl_supported(hidden: torch.Tensor, weight: torch.Tensor) -> bool: + if hidden.device.type != "cuda" or hidden.dtype not in (torch.bfloat16, torch.float16): + return False + if weight.device != hidden.device or weight.dtype != hidden.dtype: + return False + try: + return torch.cuda.get_device_capability(hidden.device)[0] >= 10 + except (AssertionError, RuntimeError): + return False + + +def _cutedsl_projection( + hidden: torch.Tensor, + weight: torch.Tensor, +) -> torch.Tensor: + padding = (-hidden.shape[1]) % K_ALIGNMENT + if padding: + hidden = F.pad(hidden, (0, padding)) + weight = F.pad(weight, (0, padding)) + logits = torch.empty( + hidden.shape[0], + weight.shape[0], + device=hidden.device, + dtype=hidden.dtype, + ) + run_epilogue_gemm(hidden, weight, logits, _identity_epilogue) + return logits + + +def _materialized_backward(ctx, grad_output: torch.Tensor): + from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps + from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_backward_kernel + + hidden, weight, exp_buffer, sum_exp, target = ctx.saved_tensors + grad_output_1d = grad_output.contiguous().reshape(-1).float() + num_warps = _get_num_warps(ctx.ce_block_size) + liger_vocab_parallel_ce_backward_kernel[(hidden.shape[0],)]( + EXP_ptr=exp_buffer, + EXP_stride=exp_buffer.stride(0), + sum_exp_ptr=sum_exp, + Y_ptr=target, + grad_out_ptr=grad_output_1d, + vocab_start=ctx.vocab_start, + n_cols=weight.shape[0], + ignore_index=ctx.ignore_index, + alpha_eff=0.0, + eps_eff=0.0, + HAS_LABEL_SMOOTHING=False, + BLOCK_SIZE=ctx.ce_block_size, + num_warps=num_warps, + ) + + if _SUPPORTS_OUT_DTYPE: + grad_hidden = torch.mm(exp_buffer, weight, out_dtype=torch.float32) + else: + grad_hidden = exp_buffer.float() @ weight.float() + reduce_work = ( + dist.all_reduce( + grad_hidden, + op=dist.ReduceOp.SUM, + group=ctx.tp_group, + async_op=True, + ) + if ctx.tp_world > 1 + else None + ) + grad_weight = exp_buffer.t() @ hidden + grad_bias = exp_buffer.sum(dim=0, dtype=torch.float32).to(ctx.bias_dtype) if ctx.has_bias else None + if reduce_work is not None: + reduce_work.wait() + + grad_hidden = grad_hidden.to(ctx.hidden_dtype).reshape(ctx.original_hidden_shape) + return grad_hidden, grad_weight, grad_bias + + +class LigerMegatronFusedLinearCrossEntropyFunction(torch.autograd.Function): + """Megatron FLCE using a persistent CuTe DSL SM100 projection.""" + + @staticmethod + def forward( + ctx, + hidden: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + bias: torch.Tensor | None, + tp_group, + ignore_index: int, + ) -> torch.Tensor: + if hidden.ndim < 2: + raise ValueError(f"hidden must have at least 2 dimensions, got shape {tuple(hidden.shape)}.") + if weight.ndim != 2: + raise ValueError(f"weight must be 2-D [V_local, H], got shape {tuple(weight.shape)}.") + if tuple(target.shape) != tuple(hidden.shape[:-1]): + raise ValueError( + f"target shape must equal hidden.shape[:-1]; got target={tuple(target.shape)}, " + f"hidden={tuple(hidden.shape)}." + ) + if hidden.shape[-1] != weight.shape[1]: + raise ValueError(f"hidden size mismatch: hidden has H={hidden.shape[-1]}, weight has H={weight.shape[1]}.") + if hidden.dtype != weight.dtype: + raise TypeError(f"hidden and weight must have the same dtype, got {hidden.dtype} and {weight.dtype}.") + if hidden.device != weight.device or hidden.device != target.device: + raise ValueError("hidden, weight, and target must be on the same device.") + if bias is not None: + if bias.ndim != 1 or bias.shape[0] != weight.shape[0]: + raise ValueError(f"bias must have shape ({weight.shape[0]},), got {tuple(bias.shape)}.") + if bias.device != hidden.device or bias.dtype != hidden.dtype: + raise TypeError("bias must have the same device and dtype as hidden.") + + tp_rank, tp_world = _tp_rank_and_world(tp_group) + vocab_local = weight.shape[0] + vocab_global = vocab_local * tp_world + vocab_start = tp_rank * vocab_local + flat_target = target.reshape(-1).to(torch.int64).contiguous() + valid = flat_target != ignore_index + invalid = valid & ((flat_target < 0) | (flat_target >= vocab_global)) + valid_targets = ~torch.any(invalid) + if hasattr(torch, "_assert_async"): + torch._assert_async(valid_targets, f"non-ignored targets must be in [0, {vocab_global}).") + elif not valid_targets.item(): + raise ValueError(f"non-ignored targets must be in [0, {vocab_global}).") + + original_hidden_shape = hidden.shape + hidden_2d = hidden.reshape(-1, hidden.shape[-1]).contiguous() + weight_2d = weight.contiguous() + bias_1d = bias.contiguous() if bias is not None else None + logits = _cutedsl_projection(hidden_2d, weight_2d) + if bias_1d is not None: + logits.add_(bias_1d) + + logits_max = logits.amax(dim=-1).float() + if tp_world > 1: + dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group) + + from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps + from liger_kernel.ops.vocab_parallel_cross_entropy import _select_block_size + from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_forward_kernel + + exp_buffer = logits + stats = torch.empty((2, hidden_2d.shape[0]), device=hidden.device, dtype=torch.float32) + predicted_logit = stats[0] + sum_exp = stats[1] + ce_block_size = _select_block_size(vocab_local) + num_warps = _get_num_warps(ce_block_size) + liger_vocab_parallel_ce_forward_kernel[(hidden_2d.shape[0],)]( + X_ptr=logits, + X_stride=logits.stride(0), + EXP_ptr=exp_buffer, + EXP_stride=exp_buffer.stride(0), + logits_max_ptr=logits_max, + Y_ptr=flat_target, + pred_ptr=predicted_logit, + sum_exp_ptr=sum_exp, + vocab_start=vocab_start, + n_cols=vocab_local, + ignore_index=ignore_index, + BLOCK_SIZE=ce_block_size, + num_warps=num_warps, + ) + if tp_world > 1: + dist.all_reduce(stats, op=dist.ReduceOp.SUM, group=tp_group) + + loss = torch.log(sum_exp) - predicted_logit + loss = torch.where(valid, loss, torch.zeros_like(loss)) + ctx.save_for_backward(hidden_2d, weight_2d, exp_buffer, sum_exp, flat_target) + ctx.has_bias = bias is not None + ctx.bias_dtype = bias.dtype if bias is not None else None + ctx.tp_group = tp_group + ctx.tp_world = tp_world + ctx.vocab_start = vocab_start + ctx.ignore_index = ignore_index + ctx.ce_block_size = ce_block_size + ctx.original_hidden_shape = original_hidden_shape + ctx.hidden_dtype = hidden.dtype + return loss.reshape(target.shape) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + grad_hidden, grad_weight, grad_bias = _materialized_backward(ctx, grad_output) + return grad_hidden, grad_weight, None, grad_bias, None, None + + +def liger_megatron_fused_linear_cross_entropy( + hidden: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + bias: torch.Tensor | None = None, + tp_group=None, + ignore_index: int = -100, +) -> torch.Tensor: + """Compute Megatron FLCE with a CuTe DSL projection and NCCL TP collectives.""" + if not _native_cutedsl_supported(hidden, weight): + return default_megatron_fused_linear_cross_entropy( + hidden, + weight, + target, + bias=bias, + tp_group=tp_group, + ignore_index=ignore_index, + ) + return LigerMegatronFusedLinearCrossEntropyFunction.apply( + hidden, + weight, + target, + bias, + tp_group, + ignore_index, + ) diff --git a/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py index c28e2cad4..81e1086b0 100644 --- a/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py @@ -55,11 +55,21 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): grad_hidden = torch.mm(exp_buf, weight, out_dtype=torch.float32) else: grad_hidden = exp_buf.float() @ weight.float() + reduce_work = ( + dist.all_reduce( + grad_hidden, + op=dist.ReduceOp.SUM, + group=ctx.tp_group, + async_op=True, + ) + if ctx.tp_world > 1 + else None + ) grad_weight = exp_buf.t() @ hidden grad_bias = exp_buf.sum(dim=0, dtype=torch.float32).to(ctx.bias_dtype) if ctx.has_bias else None - if ctx.tp_world > 1: - dist.all_reduce(grad_hidden, op=dist.ReduceOp.SUM, group=ctx.tp_group) + if reduce_work is not None: + reduce_work.wait() grad_hidden = grad_hidden.to(ctx.hidden_dtype).reshape(ctx.original_hidden_shape) return grad_hidden, grad_weight, grad_bias @@ -131,14 +141,10 @@ def forward( from liger_kernel.ops.vocab_parallel_cross_entropy import _select_block_size from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_forward_kernel - exp_buf = torch.empty( - hidden_2d.shape[0], - vocab_local, - device=hidden.device, - dtype=hidden.dtype, - ) - predicted_logit = torch.empty(hidden_2d.shape[0], device=hidden.device, dtype=torch.float32) - sum_exp = torch.empty_like(predicted_logit) + exp_buf = logits + stats = torch.empty((2, hidden_2d.shape[0]), device=hidden.device, dtype=torch.float32) + predicted_logit = stats[0] + sum_exp = stats[1] ce_block_size = _select_block_size(vocab_local) num_warps = _get_num_warps(ce_block_size) liger_vocab_parallel_ce_forward_kernel[(hidden_2d.shape[0],)]( @@ -157,8 +163,7 @@ def forward( num_warps=num_warps, ) if tp_world > 1: - dist.all_reduce(predicted_logit, op=dist.ReduceOp.SUM, group=tp_group) - dist.all_reduce(sum_exp, op=dist.ReduceOp.SUM, group=tp_group) + dist.all_reduce(stats, op=dist.ReduceOp.SUM, group=tp_group) loss = torch.log(sum_exp) - predicted_logit loss = torch.where(valid, loss, torch.zeros_like(loss)) diff --git a/src/liger_kernel/ops/triton/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/triton/ops/megatron_fused_linear_cross_entropy.py index 5ffa8947e..5434874bc 100644 --- a/src/liger_kernel/ops/triton/ops/megatron_fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/triton/ops/megatron_fused_linear_cross_entropy.py @@ -459,11 +459,21 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): ) grad_hidden = _triton_dx_matmul(exp_buf, weight) + reduce_work = ( + dist.all_reduce( + grad_hidden, + op=dist.ReduceOp.SUM, + group=ctx.tp_group, + async_op=True, + ) + if ctx.tp_world > 1 + else None + ) grad_weight = _triton_matmul(exp_buf.t(), hidden) grad_bias = _triton_column_sum(exp_buf).to(ctx.bias_dtype) if ctx.has_bias else None - if ctx.tp_world > 1: - dist.all_reduce(grad_hidden, op=dist.ReduceOp.SUM, group=ctx.tp_group) + if reduce_work is not None: + reduce_work.wait() grad_hidden = grad_hidden.to(ctx.hidden_dtype).reshape(ctx.original_hidden_shape) return grad_hidden, grad_weight, grad_bias @@ -533,14 +543,10 @@ def forward( if tp_world > 1: dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group) - exp_buf = torch.empty( - hidden_2d.shape[0], - vocab_local, - device=hidden.device, - dtype=hidden.dtype, - ) - predicted_logit = torch.empty(hidden_2d.shape[0], device=hidden.device, dtype=torch.float32) - sum_exp = torch.empty_like(predicted_logit) + exp_buf = logits + stats = torch.empty((2, hidden_2d.shape[0]), device=hidden.device, dtype=torch.float32) + predicted_logit = stats[0] + sum_exp = stats[1] num_warps = _get_num_warps(ce_block_size) liger_vocab_parallel_ce_forward_kernel[(hidden_2d.shape[0],)]( X_ptr=logits, @@ -558,8 +564,7 @@ def forward( num_warps=num_warps, ) if tp_world > 1: - dist.all_reduce(predicted_logit, op=dist.ReduceOp.SUM, group=tp_group) - dist.all_reduce(sum_exp, op=dist.ReduceOp.SUM, group=tp_group) + dist.all_reduce(stats, op=dist.ReduceOp.SUM, group=tp_group) loss = _triton_loss(sum_exp, predicted_logit, flat_target, ignore_index) diff --git a/test/megatron/test_cutedsl_fused_linear_cross_entropy.py b/test/megatron/test_cutedsl_fused_linear_cross_entropy.py new file mode 100644 index 000000000..af4b57175 --- /dev/null +++ b/test/megatron/test_cutedsl_fused_linear_cross_entropy.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +import os +import subprocess +import sys +import tempfile + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F + +pytest.importorskip("cutlass.cute") + +from liger_kernel.ops.cutedsl.ops.megatron_fused_linear_cross_entropy import ( # noqa: E402 + liger_megatron_fused_linear_cross_entropy, +) + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CuTe DSL FLCE requires CUDA"), + pytest.mark.skipif( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] < 10, + reason="native CuTe DSL Megatron FLCE requires Blackwell", + ), +] + + +def _reference_loss(hidden, weight, target, bias=None, ignore_index=-100): + logits = hidden.float() @ weight.float().t() + if bias is not None: + logits = logits + bias.float() + return F.cross_entropy( + logits.reshape(-1, logits.shape[-1]), + target.reshape(-1), + reduction="none", + ignore_index=ignore_index, + ).reshape(target.shape) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_cutedsl_megatron_flce_tp1_matches_pytorch(dtype, with_bias): + torch.manual_seed(42) + hidden_base = torch.randn(3, 2, 65, device="cuda", dtype=dtype) + weight_base = torch.randn(33, 65, device="cuda", dtype=dtype) * 0.02 + bias_base = torch.randn(33, device="cuda", dtype=dtype) * 0.02 if with_bias else None + target = torch.randint(0, 33, (3, 2), device="cuda") + target[0, 0] = -100 + upstream = torch.randn(3, 2, device="cuda") + + hidden_ref = hidden_base.clone().requires_grad_(True) + weight_ref = weight_base.clone().requires_grad_(True) + bias_ref = bias_base.clone().requires_grad_(True) if bias_base is not None else None + hidden_cutedsl = hidden_base.clone().requires_grad_(True) + weight_cutedsl = weight_base.clone().requires_grad_(True) + bias_cutedsl = bias_base.clone().requires_grad_(True) if bias_base is not None else None + + reference = _reference_loss(hidden_ref, weight_ref, target, bias_ref) + actual = liger_megatron_fused_linear_cross_entropy( + hidden_cutedsl, + weight_cutedsl, + target, + bias=bias_cutedsl, + ) + + torch.testing.assert_close(actual, reference, atol=5e-3, rtol=5e-2) + reference.backward(upstream) + actual.backward(upstream) + torch.testing.assert_close(hidden_cutedsl.grad, hidden_ref.grad, atol=5e-3, rtol=5e-2) + torch.testing.assert_close(weight_cutedsl.grad, weight_ref.grad, atol=5e-3, rtol=5e-2) + if with_bias: + torch.testing.assert_close(bias_cutedsl.grad, bias_ref.grad, atol=5e-3, rtol=5e-2) + + +def test_cutedsl_megatron_flce_backend_dispatch(): + env = os.environ.copy() + env["LIGER_KERNEL_IMPL"] = "cutedsl" + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "from liger_kernel.megatron.fused_linear_cross_entropy import " + "liger_megatron_fused_linear_cross_entropy as fn; print(fn.__module__)" + ), + ], + check=True, + capture_output=True, + text=True, + env=env, + ) + assert result.stdout.strip() == "liger_kernel.ops.cutedsl.ops.megatron_fused_linear_cross_entropy" + + +def _tp_worker(rank, world_size, file_name, dtype): + dist.init_process_group( + backend="nccl", + init_method=f"file://{file_name}", + rank=rank, + world_size=world_size, + ) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + tp_group = dist.group.WORLD + vocab_global = 64 + vocab_local = vocab_global // world_size + + torch.manual_seed(123) + hidden_base = torch.randn(3, 2, 65, device=device, dtype=dtype) + weight_global = torch.randn(vocab_global, 65, device=device, dtype=dtype) * 0.02 + bias_global = torch.randn(vocab_global, device=device, dtype=dtype) * 0.02 + target = torch.randint(0, vocab_global, (3, 2), device=device) + upstream = torch.randn(3, 2, device=device) + target[0, 0] = -100 + for tensor in (hidden_base, weight_global, bias_global, target, upstream): + dist.broadcast(tensor, src=0, group=tp_group) + + start = rank * vocab_local + end = start + vocab_local + hidden_cutedsl = hidden_base.clone().requires_grad_(True) + weight_local = weight_global[start:end].clone().requires_grad_(True) + bias_local = bias_global[start:end].clone().requires_grad_(True) + actual = liger_megatron_fused_linear_cross_entropy( + hidden_cutedsl, + weight_local, + target, + bias=bias_local, + tp_group=tp_group, + ) + + hidden_ref = hidden_base.clone().requires_grad_(True) + weight_ref = weight_global.clone().requires_grad_(True) + bias_ref = bias_global.clone().requires_grad_(True) + reference = _reference_loss(hidden_ref, weight_ref, target, bias_ref) + + torch.testing.assert_close(actual, reference, atol=5e-3, rtol=5e-2) + actual.backward(upstream) + reference.backward(upstream) + torch.testing.assert_close(hidden_cutedsl.grad, hidden_ref.grad, atol=5e-3, rtol=5e-2) + torch.testing.assert_close(weight_local.grad, weight_ref.grad[start:end], atol=5e-3, rtol=5e-2) + torch.testing.assert_close(bias_local.grad, bias_ref.grad[start:end], atol=5e-3, rtol=5e-2) + dist.destroy_process_group() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least two CUDA GPUs") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_cutedsl_megatron_flce_tp2_matches_global_reference(dtype): + with tempfile.NamedTemporaryFile() as rendezvous: + mp.spawn( + _tp_worker, + args=(2, rendezvous.name, dtype), + nprocs=2, + join=True, + ) From 60d578433f4ca0a1d0c09c31cb09cc0301a2d6b7 Mon Sep 17 00:00:00 2001 From: Justin Hu Date: Tue, 11 Aug 2026 17:57:05 +0000 Subject: [PATCH 3/7] refactor: keep Megatron Triton kernel in flat ops layout Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- benchmark/README.md | 11 +- ...ark_megatron_fused_linear_cross_entropy.py | 21 +- docs/High-Level-APIs.md | 7 +- .../megatron_fused_linear_cross_entropy.py | 450 ++++++++++++- src/liger_kernel/ops/triton/__init__.py | 12 - src/liger_kernel/ops/triton/ops/__init__.py | 9 - .../megatron_fused_linear_cross_entropy.py | 605 ------------------ .../test_fused_linear_cross_entropy.py | 52 +- 8 files changed, 446 insertions(+), 721 deletions(-) delete mode 100644 src/liger_kernel/ops/triton/__init__.py delete mode 100644 src/liger_kernel/ops/triton/ops/__init__.py delete mode 100644 src/liger_kernel/ops/triton/ops/megatron_fused_linear_cross_entropy.py diff --git a/benchmark/README.md b/benchmark/README.md index 7afcaa95f..c8f048e50 100644 --- a/benchmark/README.md +++ b/benchmark/README.md @@ -191,14 +191,19 @@ python ../benchmarks_visualizer.py \ ### Megatron fused linear cross-entropy The Megatron FLCE benchmark checks loss and gradient parity before measuring -the default cuBLAS/Triton implementation and the opt-in Triton, CuTile, and -CuTe DSL backends: +the default Triton implementation and the opt-in CuTile and CuTe DSL backends: ```bash cd benchmark/scripts python benchmark_megatron_fused_linear_cross_entropy.py \ --tp-size 4 --token-counts 512 2048 --vocab-sizes 32000 128256 \ - --providers megatron-compatible liger liger-triton liger-cutile + --providers megatron-compatible liger liger-cutile + +# 7B/Llama-3-scale output layer +python benchmark_megatron_fused_linear_cross_entropy.py \ + --tp-size 4 --token-counts 16384 --hidden-size 4096 \ + --vocab-sizes 128256 \ + --providers megatron-compatible liger liger-cutile # SM100 with nvidia-cutlass-dsl installed python benchmark_megatron_fused_linear_cross_entropy.py \ diff --git a/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py b/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py index d13158afc..08a4d8f87 100644 --- a/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py +++ b/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py @@ -6,9 +6,8 @@ vocab-parallel linear -> materialized local logits -> fused vocab-parallel CE This script compares that path with ``LigerMegatronFusedLinearCrossEntropy``, -which saves low-precision CE state to avoid projection recomputation. The -``liger-triton`` provider additionally replaces all three local GEMMs with -portable Triton kernels. ``liger-cutile`` replaces all local compute with +which uses portable Triton kernels and saves low-precision CE state to avoid +projection recomputation. ``liger-cutile`` replaces all local compute with CuTile, while ``liger-cutedsl`` uses a persistent SM100 CuTe DSL projection. When Megatron-Core is installed, the ``megatron-core`` provider uses its fused CE. The always-available ``megatron-compatible`` provider uses Liger's drop-in @@ -48,9 +47,6 @@ from liger_kernel.megatron import LigerMegatronCrossEntropy from liger_kernel.megatron import LigerMegatronFusedLinearCrossEntropy -from liger_kernel.ops.triton.ops.megatron_fused_linear_cross_entropy import ( - liger_megatron_fused_linear_cross_entropy as triton_megatron_fused_linear_cross_entropy, -) try: from liger_kernel.ops.cutedsl.ops.megatron_fused_linear_cross_entropy import ( @@ -125,14 +121,6 @@ def _make_state( if provider == "liger": loss = LigerMegatronFusedLinearCrossEntropy() forward = lambda: loss(hidden, weight, target, bias=bias, tp_group=tp_group) - elif provider == "liger-triton": - forward = lambda: triton_megatron_fused_linear_cross_entropy( - hidden, - weight, - target, - bias=bias, - tp_group=tp_group, - ) elif provider == "liger-cutedsl": if not _CUTEDSL_AVAILABLE: raise RuntimeError("provider 'liger-cutedsl' requires nvidia-cutlass-dsl.") @@ -337,9 +325,7 @@ def _check_correctness( dist.broadcast(upstream, src=0) outputs = {} correctness_providers = ["megatron-compatible"] - correctness_providers.extend( - provider for provider in ("liger", "liger-triton", "liger-cutile") if provider in providers - ) + correctness_providers.extend(provider for provider in ("liger", "liger-cutile") if provider in providers) for provider in correctness_providers: state = _make_state(provider, hidden, weight, bias, target, tp_group, tp_size) loss = state.forward() @@ -503,7 +489,6 @@ def main(): nargs="+", choices=[ "liger", - "liger-triton", "liger-cutedsl", "liger-cutile", "megatron-compatible", diff --git a/docs/High-Level-APIs.md b/docs/High-Level-APIs.md index dae52662d..d0fb89d71 100644 --- a/docs/High-Level-APIs.md +++ b/docs/High-Level-APIs.md @@ -110,10 +110,9 @@ shards. `LigerMegatronFusedLinearCrossEntropy` accepts replicated hidden states, the calling rank's contiguous `[V_local, H]` output-weight shard, and global target indices. It supports TP1 and TP>1 through the supplied process group. -The default implementation uses cuBLAS, Triton CE, and NCCL. Set -`LIGER_KERNEL_IMPL=triton`, `cutile`, or `cutedsl` before importing Liger to -select an all-Triton local path, a CuTile local path, or the SM100 CuTe DSL -persistent projection. +The default implementation uses Triton local kernels and NCCL. Set +`LIGER_KERNEL_IMPL=cutile` or `cutedsl` before importing Liger to select a +CuTile local path or the SM100 CuTe DSL persistent projection. **Usage**: diff --git a/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py index 81e1086b0..5434874bc 100644 --- a/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py @@ -1,21 +1,428 @@ -"""Materialized tensor-parallel fused linear cross entropy for Megatron. +"""Portable all-Triton tensor-parallel fused linear cross entropy for Megatron. Each tensor-parallel rank owns a contiguous vocabulary shard. Forward performs -one local projection GEMM, computes globally normalized cross entropy, and saves -shifted exponentials in the projection dtype. Backward converts that buffer to -dlogits in-place, avoiding projection recomputation before forming dX and dW. +one Triton projection GEMM, computes globally normalized cross entropy, and +saves shifted exponentials in the projection dtype. Backward converts that +buffer to dlogits in-place before Triton dX and dW GEMMs. Tensor-parallel +collectives remain NCCL/RCCL calls between architecture-independent kernels. """ from __future__ import annotations -import operator - import torch import torch.distributed as dist +import triton +import triton.language as tl + + +def _matmul_autotune_configs(): + return [ + triton.Config( + {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 8}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, + num_stages=3, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 4}, + num_stages=3, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 4}, + num_stages=3, + num_warps=4, + ), + ] + + +def _split_k_matmul_autotune_configs(): + return [ + triton.Config( + { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 2, + }, + num_stages=3, + num_warps=4, + ), + triton.Config( + { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 4, + }, + num_stages=3, + num_warps=4, + ), + triton.Config( + { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 8, + }, + num_stages=3, + num_warps=4, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 4, + }, + num_stages=3, + num_warps=8, + ), + triton.Config( + { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 4, + }, + num_stages=3, + num_warps=8, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 4, + }, + num_stages=3, + num_warps=8, + ), + ] + + +@triton.autotune(configs=_matmul_autotune_configs(), key=["M", "N", "K"]) +@triton.jit +def _matmul_kernel( + a_ptr, + b_ptr, + bias_ptr, + output_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_om, + stride_on, + HAS_BIAS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % num_pid_in_group) % group_size_m + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k_start in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + k_remaining = K - k_start * BLOCK_SIZE_K + a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < k_remaining), other=0.0) + b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_remaining) & (offs_n[None, :] < N), other=0.0) + accumulator = tl.dot(a, b, accumulator) + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if HAS_BIAS: + bias = tl.load(bias_ptr + offs_n, mask=offs_n < N, other=0.0).to(tl.float32) + accumulator += bias[None, :] + + output_ptrs = output_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on + tl.store(output_ptrs, accumulator, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N)) + + +@triton.autotune( + configs=_split_k_matmul_autotune_configs(), + key=["M", "N", "K"], + reset_to_zero=["output_ptr"], +) +@triton.jit +def _split_k_matmul_kernel( + a_ptr, + b_ptr, + output_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_om, + stride_on, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + SPLIT_K: tl.constexpr, +): + pid = tl.program_id(0) + split_k_id = tl.program_id(1) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % num_pid_in_group) % group_size_m + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = split_k_id * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k_start in range(0, tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)): + k_remaining = K - k_start * BLOCK_SIZE_K * SPLIT_K + a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < k_remaining), other=0.0) + b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_remaining) & (offs_n[None, :] < N), other=0.0) + accumulator = tl.dot(a, b, accumulator) + a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak + b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk + + output_ptrs = output_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on + tl.atomic_add(output_ptrs, accumulator, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N)) + + +@triton.jit +def _row_max_kernel( + input_ptr, + output_ptr, + n_cols, + input_row_stride, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + row_ptr = input_ptr + row * input_row_stride + row_max = -float("inf") + for start in range(0, n_cols, BLOCK_SIZE): + offsets = start + tl.arange(0, BLOCK_SIZE) + values = tl.load(row_ptr + offsets, mask=offsets < n_cols, other=-float("inf")).to(tl.float32) + row_max = tl.maximum(row_max, tl.max(values)) + tl.store(output_ptr + row, row_max) + + +@triton.jit +def _loss_kernel( + sum_exp_ptr, + predicted_logit_ptr, + target_ptr, + loss_ptr, + n_rows, + ignore_index, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0).to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_rows + sum_exp = tl.load(sum_exp_ptr + offsets, mask=mask, other=1.0) + predicted_logit = tl.load(predicted_logit_ptr + offsets, mask=mask, other=0.0) + target = tl.load(target_ptr + offsets, mask=mask, other=ignore_index) + loss = tl.log(sum_exp) - predicted_logit + loss = tl.where(target == ignore_index, 0.0, loss) + tl.store(loss_ptr + offsets, loss, mask=mask) + + +@triton.jit +def _column_sum_kernel( + input_ptr, + output_ptr, + n_rows, + input_row_stride, + BLOCK_SIZE: tl.constexpr, +): + col = tl.program_id(0).to(tl.int64) + offsets = tl.arange(0, BLOCK_SIZE) + total = 0.0 + for start in range(0, n_rows, BLOCK_SIZE): + rows = start + offsets + values = tl.load(input_ptr + rows * input_row_stride + col, mask=rows < n_rows, other=0.0) + total += tl.sum(values.to(tl.float32)) + tl.store(output_ptr + col, total) + + +def _triton_matmul( + a: torch.Tensor, + b: torch.Tensor, + *, + bias: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, +) -> torch.Tensor: + if a.ndim != 2 or b.ndim != 2 or a.shape[1] != b.shape[0]: + raise ValueError(f"matmul expects [M, K] @ [K, N], got {tuple(a.shape)} and {tuple(b.shape)}.") + m, k = a.shape + n = b.shape[1] + output = torch.empty((m, n), device=a.device, dtype=output_dtype or a.dtype) + grid = lambda meta: (triton.cdiv(m, meta["BLOCK_SIZE_M"]) * triton.cdiv(n, meta["BLOCK_SIZE_N"]),) + _matmul_kernel[grid]( + a, + b, + bias if bias is not None else output, + output, + m, + n, + k, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + output.stride(0), + output.stride(1), + HAS_BIAS=bias is not None, + ) + return output + -from liger_kernel.ops.utils import compare_version +def _triton_dx_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + m, k = a.shape + n = b.shape[1] + if m > 1024 or k < 4096: + return _triton_matmul(a, b, output_dtype=torch.float32) -_SUPPORTS_OUT_DTYPE = compare_version("torch", operator.ge, "2.8.0") + output = torch.zeros((m, n), device=a.device, dtype=torch.float32) + grid = lambda meta: ( + triton.cdiv(m, meta["BLOCK_SIZE_M"]) * triton.cdiv(n, meta["BLOCK_SIZE_N"]), + meta["SPLIT_K"], + ) + _split_k_matmul_kernel[grid]( + a, + b, + output, + m, + n, + k, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + output.stride(0), + output.stride(1), + ) + return output + + +def _triton_row_max(input: torch.Tensor, block_size: int) -> torch.Tensor: + from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps + + output = torch.empty(input.shape[0], device=input.device, dtype=torch.float32) + _row_max_kernel[(input.shape[0],)]( + input, + output, + input.shape[1], + input.stride(0), + BLOCK_SIZE=block_size, + num_warps=_get_num_warps(block_size), + ) + return output + + +def _triton_loss( + sum_exp: torch.Tensor, + predicted_logit: torch.Tensor, + target: torch.Tensor, + ignore_index: int, +) -> torch.Tensor: + output = torch.empty_like(sum_exp) + block_size = 256 + _loss_kernel[(triton.cdiv(target.numel(), block_size),)]( + sum_exp, + predicted_logit, + target, + output, + target.numel(), + ignore_index, + BLOCK_SIZE=block_size, + num_warps=4, + ) + return output + + +def _triton_column_sum(input: torch.Tensor) -> torch.Tensor: + from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps + from liger_kernel.ops.vocab_parallel_cross_entropy import _select_block_size + + block_size = _select_block_size(input.shape[0]) + output = torch.empty(input.shape[1], device=input.device, dtype=torch.float32) + _column_sum_kernel[(input.shape[1],)]( + input, + output, + input.shape[0], + input.stride(0), + BLOCK_SIZE=block_size, + num_warps=_get_num_warps(block_size), + ) + return output def _tp_rank_and_world(tp_group) -> tuple[int, int]: @@ -51,10 +458,7 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): num_warps=num_warps, ) - if _SUPPORTS_OUT_DTYPE: - grad_hidden = torch.mm(exp_buf, weight, out_dtype=torch.float32) - else: - grad_hidden = exp_buf.float() @ weight.float() + grad_hidden = _triton_dx_matmul(exp_buf, weight) reduce_work = ( dist.all_reduce( grad_hidden, @@ -65,8 +469,8 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): if ctx.tp_world > 1 else None ) - grad_weight = exp_buf.t() @ hidden - grad_bias = exp_buf.sum(dim=0, dtype=torch.float32).to(ctx.bias_dtype) if ctx.has_bias else None + grad_weight = _triton_matmul(exp_buf.t(), hidden) + grad_bias = _triton_column_sum(exp_buf).to(ctx.bias_dtype) if ctx.has_bias else None if reduce_work is not None: reduce_work.wait() @@ -129,23 +533,20 @@ def forward( weight_2d = weight.contiguous() bias_1d = bias.contiguous() if bias is not None else None - logits = torch.mm(hidden_2d, weight_2d.t()) - if bias_1d is not None: - logits.add_(bias_1d) - - logits_max = logits.amax(dim=-1).float() - if tp_world > 1: - dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group) - from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps from liger_kernel.ops.vocab_parallel_cross_entropy import _select_block_size from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_forward_kernel + logits = _triton_matmul(hidden_2d, weight_2d.t(), bias=bias_1d) + ce_block_size = _select_block_size(vocab_local) + logits_max = _triton_row_max(logits, ce_block_size) + if tp_world > 1: + dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group) + exp_buf = logits stats = torch.empty((2, hidden_2d.shape[0]), device=hidden.device, dtype=torch.float32) predicted_logit = stats[0] sum_exp = stats[1] - ce_block_size = _select_block_size(vocab_local) num_warps = _get_num_warps(ce_block_size) liger_vocab_parallel_ce_forward_kernel[(hidden_2d.shape[0],)]( X_ptr=logits, @@ -165,8 +566,7 @@ def forward( if tp_world > 1: dist.all_reduce(stats, op=dist.ReduceOp.SUM, group=tp_group) - loss = torch.log(sum_exp) - predicted_logit - loss = torch.where(valid, loss, torch.zeros_like(loss)) + loss = _triton_loss(sum_exp, predicted_logit, flat_target, ignore_index) ctx.save_for_backward(hidden_2d, weight_2d, exp_buf, sum_exp, flat_target) ctx.has_bias = bias is not None diff --git a/src/liger_kernel/ops/triton/__init__.py b/src/liger_kernel/ops/triton/__init__.py deleted file mode 100644 index 0de2a07af..000000000 --- a/src/liger_kernel/ops/triton/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -"""Opt-in portable Triton implementations.""" - -from liger_kernel.ops.backends.registry import ImplInfo -from liger_kernel.ops.backends.registry import register_impl - -register_impl( - ImplInfo( - name="triton", - devices=("cuda",), - module_path=f"{__name__}.ops", - ) -) diff --git a/src/liger_kernel/ops/triton/ops/__init__.py b/src/liger_kernel/ops/triton/ops/__init__.py deleted file mode 100644 index f757b8cf8..000000000 --- a/src/liger_kernel/ops/triton/ops/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -"""Operators replaced by the opt-in all-Triton implementation.""" - -from liger_kernel.ops.triton.ops.megatron_fused_linear_cross_entropy import LigerMegatronFusedLinearCrossEntropyFunction -from liger_kernel.ops.triton.ops.megatron_fused_linear_cross_entropy import liger_megatron_fused_linear_cross_entropy - -__all__ = [ - "LigerMegatronFusedLinearCrossEntropyFunction", - "liger_megatron_fused_linear_cross_entropy", -] diff --git a/src/liger_kernel/ops/triton/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/triton/ops/megatron_fused_linear_cross_entropy.py deleted file mode 100644 index 5434874bc..000000000 --- a/src/liger_kernel/ops/triton/ops/megatron_fused_linear_cross_entropy.py +++ /dev/null @@ -1,605 +0,0 @@ -"""Portable all-Triton tensor-parallel fused linear cross entropy for Megatron. - -Each tensor-parallel rank owns a contiguous vocabulary shard. Forward performs -one Triton projection GEMM, computes globally normalized cross entropy, and -saves shifted exponentials in the projection dtype. Backward converts that -buffer to dlogits in-place before Triton dX and dW GEMMs. Tensor-parallel -collectives remain NCCL/RCCL calls between architecture-independent kernels. -""" - -from __future__ import annotations - -import torch -import torch.distributed as dist -import triton -import triton.language as tl - - -def _matmul_autotune_configs(): - return [ - triton.Config( - {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 8}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, - num_stages=3, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 4}, - num_stages=3, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 4}, - num_stages=3, - num_warps=4, - ), - ] - - -def _split_k_matmul_autotune_configs(): - return [ - triton.Config( - { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 64, - "GROUP_SIZE_M": 8, - "SPLIT_K": 2, - }, - num_stages=3, - num_warps=4, - ), - triton.Config( - { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 64, - "GROUP_SIZE_M": 8, - "SPLIT_K": 4, - }, - num_stages=3, - num_warps=4, - ), - triton.Config( - { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 64, - "GROUP_SIZE_M": 8, - "SPLIT_K": 8, - }, - num_stages=3, - num_warps=4, - ), - triton.Config( - { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 64, - "GROUP_SIZE_M": 8, - "SPLIT_K": 4, - }, - num_stages=3, - num_warps=8, - ), - triton.Config( - { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 256, - "BLOCK_SIZE_K": 64, - "GROUP_SIZE_M": 8, - "SPLIT_K": 4, - }, - num_stages=3, - num_warps=8, - ), - triton.Config( - { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 256, - "BLOCK_SIZE_K": 64, - "GROUP_SIZE_M": 8, - "SPLIT_K": 4, - }, - num_stages=3, - num_warps=8, - ), - ] - - -@triton.autotune(configs=_matmul_autotune_configs(), key=["M", "N", "K"]) -@triton.jit -def _matmul_kernel( - a_ptr, - b_ptr, - bias_ptr, - output_ptr, - M, - N, - K, - stride_am, - stride_ak, - stride_bk, - stride_bn, - stride_om, - stride_on, - HAS_BIAS: tl.constexpr, - BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr, -): - pid = tl.program_id(0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + (pid % num_pid_in_group) % group_size_m - pid_n = (pid % num_pid_in_group) // group_size_m - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - offs_k = tl.arange(0, BLOCK_SIZE_K) - a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak - b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for k_start in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - k_remaining = K - k_start * BLOCK_SIZE_K - a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < k_remaining), other=0.0) - b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_remaining) & (offs_n[None, :] < N), other=0.0) - accumulator = tl.dot(a, b, accumulator) - a_ptrs += BLOCK_SIZE_K * stride_ak - b_ptrs += BLOCK_SIZE_K * stride_bk - - if HAS_BIAS: - bias = tl.load(bias_ptr + offs_n, mask=offs_n < N, other=0.0).to(tl.float32) - accumulator += bias[None, :] - - output_ptrs = output_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on - tl.store(output_ptrs, accumulator, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N)) - - -@triton.autotune( - configs=_split_k_matmul_autotune_configs(), - key=["M", "N", "K"], - reset_to_zero=["output_ptr"], -) -@triton.jit -def _split_k_matmul_kernel( - a_ptr, - b_ptr, - output_ptr, - M, - N, - K, - stride_am, - stride_ak, - stride_bk, - stride_bn, - stride_om, - stride_on, - BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr, - SPLIT_K: tl.constexpr, -): - pid = tl.program_id(0) - split_k_id = tl.program_id(1) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + (pid % num_pid_in_group) % group_size_m - pid_n = (pid % num_pid_in_group) // group_size_m - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - offs_k = split_k_id * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) - a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak - b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for k_start in range(0, tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)): - k_remaining = K - k_start * BLOCK_SIZE_K * SPLIT_K - a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < k_remaining), other=0.0) - b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_remaining) & (offs_n[None, :] < N), other=0.0) - accumulator = tl.dot(a, b, accumulator) - a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak - b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk - - output_ptrs = output_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on - tl.atomic_add(output_ptrs, accumulator, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N)) - - -@triton.jit -def _row_max_kernel( - input_ptr, - output_ptr, - n_cols, - input_row_stride, - BLOCK_SIZE: tl.constexpr, -): - row = tl.program_id(0).to(tl.int64) - row_ptr = input_ptr + row * input_row_stride - row_max = -float("inf") - for start in range(0, n_cols, BLOCK_SIZE): - offsets = start + tl.arange(0, BLOCK_SIZE) - values = tl.load(row_ptr + offsets, mask=offsets < n_cols, other=-float("inf")).to(tl.float32) - row_max = tl.maximum(row_max, tl.max(values)) - tl.store(output_ptr + row, row_max) - - -@triton.jit -def _loss_kernel( - sum_exp_ptr, - predicted_logit_ptr, - target_ptr, - loss_ptr, - n_rows, - ignore_index, - BLOCK_SIZE: tl.constexpr, -): - offsets = tl.program_id(0).to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_rows - sum_exp = tl.load(sum_exp_ptr + offsets, mask=mask, other=1.0) - predicted_logit = tl.load(predicted_logit_ptr + offsets, mask=mask, other=0.0) - target = tl.load(target_ptr + offsets, mask=mask, other=ignore_index) - loss = tl.log(sum_exp) - predicted_logit - loss = tl.where(target == ignore_index, 0.0, loss) - tl.store(loss_ptr + offsets, loss, mask=mask) - - -@triton.jit -def _column_sum_kernel( - input_ptr, - output_ptr, - n_rows, - input_row_stride, - BLOCK_SIZE: tl.constexpr, -): - col = tl.program_id(0).to(tl.int64) - offsets = tl.arange(0, BLOCK_SIZE) - total = 0.0 - for start in range(0, n_rows, BLOCK_SIZE): - rows = start + offsets - values = tl.load(input_ptr + rows * input_row_stride + col, mask=rows < n_rows, other=0.0) - total += tl.sum(values.to(tl.float32)) - tl.store(output_ptr + col, total) - - -def _triton_matmul( - a: torch.Tensor, - b: torch.Tensor, - *, - bias: torch.Tensor | None = None, - output_dtype: torch.dtype | None = None, -) -> torch.Tensor: - if a.ndim != 2 or b.ndim != 2 or a.shape[1] != b.shape[0]: - raise ValueError(f"matmul expects [M, K] @ [K, N], got {tuple(a.shape)} and {tuple(b.shape)}.") - m, k = a.shape - n = b.shape[1] - output = torch.empty((m, n), device=a.device, dtype=output_dtype or a.dtype) - grid = lambda meta: (triton.cdiv(m, meta["BLOCK_SIZE_M"]) * triton.cdiv(n, meta["BLOCK_SIZE_N"]),) - _matmul_kernel[grid]( - a, - b, - bias if bias is not None else output, - output, - m, - n, - k, - a.stride(0), - a.stride(1), - b.stride(0), - b.stride(1), - output.stride(0), - output.stride(1), - HAS_BIAS=bias is not None, - ) - return output - - -def _triton_dx_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - m, k = a.shape - n = b.shape[1] - if m > 1024 or k < 4096: - return _triton_matmul(a, b, output_dtype=torch.float32) - - output = torch.zeros((m, n), device=a.device, dtype=torch.float32) - grid = lambda meta: ( - triton.cdiv(m, meta["BLOCK_SIZE_M"]) * triton.cdiv(n, meta["BLOCK_SIZE_N"]), - meta["SPLIT_K"], - ) - _split_k_matmul_kernel[grid]( - a, - b, - output, - m, - n, - k, - a.stride(0), - a.stride(1), - b.stride(0), - b.stride(1), - output.stride(0), - output.stride(1), - ) - return output - - -def _triton_row_max(input: torch.Tensor, block_size: int) -> torch.Tensor: - from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps - - output = torch.empty(input.shape[0], device=input.device, dtype=torch.float32) - _row_max_kernel[(input.shape[0],)]( - input, - output, - input.shape[1], - input.stride(0), - BLOCK_SIZE=block_size, - num_warps=_get_num_warps(block_size), - ) - return output - - -def _triton_loss( - sum_exp: torch.Tensor, - predicted_logit: torch.Tensor, - target: torch.Tensor, - ignore_index: int, -) -> torch.Tensor: - output = torch.empty_like(sum_exp) - block_size = 256 - _loss_kernel[(triton.cdiv(target.numel(), block_size),)]( - sum_exp, - predicted_logit, - target, - output, - target.numel(), - ignore_index, - BLOCK_SIZE=block_size, - num_warps=4, - ) - return output - - -def _triton_column_sum(input: torch.Tensor) -> torch.Tensor: - from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps - from liger_kernel.ops.vocab_parallel_cross_entropy import _select_block_size - - block_size = _select_block_size(input.shape[0]) - output = torch.empty(input.shape[1], device=input.device, dtype=torch.float32) - _column_sum_kernel[(input.shape[1],)]( - input, - output, - input.shape[0], - input.stride(0), - BLOCK_SIZE=block_size, - num_warps=_get_num_warps(block_size), - ) - return output - - -def _tp_rank_and_world(tp_group) -> tuple[int, int]: - if tp_group is None: - return 0, 1 - world = dist.get_world_size(tp_group) - if world == 1: - return 0, 1 - return dist.get_rank(tp_group), world - - -def _materialized_backward(ctx, grad_output: torch.Tensor): - """Convert saved CE state to dlogits and form projection gradients.""" - from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps - from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_backward_kernel - - hidden, weight, exp_buf, sum_exp_global, target = ctx.saved_tensors - grad_out = grad_output.contiguous().reshape(-1).float() - num_warps = _get_num_warps(ctx.ce_block_size) - liger_vocab_parallel_ce_backward_kernel[(hidden.shape[0],)]( - EXP_ptr=exp_buf, - EXP_stride=exp_buf.stride(0), - sum_exp_ptr=sum_exp_global, - Y_ptr=target, - grad_out_ptr=grad_out, - vocab_start=ctx.vocab_start, - n_cols=weight.shape[0], - ignore_index=ctx.ignore_index, - alpha_eff=0.0, - eps_eff=0.0, - HAS_LABEL_SMOOTHING=False, - BLOCK_SIZE=ctx.ce_block_size, - num_warps=num_warps, - ) - - grad_hidden = _triton_dx_matmul(exp_buf, weight) - reduce_work = ( - dist.all_reduce( - grad_hidden, - op=dist.ReduceOp.SUM, - group=ctx.tp_group, - async_op=True, - ) - if ctx.tp_world > 1 - else None - ) - grad_weight = _triton_matmul(exp_buf.t(), hidden) - grad_bias = _triton_column_sum(exp_buf).to(ctx.bias_dtype) if ctx.has_bias else None - - if reduce_work is not None: - reduce_work.wait() - grad_hidden = grad_hidden.to(ctx.hidden_dtype).reshape(ctx.original_hidden_shape) - return grad_hidden, grad_weight, grad_bias - - -class LigerMegatronFusedLinearCrossEntropyFunction(torch.autograd.Function): - """Hidden-to-loss tensor-parallel FLCE with saved low-precision CE state.""" - - @staticmethod - def forward( - ctx, - hidden: torch.Tensor, - weight: torch.Tensor, - target: torch.Tensor, - bias: torch.Tensor | None, - tp_group, - ignore_index: int, - ) -> torch.Tensor: - if hidden.ndim < 2: - raise ValueError(f"hidden must have at least 2 dimensions, got shape {tuple(hidden.shape)}.") - if weight.ndim != 2: - raise ValueError(f"weight must be 2-D [V_local, H], got shape {tuple(weight.shape)}.") - if tuple(target.shape) != tuple(hidden.shape[:-1]): - raise ValueError( - f"target shape must equal hidden.shape[:-1]; got target={tuple(target.shape)}, " - f"hidden={tuple(hidden.shape)}." - ) - if hidden.shape[-1] != weight.shape[1]: - raise ValueError(f"hidden size mismatch: hidden has H={hidden.shape[-1]}, weight has H={weight.shape[1]}.") - if hidden.dtype != weight.dtype: - raise TypeError(f"hidden and weight must have the same dtype, got {hidden.dtype} and {weight.dtype}.") - if hidden.device != weight.device or hidden.device != target.device: - raise ValueError("hidden, weight, and target must be on the same device.") - if bias is not None: - if bias.ndim != 1 or bias.shape[0] != weight.shape[0]: - raise ValueError(f"bias must have shape ({weight.shape[0]},), got {tuple(bias.shape)}.") - if bias.device != hidden.device or bias.dtype != hidden.dtype: - raise TypeError("bias must have the same device and dtype as hidden.") - if hidden.device.type != "cuda" or hidden.dtype not in (torch.bfloat16, torch.float16): - raise RuntimeError("Megatron FLCE requires a CUDA GPU and float16 or bfloat16 inputs.") - - tp_rank, tp_world = _tp_rank_and_world(tp_group) - vocab_local = weight.shape[0] - vocab_global = vocab_local * tp_world - vocab_start = tp_rank * vocab_local - - flat_target = target.reshape(-1).to(torch.int64).contiguous() - valid = flat_target != ignore_index - invalid = valid & ((flat_target < 0) | (flat_target >= vocab_global)) - valid_targets = ~torch.any(invalid) - if hasattr(torch, "_assert_async"): - torch._assert_async(valid_targets, f"non-ignored targets must be in [0, {vocab_global}).") - elif not valid_targets.item(): - raise ValueError(f"non-ignored targets must be in [0, {vocab_global}).") - - original_hidden_shape = hidden.shape - hidden_2d = hidden.reshape(-1, hidden.shape[-1]).contiguous() - weight_2d = weight.contiguous() - bias_1d = bias.contiguous() if bias is not None else None - - from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps - from liger_kernel.ops.vocab_parallel_cross_entropy import _select_block_size - from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_forward_kernel - - logits = _triton_matmul(hidden_2d, weight_2d.t(), bias=bias_1d) - ce_block_size = _select_block_size(vocab_local) - logits_max = _triton_row_max(logits, ce_block_size) - if tp_world > 1: - dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group) - - exp_buf = logits - stats = torch.empty((2, hidden_2d.shape[0]), device=hidden.device, dtype=torch.float32) - predicted_logit = stats[0] - sum_exp = stats[1] - num_warps = _get_num_warps(ce_block_size) - liger_vocab_parallel_ce_forward_kernel[(hidden_2d.shape[0],)]( - X_ptr=logits, - X_stride=logits.stride(0), - EXP_ptr=exp_buf, - EXP_stride=exp_buf.stride(0), - logits_max_ptr=logits_max, - Y_ptr=flat_target, - pred_ptr=predicted_logit, - sum_exp_ptr=sum_exp, - vocab_start=vocab_start, - n_cols=vocab_local, - ignore_index=ignore_index, - BLOCK_SIZE=ce_block_size, - num_warps=num_warps, - ) - if tp_world > 1: - dist.all_reduce(stats, op=dist.ReduceOp.SUM, group=tp_group) - - loss = _triton_loss(sum_exp, predicted_logit, flat_target, ignore_index) - - ctx.save_for_backward(hidden_2d, weight_2d, exp_buf, sum_exp, flat_target) - ctx.has_bias = bias is not None - ctx.bias_dtype = bias.dtype if bias is not None else None - ctx.tp_group = tp_group - ctx.tp_world = tp_world - ctx.vocab_start = vocab_start - ctx.ignore_index = ignore_index - ctx.ce_block_size = ce_block_size - ctx.original_hidden_shape = original_hidden_shape - ctx.hidden_dtype = hidden.dtype - return loss.reshape(target.shape) - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - grad_hidden, grad_weight, grad_bias = _materialized_backward(ctx, grad_output) - return grad_hidden, grad_weight, None, grad_bias, None, None - - -def liger_megatron_fused_linear_cross_entropy( - hidden: torch.Tensor, - weight: torch.Tensor, - target: torch.Tensor, - bias: torch.Tensor | None = None, - tp_group=None, - ignore_index: int = -100, -) -> torch.Tensor: - """Compute per-token loss from replicated hidden states and a local vocab shard.""" - return LigerMegatronFusedLinearCrossEntropyFunction.apply( - hidden, - weight, - target, - bias, - tp_group, - ignore_index, - ) diff --git a/test/megatron/test_fused_linear_cross_entropy.py b/test/megatron/test_fused_linear_cross_entropy.py index 39453870b..0bc07d69b 100644 --- a/test/megatron/test_fused_linear_cross_entropy.py +++ b/test/megatron/test_fused_linear_cross_entropy.py @@ -1,8 +1,5 @@ from __future__ import annotations -import os -import subprocess -import sys import tempfile import pytest @@ -13,14 +10,6 @@ from liger_kernel.megatron import LigerMegatronFusedLinearCrossEntropy from liger_kernel.ops.megatron_fused_linear_cross_entropy import liger_megatron_fused_linear_cross_entropy -from liger_kernel.ops.triton.ops.megatron_fused_linear_cross_entropy import ( - liger_megatron_fused_linear_cross_entropy as triton_megatron_fused_linear_cross_entropy, -) - -_IMPLEMENTATIONS = [ - pytest.param(liger_megatron_fused_linear_cross_entropy, id="cublas"), - pytest.param(triton_megatron_fused_linear_cross_entropy, id="triton"), -] def _reference_loss(hidden, weight, target, bias=None, ignore_index=-100): @@ -39,8 +28,7 @@ def _reference_loss(hidden, weight, target, bias=None, ignore_index=-100): @pytest.mark.parametrize("shape", [(2, 3, 8, 16), (3, 2, 17, 32)]) @pytest.mark.parametrize("with_bias", [False, True]) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -@pytest.mark.parametrize("implementation", _IMPLEMENTATIONS) -def test_megatron_flce_tp1_matches_pytorch(shape, with_bias, dtype, implementation): +def test_megatron_flce_tp1_matches_pytorch(shape, with_bias, dtype): s, b, h, v = shape torch.manual_seed(42) hidden_base = torch.randn(s, b, h, device="cuda", dtype=dtype) @@ -58,7 +46,7 @@ def test_megatron_flce_tp1_matches_pytorch(shape, with_bias, dtype, implementati bias_liger = bias_base.clone().requires_grad_(True) if bias_base is not None else None reference = _reference_loss(hidden_ref, weight_ref, target, bias_ref) - actual = implementation( + actual = liger_megatron_fused_linear_cross_entropy( hidden_liger, weight_liger, target, @@ -103,7 +91,7 @@ def test_megatron_flce_triton_split_k_dx_matches_pytorch(dtype): weight_triton = weight_base.clone().requires_grad_(True) reference = _reference_loss(hidden_ref, weight_ref, target) - actual = triton_megatron_fused_linear_cross_entropy(hidden_triton, weight_triton, target) + actual = liger_megatron_fused_linear_cross_entropy(hidden_triton, weight_triton, target) torch.testing.assert_close(actual, reference, atol=5e-3, rtol=5e-2) reference.backward(upstream) @@ -112,26 +100,6 @@ def test_megatron_flce_triton_split_k_dx_matches_pytorch(dtype): torch.testing.assert_close(weight_triton.grad, weight_ref.grad, atol=5e-3, rtol=5e-2) -def test_megatron_flce_triton_backend_dispatch(): - env = os.environ.copy() - env["LIGER_KERNEL_IMPL"] = "triton" - result = subprocess.run( - [ - sys.executable, - "-c", - ( - "from liger_kernel.megatron.fused_linear_cross_entropy import " - "liger_megatron_fused_linear_cross_entropy as fn; print(fn.__module__)" - ), - ], - check=True, - capture_output=True, - text=True, - env=env, - ) - assert result.stdout.strip() == "liger_kernel.ops.triton.ops.megatron_fused_linear_cross_entropy" - - def test_megatron_flce_rejects_cpu_inputs(): hidden = torch.randn(2, 3, 8, dtype=torch.bfloat16) weight = torch.randn(16, 8, dtype=torch.bfloat16) @@ -141,7 +109,7 @@ def test_megatron_flce_rejects_cpu_inputs(): liger_megatron_fused_linear_cross_entropy(hidden, weight, target) -def _tp_worker(rank, world_size, file_name, dtype, implementation_name): +def _tp_worker(rank, world_size, file_name, dtype): dist.init_process_group( backend="nccl", init_method=f"file://{file_name}", @@ -151,11 +119,6 @@ def _tp_worker(rank, world_size, file_name, dtype, implementation_name): torch.cuda.set_device(rank) device = torch.device("cuda", rank) tp_group = dist.group.WORLD - implementation = ( - triton_megatron_fused_linear_cross_entropy - if implementation_name == "triton" - else liger_megatron_fused_linear_cross_entropy - ) s, b, h, v_global = 3, 2, 17, 32 v_local = v_global // world_size @@ -175,7 +138,7 @@ def _tp_worker(rank, world_size, file_name, dtype, implementation_name): weight_local = weight_global[start:end].clone().requires_grad_(True) bias_local = bias_global[start:end].clone().requires_grad_(True) - actual = implementation( + actual = liger_megatron_fused_linear_cross_entropy( hidden_liger, weight_local, target, @@ -202,12 +165,11 @@ def _tp_worker(rank, world_size, file_name, dtype, implementation_name): reason="requires at least two CUDA GPUs", ) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -@pytest.mark.parametrize("implementation_name", ["cublas", "triton"]) -def test_megatron_flce_tp2_matches_global_reference(dtype, implementation_name): +def test_megatron_flce_tp2_matches_global_reference(dtype): with tempfile.NamedTemporaryFile() as rendezvous: mp.spawn( _tp_worker, - args=(2, rendezvous.name, dtype, implementation_name), + args=(2, rendezvous.name, dtype), nprocs=2, join=True, ) From 2ee1f696743f24755b2948d1ba9692c3e39a4d1d Mon Sep 17 00:00:00 2001 From: Justin Hu Date: Tue, 11 Aug 2026 18:42:03 +0000 Subject: [PATCH 4/7] fix: preserve logits in Megatron FLCE backward Recompute FP32 exponentials from saved low-precision logits instead of persisting quantized exponentials. Align backend dispatch with the standard Function-class export and harden benchmark correctness checks. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...ark_megatron_fused_linear_cross_entropy.py | 18 +- .../megatron/fused_linear_cross_entropy.py | 10 +- src/liger_kernel/ops/__init__.py | 8 +- .../megatron_fused_linear_cross_entropy.py | 75 +++---- .../megatron_fused_linear_cross_entropy.py | 56 ++--- .../megatron_fused_linear_cross_entropy.py | 201 ++++++++++++++---- ...test_cutedsl_fused_linear_cross_entropy.py | 2 +- .../test_cutile_fused_linear_cross_entropy.py | 2 +- 8 files changed, 235 insertions(+), 137 deletions(-) diff --git a/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py b/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py index 08a4d8f87..bc6e9d29e 100644 --- a/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py +++ b/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py @@ -32,6 +32,7 @@ import os import tempfile +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path @@ -91,9 +92,9 @@ class _ProviderState: weight: torch.Tensor bias: torch.Tensor | None target: torch.Tensor - forward: object + forward: Callable[[], torch.Tensor] - def clear_grads(self): + def clear_grads(self) -> None: self.hidden.grad = None self.weight.grad = None if self.bias is not None: @@ -325,7 +326,9 @@ def _check_correctness( dist.broadcast(upstream, src=0) outputs = {} correctness_providers = ["megatron-compatible"] - correctness_providers.extend(provider for provider in ("liger", "liger-cutile") if provider in providers) + correctness_providers.extend( + provider for provider in ("liger", "liger-cutile", "liger-cutedsl") if provider in providers + ) for provider in correctness_providers: state = _make_state(provider, hidden, weight, bias, target, tp_group, tp_size) loss = state.forward() @@ -512,6 +515,15 @@ def main(): raise RuntimeError("provider 'megatron-core' requested, but megatron-core is not installed.") if "liger-cutedsl" in providers and not _CUTEDSL_AVAILABLE: raise RuntimeError("provider 'liger-cutedsl' requested, but nvidia-cutlass-dsl is not installed.") + if "liger-cutedsl" in providers: + unsupported_devices = [ + index for index in range(args.tp_size) if torch.cuda.get_device_capability(index)[0] < 10 + ] + if unsupported_devices: + raise RuntimeError( + "provider 'liger-cutedsl' requires SM100 or newer; " + f"unsupported CUDA device indices: {unsupported_devices}." + ) if "liger-cutile" in providers and not _CUTILE_AVAILABLE: raise RuntimeError("provider 'liger-cutile' requested, but cuda-tile is not installed.") if min(args.token_counts) <= 0 or args.hidden_size <= 0 or min(args.vocab_sizes) <= 0: diff --git a/src/liger_kernel/megatron/fused_linear_cross_entropy.py b/src/liger_kernel/megatron/fused_linear_cross_entropy.py index ea3fd6ed7..eb87c627c 100644 --- a/src/liger_kernel/megatron/fused_linear_cross_entropy.py +++ b/src/liger_kernel/megatron/fused_linear_cross_entropy.py @@ -5,7 +5,7 @@ import torch import torch.nn as nn -from liger_kernel.ops import liger_megatron_fused_linear_cross_entropy +from liger_kernel.ops import LigerMegatronFusedLinearCrossEntropyFunction class LigerMegatronFusedLinearCrossEntropy(nn.Module): @@ -30,13 +30,13 @@ def forward( bias: torch.Tensor | None = None, tp_group=None, ) -> torch.Tensor: - return liger_megatron_fused_linear_cross_entropy( + return LigerMegatronFusedLinearCrossEntropyFunction.apply( hidden, weight, target, - bias=bias, - tp_group=tp_group, - ignore_index=self.ignore_index, + bias, + tp_group, + self.ignore_index, ) def extra_repr(self) -> str: diff --git a/src/liger_kernel/ops/__init__.py b/src/liger_kernel/ops/__init__.py index 19b1f7676..ee78f04e0 100644 --- a/src/liger_kernel/ops/__init__.py +++ b/src/liger_kernel/ops/__init__.py @@ -67,12 +67,10 @@ from liger_kernel.ops.layer_norm import layer_norm_backward # noqa: F401 from liger_kernel.ops.layer_norm import layer_norm_forward # noqa: F401 from liger_kernel.ops.llama4_rope import LigerLlama4RopeFunction # noqa: F401 -from liger_kernel.ops.megatron_fused_linear_cross_entropy import ( # noqa: F401 - LigerMegatronFusedLinearCrossEntropyFunction as LigerMegatronFusedLinearCrossEntropyFunction, -) -from liger_kernel.ops.megatron_fused_linear_cross_entropy import ( # noqa: F401 - liger_megatron_fused_linear_cross_entropy as liger_megatron_fused_linear_cross_entropy, +from liger_kernel.ops.megatron_fused_linear_cross_entropy import ( + LigerMegatronFusedLinearCrossEntropyFunction, # noqa: F401 ) +from liger_kernel.ops.megatron_fused_linear_cross_entropy import liger_megatron_fused_linear_cross_entropy # noqa: F401 from liger_kernel.ops.mhc import LigerMHCCoeffsFunction # noqa: F401 from liger_kernel.ops.mhc import LigerMHCPostResFunction # noqa: F401 from liger_kernel.ops.mhc import LigerMHCPreFunction # noqa: F401 diff --git a/src/liger_kernel/ops/cutedsl/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/cutedsl/ops/megatron_fused_linear_cross_entropy.py index c3a7a9026..152b74efd 100644 --- a/src/liger_kernel/ops/cutedsl/ops/megatron_fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/cutedsl/ops/megatron_fused_linear_cross_entropy.py @@ -2,8 +2,8 @@ The SM100 path uses a persistent CuTe DSL GEMM for the local vocabulary projection, Triton for vocabulary-parallel cross entropy, and NCCL for -tensor-parallel collectives. Shifted exponentials overwrite the projection -buffer and are reused by backward. +tensor-parallel collectives. Backward converts the saved projection buffer to +dlogits in-place. """ from __future__ import annotations @@ -18,6 +18,8 @@ from liger_kernel.ops.cutedsl.ops._sm100_gemm import K_ALIGNMENT from liger_kernel.ops.cutedsl.ops._sm100_gemm import run_epilogue_gemm +from liger_kernel.ops.megatron_fused_linear_cross_entropy import _ce_backward_from_logits +from liger_kernel.ops.megatron_fused_linear_cross_entropy import _ce_forward_stats from liger_kernel.ops.megatron_fused_linear_cross_entropy import _tp_rank_and_world from liger_kernel.ops.megatron_fused_linear_cross_entropy import ( liger_megatron_fused_linear_cross_entropy as default_megatron_fused_linear_cross_entropy, @@ -64,32 +66,23 @@ def _cutedsl_projection( def _materialized_backward(ctx, grad_output: torch.Tensor): - from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps - from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_backward_kernel - - hidden, weight, exp_buffer, sum_exp, target = ctx.saved_tensors + hidden, weight, logits, logits_max, sum_exp, target = ctx.saved_tensors grad_output_1d = grad_output.contiguous().reshape(-1).float() - num_warps = _get_num_warps(ctx.ce_block_size) - liger_vocab_parallel_ce_backward_kernel[(hidden.shape[0],)]( - EXP_ptr=exp_buffer, - EXP_stride=exp_buffer.stride(0), - sum_exp_ptr=sum_exp, - Y_ptr=target, - grad_out_ptr=grad_output_1d, - vocab_start=ctx.vocab_start, - n_cols=weight.shape[0], - ignore_index=ctx.ignore_index, - alpha_eff=0.0, - eps_eff=0.0, - HAS_LABEL_SMOOTHING=False, - BLOCK_SIZE=ctx.ce_block_size, - num_warps=num_warps, + _ce_backward_from_logits( + logits, + logits_max, + sum_exp, + target, + grad_output_1d, + ctx.vocab_start, + ctx.ignore_index, + ctx.ce_block_size, ) if _SUPPORTS_OUT_DTYPE: - grad_hidden = torch.mm(exp_buffer, weight, out_dtype=torch.float32) + grad_hidden = torch.mm(logits, weight, out_dtype=torch.float32) else: - grad_hidden = exp_buffer.float() @ weight.float() + grad_hidden = logits.float() @ weight.float() reduce_work = ( dist.all_reduce( grad_hidden, @@ -100,8 +93,8 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): if ctx.tp_world > 1 else None ) - grad_weight = exp_buffer.t() @ hidden - grad_bias = exp_buffer.sum(dim=0, dtype=torch.float32).to(ctx.bias_dtype) if ctx.has_bias else None + grad_weight = logits.t() @ hidden + grad_bias = logits.sum(dim=0, dtype=torch.float32).to(ctx.bias_dtype) if ctx.has_bias else None if reduce_work is not None: reduce_work.wait() @@ -168,37 +161,25 @@ def forward( if tp_world > 1: dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group) - from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps from liger_kernel.ops.vocab_parallel_cross_entropy import _select_block_size - from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_forward_kernel - exp_buffer = logits - stats = torch.empty((2, hidden_2d.shape[0]), device=hidden.device, dtype=torch.float32) - predicted_logit = stats[0] - sum_exp = stats[1] ce_block_size = _select_block_size(vocab_local) - num_warps = _get_num_warps(ce_block_size) - liger_vocab_parallel_ce_forward_kernel[(hidden_2d.shape[0],)]( - X_ptr=logits, - X_stride=logits.stride(0), - EXP_ptr=exp_buffer, - EXP_stride=exp_buffer.stride(0), - logits_max_ptr=logits_max, - Y_ptr=flat_target, - pred_ptr=predicted_logit, - sum_exp_ptr=sum_exp, - vocab_start=vocab_start, - n_cols=vocab_local, - ignore_index=ignore_index, - BLOCK_SIZE=ce_block_size, - num_warps=num_warps, + stats = _ce_forward_stats( + logits, + logits_max, + flat_target, + vocab_start, + ignore_index, + ce_block_size, ) + predicted_logit = stats[0] + sum_exp = stats[1] if tp_world > 1: dist.all_reduce(stats, op=dist.ReduceOp.SUM, group=tp_group) loss = torch.log(sum_exp) - predicted_logit loss = torch.where(valid, loss, torch.zeros_like(loss)) - ctx.save_for_backward(hidden_2d, weight_2d, exp_buffer, sum_exp, flat_target) + ctx.save_for_backward(hidden_2d, weight_2d, logits, logits_max, sum_exp, flat_target) ctx.has_bias = bias is not None ctx.bias_dtype = bias.dtype if bias is not None else None ctx.tp_group = tp_group diff --git a/src/liger_kernel/ops/cutile/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/cutile/ops/megatron_fused_linear_cross_entropy.py index 2f7284a11..eec7cb2a2 100644 --- a/src/liger_kernel/ops/cutile/ops/megatron_fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/cutile/ops/megatron_fused_linear_cross_entropy.py @@ -184,12 +184,6 @@ def _vocab_parallel_ce_forward_kernel( running_sum + ct.sum(exponentials, 0, keepdims=False), dtype=ct.float32, ) - ct.scatter( - logits, - (row, columns), - ct.astype(exponentials, logits.dtype), - check_bounds=True, - ) ct.scatter(predicted_logit, row, predicted) ct.scatter(sum_exp, row, ct.sum(sum_exp_tile, 0, keepdims=False)) @@ -197,7 +191,8 @@ def _vocab_parallel_ce_forward_kernel( @ct.kernel(occupancy=4) def _vocab_parallel_ce_backward_kernel( - exp_buffer, + logits, + logits_max, sum_exp, target, grad_output, @@ -213,29 +208,31 @@ def _vocab_parallel_ce_backward_kernel( if y_global == ignore_index: for chunk in range(num_chunks): columns = ct.arange(BLOCK_SIZE, dtype=ct.int32) + chunk * BLOCK_SIZE - zeros = ct.full((BLOCK_SIZE,), 0.0, dtype=exp_buffer.dtype) - ct.scatter(exp_buffer, (row, columns), zeros, check_bounds=True) + zeros = ct.full((BLOCK_SIZE,), 0.0, dtype=logits.dtype) + ct.scatter(logits, (row, columns), zeros, check_bounds=True) return target_off_rank = (y_global < vocab_start) or (y_global >= vocab_start + n_cols) y_local = ct.astype(y_global - vocab_start, ct.int32) + maximum = ct.astype(ct.load(logits_max, row, shape=()), ct.float32) global_sum = ct.astype(ct.load(sum_exp, row, shape=()), ct.float32) upstream = ct.astype(ct.load(grad_output, row, shape=()), ct.float32) for chunk in range(num_chunks): columns = ct.arange(BLOCK_SIZE, dtype=ct.int32) + chunk * BLOCK_SIZE - exponentials = ct.astype( - ct.gather(exp_buffer, (row, columns), check_bounds=True, padding_value=0.0), + values = ct.astype( + ct.gather(logits, (row, columns), check_bounds=True, padding_value=-math.inf), ct.float32, ) + exponentials = ct.exp2((values - maximum) * LOG2E, flush_to_zero=True) gradient = exponentials / global_sum if not target_off_rank: gradient = ct.where(columns == y_local, gradient - 1.0, gradient) gradient = gradient * upstream ct.scatter( - exp_buffer, + logits, (row, columns), - ct.astype(gradient, exp_buffer.dtype), + ct.astype(gradient, logits.dtype), check_bounds=True, ) @@ -371,7 +368,7 @@ def _cutile_ce_forward( target: torch.Tensor, vocab_start: int, ignore_index: int, -) -> tuple[torch.Tensor, torch.Tensor]: +) -> torch.Tensor: rows, vocab_local = logits.shape stats = torch.empty((2, rows), device=logits.device, dtype=torch.float32) predicted_logit = stats[0] @@ -393,29 +390,31 @@ def _cutile_ce_forward( int(block_size), ), ) - return logits, stats + return stats def _cutile_ce_backward( - exp_buffer: torch.Tensor, + logits: torch.Tensor, + logits_max: torch.Tensor, sum_exp: torch.Tensor, target: torch.Tensor, grad_output: torch.Tensor, vocab_start: int, ignore_index: int, ) -> None: - block_size = min(2048, _select_row_block_size(exp_buffer.shape[1])) + block_size = min(2048, _select_row_block_size(logits.shape[1])) ct.launch( torch.cuda.current_stream(), - (exp_buffer.shape[0], 1, 1), + (logits.shape[0], 1, 1), _vocab_parallel_ce_backward_kernel, ( - exp_buffer, + logits, + logits_max, sum_exp, target, grad_output, int(vocab_start), - int(exp_buffer.shape[1]), + int(logits.shape[1]), int(ignore_index), int(block_size), ), @@ -451,10 +450,11 @@ def _cutile_column_sum(input: torch.Tensor, output_dtype: torch.dtype) -> torch. def _materialized_backward(ctx, grad_output: torch.Tensor): - hidden, weight, exp_buffer, sum_exp, target = ctx.saved_tensors + hidden, weight, logits, logits_max, sum_exp, target = ctx.saved_tensors grad_output_1d = grad_output.contiguous().reshape(-1).float() _cutile_ce_backward( - exp_buffer, + logits, + logits_max, sum_exp, target, grad_output_1d, @@ -463,10 +463,10 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): ) grad_hidden = _cutile_matmul( - exp_buffer, + logits, weight, operation="dx", - output_dtype=torch.float32 if exp_buffer.shape[0] <= 1024 else None, + output_dtype=torch.float32 if logits.shape[0] <= 1024 else None, ) reduce_work = ( dist.all_reduce( @@ -478,8 +478,8 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): if ctx.tp_world > 1 else None ) - grad_weight = _cutile_matmul(exp_buffer.t(), hidden, operation="dw") - grad_bias = _cutile_column_sum(exp_buffer, ctx.bias_dtype) if ctx.has_bias else None + grad_weight = _cutile_matmul(logits.t(), hidden, operation="dw") + grad_bias = _cutile_column_sum(logits, ctx.bias_dtype) if ctx.has_bias else None if reduce_work is not None: reduce_work.wait() @@ -547,7 +547,7 @@ def forward( if tp_world > 1: dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group) - exp_buffer, stats = _cutile_ce_forward( + stats = _cutile_ce_forward( logits, logits_max, flat_target, @@ -561,7 +561,7 @@ def forward( loss = _cutile_loss(sum_exp, predicted_logit, flat_target, ignore_index) - ctx.save_for_backward(hidden_2d, weight_2d, exp_buffer, sum_exp, flat_target) + ctx.save_for_backward(hidden_2d, weight_2d, logits, logits_max, sum_exp, flat_target) ctx.has_bias = bias is not None ctx.bias_dtype = bias.dtype if bias is not None else None ctx.tp_group = tp_group diff --git a/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py index 5434874bc..f9a9d29e9 100644 --- a/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py @@ -2,9 +2,9 @@ Each tensor-parallel rank owns a contiguous vocabulary shard. Forward performs one Triton projection GEMM, computes globally normalized cross entropy, and -saves shifted exponentials in the projection dtype. Backward converts that -buffer to dlogits in-place before Triton dX and dW GEMMs. Tensor-parallel -collectives remain NCCL/RCCL calls between architecture-independent kernels. +saves the local logits in the projection dtype. Backward converts that buffer +to dlogits in-place before Triton dX and dW GEMMs. Tensor-parallel collectives +remain NCCL/RCCL calls between architecture-independent kernels. """ from __future__ import annotations @@ -274,6 +274,80 @@ def _row_max_kernel( tl.store(output_ptr + row, row_max) +@triton.jit +def _ce_forward_stats_kernel( + logits_ptr, + logits_stride, + logits_max_ptr, + target_ptr, + predicted_logit_ptr, + sum_exp_ptr, + vocab_start, + n_cols, + ignore_index, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + row_ptr = logits_ptr + row * logits_stride + target = tl.load(target_ptr + row) + maximum = tl.load(logits_max_ptr + row).to(tl.float32) + target_off_rank = (target < vocab_start) | (target >= vocab_start + n_cols) + + if target == ignore_index or target_off_rank: + predicted_logit = 0.0 + else: + predicted_logit = tl.load(row_ptr + target - vocab_start).to(tl.float32) - maximum + + sum_exp = 0.0 + for start in range(0, n_cols, BLOCK_SIZE): + offsets = start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_cols + logits = tl.load(row_ptr + offsets, mask=mask, other=-float("inf")).to(tl.float32) + sum_exp += tl.sum(tl.exp(logits - maximum)) + + tl.store(predicted_logit_ptr + row, predicted_logit) + tl.store(sum_exp_ptr + row, sum_exp) + + +@triton.jit +def _ce_backward_from_logits_kernel( + logits_ptr, + logits_stride, + logits_max_ptr, + sum_exp_ptr, + target_ptr, + grad_output_ptr, + vocab_start, + n_cols, + ignore_index, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + row_ptr = logits_ptr + row * logits_stride + target = tl.load(target_ptr + row) + + if target == ignore_index: + for start in range(0, n_cols, BLOCK_SIZE): + offsets = start + tl.arange(0, BLOCK_SIZE) + tl.store(row_ptr + offsets, 0.0, mask=offsets < n_cols) + return + + maximum = tl.load(logits_max_ptr + row).to(tl.float32) + sum_exp = tl.load(sum_exp_ptr + row).to(tl.float32) + grad_output = tl.load(grad_output_ptr + row).to(tl.float32) + target_off_rank = (target < vocab_start) | (target >= vocab_start + n_cols) + target_local = target - vocab_start + + for start in range(0, n_cols, BLOCK_SIZE): + offsets = start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_cols + logits = tl.load(row_ptr + offsets, mask=mask, other=-float("inf")).to(tl.float32) + gradient = tl.exp(logits - maximum) / sum_exp + if not target_off_rank: + gradient = tl.where(offsets == target_local, gradient - 1.0, gradient) + tl.store(row_ptr + offsets, gradient * grad_output, mask=mask) + + @triton.jit def _loss_kernel( sum_exp_ptr, @@ -387,6 +461,60 @@ def _triton_row_max(input: torch.Tensor, block_size: int) -> torch.Tensor: return output +def _ce_forward_stats( + logits: torch.Tensor, + logits_max: torch.Tensor, + target: torch.Tensor, + vocab_start: int, + ignore_index: int, + block_size: int, +) -> torch.Tensor: + from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps + + stats = torch.empty((2, logits.shape[0]), device=logits.device, dtype=torch.float32) + _ce_forward_stats_kernel[(logits.shape[0],)]( + logits, + logits.stride(0), + logits_max, + target, + stats[0], + stats[1], + vocab_start, + logits.shape[1], + ignore_index, + BLOCK_SIZE=block_size, + num_warps=_get_num_warps(block_size), + ) + return stats + + +def _ce_backward_from_logits( + logits: torch.Tensor, + logits_max: torch.Tensor, + sum_exp: torch.Tensor, + target: torch.Tensor, + grad_output: torch.Tensor, + vocab_start: int, + ignore_index: int, + block_size: int, +) -> None: + from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps + + _ce_backward_from_logits_kernel[(logits.shape[0],)]( + logits, + logits.stride(0), + logits_max, + sum_exp, + target, + grad_output, + vocab_start, + logits.shape[1], + ignore_index, + BLOCK_SIZE=block_size, + num_warps=_get_num_warps(block_size), + ) + + def _triton_loss( sum_exp: torch.Tensor, predicted_logit: torch.Tensor, @@ -435,30 +563,21 @@ def _tp_rank_and_world(tp_group) -> tuple[int, int]: def _materialized_backward(ctx, grad_output: torch.Tensor): - """Convert saved CE state to dlogits and form projection gradients.""" - from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps - from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_backward_kernel - - hidden, weight, exp_buf, sum_exp_global, target = ctx.saved_tensors + """Convert saved logits to dlogits and form projection gradients.""" + hidden, weight, logits, logits_max, sum_exp_global, target = ctx.saved_tensors grad_out = grad_output.contiguous().reshape(-1).float() - num_warps = _get_num_warps(ctx.ce_block_size) - liger_vocab_parallel_ce_backward_kernel[(hidden.shape[0],)]( - EXP_ptr=exp_buf, - EXP_stride=exp_buf.stride(0), - sum_exp_ptr=sum_exp_global, - Y_ptr=target, - grad_out_ptr=grad_out, - vocab_start=ctx.vocab_start, - n_cols=weight.shape[0], - ignore_index=ctx.ignore_index, - alpha_eff=0.0, - eps_eff=0.0, - HAS_LABEL_SMOOTHING=False, - BLOCK_SIZE=ctx.ce_block_size, - num_warps=num_warps, + _ce_backward_from_logits( + logits, + logits_max, + sum_exp_global, + target, + grad_out, + ctx.vocab_start, + ctx.ignore_index, + ctx.ce_block_size, ) - grad_hidden = _triton_dx_matmul(exp_buf, weight) + grad_hidden = _triton_dx_matmul(logits, weight) reduce_work = ( dist.all_reduce( grad_hidden, @@ -469,8 +588,8 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): if ctx.tp_world > 1 else None ) - grad_weight = _triton_matmul(exp_buf.t(), hidden) - grad_bias = _triton_column_sum(exp_buf).to(ctx.bias_dtype) if ctx.has_bias else None + grad_weight = _triton_matmul(logits.t(), hidden) + grad_bias = _triton_column_sum(logits).to(ctx.bias_dtype) if ctx.has_bias else None if reduce_work is not None: reduce_work.wait() @@ -533,9 +652,7 @@ def forward( weight_2d = weight.contiguous() bias_1d = bias.contiguous() if bias is not None else None - from liger_kernel.ops.vocab_parallel_cross_entropy import _get_num_warps from liger_kernel.ops.vocab_parallel_cross_entropy import _select_block_size - from liger_kernel.ops.vocab_parallel_cross_entropy import liger_vocab_parallel_ce_forward_kernel logits = _triton_matmul(hidden_2d, weight_2d.t(), bias=bias_1d) ce_block_size = _select_block_size(vocab_local) @@ -543,32 +660,22 @@ def forward( if tp_world > 1: dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group) - exp_buf = logits - stats = torch.empty((2, hidden_2d.shape[0]), device=hidden.device, dtype=torch.float32) + stats = _ce_forward_stats( + logits, + logits_max, + flat_target, + vocab_start, + ignore_index, + ce_block_size, + ) predicted_logit = stats[0] sum_exp = stats[1] - num_warps = _get_num_warps(ce_block_size) - liger_vocab_parallel_ce_forward_kernel[(hidden_2d.shape[0],)]( - X_ptr=logits, - X_stride=logits.stride(0), - EXP_ptr=exp_buf, - EXP_stride=exp_buf.stride(0), - logits_max_ptr=logits_max, - Y_ptr=flat_target, - pred_ptr=predicted_logit, - sum_exp_ptr=sum_exp, - vocab_start=vocab_start, - n_cols=vocab_local, - ignore_index=ignore_index, - BLOCK_SIZE=ce_block_size, - num_warps=num_warps, - ) if tp_world > 1: dist.all_reduce(stats, op=dist.ReduceOp.SUM, group=tp_group) loss = _triton_loss(sum_exp, predicted_logit, flat_target, ignore_index) - ctx.save_for_backward(hidden_2d, weight_2d, exp_buf, sum_exp, flat_target) + ctx.save_for_backward(hidden_2d, weight_2d, logits, logits_max, sum_exp, flat_target) ctx.has_bias = bias is not None ctx.bias_dtype = bias.dtype if bias is not None else None ctx.tp_group = tp_group diff --git a/test/megatron/test_cutedsl_fused_linear_cross_entropy.py b/test/megatron/test_cutedsl_fused_linear_cross_entropy.py index af4b57175..40ff34f5a 100644 --- a/test/megatron/test_cutedsl_fused_linear_cross_entropy.py +++ b/test/megatron/test_cutedsl_fused_linear_cross_entropy.py @@ -82,7 +82,7 @@ def test_cutedsl_megatron_flce_backend_dispatch(): "-c", ( "from liger_kernel.megatron.fused_linear_cross_entropy import " - "liger_megatron_fused_linear_cross_entropy as fn; print(fn.__module__)" + "LigerMegatronFusedLinearCrossEntropyFunction as fn; print(fn.__module__)" ), ], check=True, diff --git a/test/megatron/test_cutile_fused_linear_cross_entropy.py b/test/megatron/test_cutile_fused_linear_cross_entropy.py index ac32c5287..853c9a3e1 100644 --- a/test/megatron/test_cutile_fused_linear_cross_entropy.py +++ b/test/megatron/test_cutile_fused_linear_cross_entropy.py @@ -75,7 +75,7 @@ def test_cutile_megatron_flce_backend_dispatch(): "-c", ( "from liger_kernel.megatron.fused_linear_cross_entropy import " - "liger_megatron_fused_linear_cross_entropy as fn; print(fn.__module__)" + "LigerMegatronFusedLinearCrossEntropyFunction as fn; print(fn.__module__)" ), ], check=True, From 4ec705fcc2d6b917c31e04907b865bbc665b5c7c Mon Sep 17 00:00:00 2001 From: Justin Hu Date: Tue, 11 Aug 2026 18:57:24 +0000 Subject: [PATCH 5/7] fix: align Megatron FLCE with upstream conventions Reuse shared validation, require the native CuTe path to run only on exact SM100 hardware, communicate dX in the projection dtype, preserve the monkey-patch TP limitation in the docs, and make benchmark memory collection discard stale autograd graphs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...ark_megatron_fused_linear_cross_entropy.py | 12 ++-- docs/High-Level-APIs.md | 18 ++++-- .../megatron_fused_linear_cross_entropy.py | 52 +++------------- .../megatron_fused_linear_cross_entropy.py | 28 ++------- .../megatron_fused_linear_cross_entropy.py | 61 +++++++++++-------- .../test_fused_linear_cross_entropy.py | 10 +++ 6 files changed, 79 insertions(+), 102 deletions(-) diff --git a/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py b/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py index bc6e9d29e..b4e846e8f 100644 --- a/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py +++ b/benchmark/scripts/benchmark_megatron_fused_linear_cross_entropy.py @@ -29,6 +29,7 @@ from __future__ import annotations import argparse +import gc import os import tempfile @@ -245,9 +246,12 @@ def full_step(): full_step() torch.cuda.synchronize() + state.clear_grads() + gc.collect() samples = [] for _ in range(_MEMORY_SAMPLES): state.clear_grads() + gc.collect() torch.cuda.reset_peak_memory_stats() full_step() torch.cuda.synchronize() @@ -326,9 +330,7 @@ def _check_correctness( dist.broadcast(upstream, src=0) outputs = {} correctness_providers = ["megatron-compatible"] - correctness_providers.extend( - provider for provider in ("liger", "liger-cutile", "liger-cutedsl") if provider in providers - ) + correctness_providers.extend(provider for provider in providers if provider != "megatron-compatible") for provider in correctness_providers: state = _make_state(provider, hidden, weight, bias, target, tp_group, tp_size) loss = state.forward() @@ -517,11 +519,11 @@ def main(): raise RuntimeError("provider 'liger-cutedsl' requested, but nvidia-cutlass-dsl is not installed.") if "liger-cutedsl" in providers: unsupported_devices = [ - index for index in range(args.tp_size) if torch.cuda.get_device_capability(index)[0] < 10 + index for index in range(args.tp_size) if torch.cuda.get_device_capability(index) != (10, 0) ] if unsupported_devices: raise RuntimeError( - "provider 'liger-cutedsl' requires SM100 or newer; " + "provider 'liger-cutedsl' requires SM100 (compute capability 10.0); " f"unsupported CUDA device indices: {unsupported_devices}." ) if "liger-cutile" in providers and not _CUTILE_AVAILABLE: diff --git a/docs/High-Level-APIs.md b/docs/High-Level-APIs.md index d0fb89d71..1f21910df 100644 --- a/docs/High-Level-APIs.md +++ b/docs/High-Level-APIs.md @@ -107,12 +107,18 @@ shards. | Megatron-LM | `liger_kernel.megatron.apply_liger_kernel_to_megatron` | RMSNorm, CrossEntropyLoss | | Megatron-LM | `liger_kernel.megatron.LigerMegatronFusedLinearCrossEntropy` | Fused output projection + CrossEntropyLoss | -`LigerMegatronFusedLinearCrossEntropy` accepts replicated hidden states, -the calling rank's contiguous `[V_local, H]` output-weight shard, and global -target indices. It supports TP1 and TP>1 through the supplied process group. -The default implementation uses Triton local kernels and NCCL. Set -`LIGER_KERNEL_IMPL=cutile` or `cutedsl` before importing Liger to select a -CuTile local path or the SM100 CuTe DSL persistent projection. +**Scope**: The monkey patch supports `tensor_model_parallel_size=1` only for +cross-entropy. Vocab-parallel cross-entropy patching (TP>1) remains follow-up +work; the patch raises a `RuntimeError` at patch time or call time if TP>1 is +detected. + +The separately wired `LigerMegatronFusedLinearCrossEntropy` module accepts +replicated hidden states, the calling rank's contiguous `[V_local, H]` +output-weight shard, and global target indices. It supports TP1 and TP>1 +through the supplied process group. The default implementation uses Triton +local kernels and NCCL. Set `LIGER_KERNEL_IMPL=cutile` or `cutedsl` before +importing Liger to select a CuTile local path or the SM100 CuTe DSL persistent +projection. **Usage**: diff --git a/src/liger_kernel/ops/cutedsl/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/cutedsl/ops/megatron_fused_linear_cross_entropy.py index 152b74efd..1f0f8fdee 100644 --- a/src/liger_kernel/ops/cutedsl/ops/megatron_fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/cutedsl/ops/megatron_fused_linear_cross_entropy.py @@ -8,8 +8,6 @@ from __future__ import annotations -import operator - import cutlass import cutlass.cute as cute import torch @@ -21,12 +19,7 @@ from liger_kernel.ops.megatron_fused_linear_cross_entropy import _ce_backward_from_logits from liger_kernel.ops.megatron_fused_linear_cross_entropy import _ce_forward_stats from liger_kernel.ops.megatron_fused_linear_cross_entropy import _tp_rank_and_world -from liger_kernel.ops.megatron_fused_linear_cross_entropy import ( - liger_megatron_fused_linear_cross_entropy as default_megatron_fused_linear_cross_entropy, -) -from liger_kernel.ops.utils import compare_version - -_SUPPORTS_OUT_DTYPE = compare_version("torch", operator.ge, "2.8.0") +from liger_kernel.ops.megatron_fused_linear_cross_entropy import _validate_megatron_flce_inputs @cute.jit @@ -42,7 +35,7 @@ def _native_cutedsl_supported(hidden: torch.Tensor, weight: torch.Tensor) -> boo if weight.device != hidden.device or weight.dtype != hidden.dtype: return False try: - return torch.cuda.get_device_capability(hidden.device)[0] >= 10 + return torch.cuda.get_device_capability(hidden.device) == (10, 0) except (AssertionError, RuntimeError): return False @@ -79,10 +72,7 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): ctx.ce_block_size, ) - if _SUPPORTS_OUT_DTYPE: - grad_hidden = torch.mm(logits, weight, out_dtype=torch.float32) - else: - grad_hidden = logits.float() @ weight.float() + grad_hidden = torch.mm(logits, weight) reduce_work = ( dist.all_reduce( grad_hidden, @@ -98,7 +88,7 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): if reduce_work is not None: reduce_work.wait() - grad_hidden = grad_hidden.to(ctx.hidden_dtype).reshape(ctx.original_hidden_shape) + grad_hidden = grad_hidden.reshape(ctx.original_hidden_shape) return grad_hidden, grad_weight, grad_bias @@ -115,32 +105,15 @@ def forward( tp_group, ignore_index: int, ) -> torch.Tensor: - if hidden.ndim < 2: - raise ValueError(f"hidden must have at least 2 dimensions, got shape {tuple(hidden.shape)}.") - if weight.ndim != 2: - raise ValueError(f"weight must be 2-D [V_local, H], got shape {tuple(weight.shape)}.") - if tuple(target.shape) != tuple(hidden.shape[:-1]): - raise ValueError( - f"target shape must equal hidden.shape[:-1]; got target={tuple(target.shape)}, " - f"hidden={tuple(hidden.shape)}." - ) - if hidden.shape[-1] != weight.shape[1]: - raise ValueError(f"hidden size mismatch: hidden has H={hidden.shape[-1]}, weight has H={weight.shape[1]}.") - if hidden.dtype != weight.dtype: - raise TypeError(f"hidden and weight must have the same dtype, got {hidden.dtype} and {weight.dtype}.") - if hidden.device != weight.device or hidden.device != target.device: - raise ValueError("hidden, weight, and target must be on the same device.") - if bias is not None: - if bias.ndim != 1 or bias.shape[0] != weight.shape[0]: - raise ValueError(f"bias must have shape ({weight.shape[0]},), got {tuple(bias.shape)}.") - if bias.device != hidden.device or bias.dtype != hidden.dtype: - raise TypeError("bias must have the same device and dtype as hidden.") + _validate_megatron_flce_inputs(hidden, weight, target, bias) + if not _native_cutedsl_supported(hidden, weight): + raise RuntimeError("CuTe DSL Megatron FLCE requires an SM100 GPU and float16 or bfloat16 inputs.") tp_rank, tp_world = _tp_rank_and_world(tp_group) vocab_local = weight.shape[0] vocab_global = vocab_local * tp_world vocab_start = tp_rank * vocab_local - flat_target = target.reshape(-1).to(torch.int64).contiguous() + flat_target = target.reshape(-1).contiguous() valid = flat_target != ignore_index invalid = valid & ((flat_target < 0) | (flat_target >= vocab_global)) valid_targets = ~torch.any(invalid) @@ -206,15 +179,6 @@ def liger_megatron_fused_linear_cross_entropy( ignore_index: int = -100, ) -> torch.Tensor: """Compute Megatron FLCE with a CuTe DSL projection and NCCL TP collectives.""" - if not _native_cutedsl_supported(hidden, weight): - return default_megatron_fused_linear_cross_entropy( - hidden, - weight, - target, - bias=bias, - tp_group=tp_group, - ignore_index=ignore_index, - ) return LigerMegatronFusedLinearCrossEntropyFunction.apply( hidden, weight, diff --git a/src/liger_kernel/ops/cutile/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/cutile/ops/megatron_fused_linear_cross_entropy.py index eec7cb2a2..a31f26a82 100644 --- a/src/liger_kernel/ops/cutile/ops/megatron_fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/cutile/ops/megatron_fused_linear_cross_entropy.py @@ -10,6 +10,7 @@ from liger_kernel.ops.cutile.ops.utils import _next_power_of_2 from liger_kernel.ops.megatron_fused_linear_cross_entropy import _tp_rank_and_world +from liger_kernel.ops.megatron_fused_linear_cross_entropy import _validate_megatron_flce_inputs ConstBool = ct.Constant[bool] ConstInt = ct.Constant[int] @@ -467,7 +468,7 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): weight, operation="dx", output_dtype=torch.float32 if logits.shape[0] <= 1024 else None, - ) + ).to(ctx.hidden_dtype) reduce_work = ( dist.all_reduce( grad_hidden, @@ -483,7 +484,7 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): if reduce_work is not None: reduce_work.wait() - grad_hidden = grad_hidden.to(ctx.hidden_dtype).reshape(ctx.original_hidden_shape) + grad_hidden = grad_hidden.reshape(ctx.original_hidden_shape) return grad_hidden, grad_weight, grad_bias @@ -500,26 +501,7 @@ def forward( tp_group, ignore_index: int, ) -> torch.Tensor: - if hidden.ndim < 2: - raise ValueError(f"hidden must have at least 2 dimensions, got shape {tuple(hidden.shape)}.") - if weight.ndim != 2: - raise ValueError(f"weight must be 2-D [V_local, H], got shape {tuple(weight.shape)}.") - if tuple(target.shape) != tuple(hidden.shape[:-1]): - raise ValueError( - f"target shape must equal hidden.shape[:-1]; got target={tuple(target.shape)}, " - f"hidden={tuple(hidden.shape)}." - ) - if hidden.shape[-1] != weight.shape[1]: - raise ValueError(f"hidden size mismatch: hidden has H={hidden.shape[-1]}, weight has H={weight.shape[1]}.") - if hidden.dtype != weight.dtype: - raise TypeError(f"hidden and weight must have the same dtype, got {hidden.dtype} and {weight.dtype}.") - if hidden.device != weight.device or hidden.device != target.device: - raise ValueError("hidden, weight, and target must be on the same device.") - if bias is not None: - if bias.ndim != 1 or bias.shape[0] != weight.shape[0]: - raise ValueError(f"bias must have shape ({weight.shape[0]},), got {tuple(bias.shape)}.") - if bias.device != hidden.device or bias.dtype != hidden.dtype: - raise TypeError("bias must have the same device and dtype as hidden.") + _validate_megatron_flce_inputs(hidden, weight, target, bias) if hidden.device.type != "cuda" or hidden.dtype not in (torch.bfloat16, torch.float16): raise RuntimeError("CuTile Megatron FLCE requires a CUDA GPU and float16 or bfloat16 inputs.") @@ -528,7 +510,7 @@ def forward( vocab_global = vocab_local * tp_world vocab_start = tp_rank * vocab_local - flat_target = target.reshape(-1).to(torch.int64).contiguous() + flat_target = target.reshape(-1).contiguous() valid = flat_target != ignore_index invalid = valid & ((flat_target < 0) | (flat_target >= vocab_global)) valid_targets = ~torch.any(invalid) diff --git a/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py b/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py index f9a9d29e9..e4c9ed9d6 100644 --- a/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/megatron_fused_linear_cross_entropy.py @@ -422,7 +422,7 @@ def _triton_dx_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: m, k = a.shape n = b.shape[1] if m > 1024 or k < 4096: - return _triton_matmul(a, b, output_dtype=torch.float32) + return _triton_matmul(a, b) output = torch.zeros((m, n), device=a.device, dtype=torch.float32) grid = lambda meta: ( @@ -562,6 +562,38 @@ def _tp_rank_and_world(tp_group) -> tuple[int, int]: return dist.get_rank(tp_group), world +def _validate_megatron_flce_inputs( + hidden: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + bias: torch.Tensor | None, +) -> None: + if hidden.ndim < 2: + raise ValueError(f"hidden must have at least 2 dimensions, got shape {tuple(hidden.shape)}.") + if weight.ndim != 2: + raise ValueError(f"weight must be 2-D [V_local, H], got shape {tuple(weight.shape)}.") + if tuple(target.shape) != tuple(hidden.shape[:-1]): + raise ValueError( + f"target shape must equal hidden.shape[:-1]; got target={tuple(target.shape)}, " + f"hidden={tuple(hidden.shape)}." + ) + if target.dtype != torch.long: + raise TypeError(f"target must have dtype torch.long, got {target.dtype}.") + if hidden.shape[-1] != weight.shape[1]: + raise ValueError(f"hidden size mismatch: hidden has H={hidden.shape[-1]}, weight has H={weight.shape[1]}.") + if target.numel() == 0 or hidden.shape[-1] == 0 or weight.shape[0] == 0: + raise ValueError("hidden, weight, and target dimensions must be non-empty.") + if hidden.dtype != weight.dtype: + raise TypeError(f"hidden and weight must have the same dtype, got {hidden.dtype} and {weight.dtype}.") + if hidden.device != weight.device or hidden.device != target.device: + raise ValueError("hidden, weight, and target must be on the same device.") + if bias is not None: + if bias.ndim != 1 or bias.shape[0] != weight.shape[0]: + raise ValueError(f"bias must have shape ({weight.shape[0]},), got {tuple(bias.shape)}.") + if bias.device != hidden.device or bias.dtype != hidden.dtype: + raise TypeError("bias must have the same device and dtype as hidden.") + + def _materialized_backward(ctx, grad_output: torch.Tensor): """Convert saved logits to dlogits and form projection gradients.""" hidden, weight, logits, logits_max, sum_exp_global, target = ctx.saved_tensors @@ -577,7 +609,7 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): ctx.ce_block_size, ) - grad_hidden = _triton_dx_matmul(logits, weight) + grad_hidden = _triton_dx_matmul(logits, weight).to(ctx.hidden_dtype) reduce_work = ( dist.all_reduce( grad_hidden, @@ -593,7 +625,7 @@ def _materialized_backward(ctx, grad_output: torch.Tensor): if reduce_work is not None: reduce_work.wait() - grad_hidden = grad_hidden.to(ctx.hidden_dtype).reshape(ctx.original_hidden_shape) + grad_hidden = grad_hidden.reshape(ctx.original_hidden_shape) return grad_hidden, grad_weight, grad_bias @@ -610,26 +642,7 @@ def forward( tp_group, ignore_index: int, ) -> torch.Tensor: - if hidden.ndim < 2: - raise ValueError(f"hidden must have at least 2 dimensions, got shape {tuple(hidden.shape)}.") - if weight.ndim != 2: - raise ValueError(f"weight must be 2-D [V_local, H], got shape {tuple(weight.shape)}.") - if tuple(target.shape) != tuple(hidden.shape[:-1]): - raise ValueError( - f"target shape must equal hidden.shape[:-1]; got target={tuple(target.shape)}, " - f"hidden={tuple(hidden.shape)}." - ) - if hidden.shape[-1] != weight.shape[1]: - raise ValueError(f"hidden size mismatch: hidden has H={hidden.shape[-1]}, weight has H={weight.shape[1]}.") - if hidden.dtype != weight.dtype: - raise TypeError(f"hidden and weight must have the same dtype, got {hidden.dtype} and {weight.dtype}.") - if hidden.device != weight.device or hidden.device != target.device: - raise ValueError("hidden, weight, and target must be on the same device.") - if bias is not None: - if bias.ndim != 1 or bias.shape[0] != weight.shape[0]: - raise ValueError(f"bias must have shape ({weight.shape[0]},), got {tuple(bias.shape)}.") - if bias.device != hidden.device or bias.dtype != hidden.dtype: - raise TypeError("bias must have the same device and dtype as hidden.") + _validate_megatron_flce_inputs(hidden, weight, target, bias) if hidden.device.type != "cuda" or hidden.dtype not in (torch.bfloat16, torch.float16): raise RuntimeError("Megatron FLCE requires a CUDA GPU and float16 or bfloat16 inputs.") @@ -638,7 +651,7 @@ def forward( vocab_global = vocab_local * tp_world vocab_start = tp_rank * vocab_local - flat_target = target.reshape(-1).to(torch.int64).contiguous() + flat_target = target.reshape(-1).contiguous() valid = flat_target != ignore_index invalid = valid & ((flat_target < 0) | (flat_target >= vocab_global)) valid_targets = ~torch.any(invalid) diff --git a/test/megatron/test_fused_linear_cross_entropy.py b/test/megatron/test_fused_linear_cross_entropy.py index 0bc07d69b..0b4fe2c7f 100644 --- a/test/megatron/test_fused_linear_cross_entropy.py +++ b/test/megatron/test_fused_linear_cross_entropy.py @@ -109,6 +109,16 @@ def test_megatron_flce_rejects_cpu_inputs(): liger_megatron_fused_linear_cross_entropy(hidden, weight, target) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Megatron FLCE requires CUDA") +def test_megatron_flce_rejects_non_long_targets(): + hidden = torch.randn(2, 3, 8, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(16, 8, device="cuda", dtype=torch.bfloat16) + target = torch.zeros(2, 3, device="cuda", dtype=torch.int32) + + with pytest.raises(TypeError, match="target must have dtype torch.long"): + liger_megatron_fused_linear_cross_entropy(hidden, weight, target) + + def _tp_worker(rank, world_size, file_name, dtype): dist.init_process_group( backend="nccl", From 36bfa0ec3ebc9ffcc729671305140d12c29ef5df Mon Sep 17 00:00:00 2001 From: Justin Hu Date: Tue, 11 Aug 2026 20:14:07 +0000 Subject: [PATCH 6/7] fix: wire Megatron loss patches into GPT Repair Megatron's loaded cross-entropy bindings and add an opt-in GPT output-processor hook for fused linear cross entropy. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/High-Level-APIs.md | 56 ++++--- examples/megatron/README.md | 5 + src/liger_kernel/megatron/__init__.py | 6 +- .../megatron/fused_linear_cross_entropy.py | 57 +++++++ src/liger_kernel/megatron/monkey_patch.py | 97 +++++++++++- .../test_fused_linear_cross_entropy.py | 68 ++++++++ test/megatron/test_monkey_patch.py | 147 +++++++++++++++++- 7 files changed, 406 insertions(+), 30 deletions(-) diff --git a/docs/High-Level-APIs.md b/docs/High-Level-APIs.md index 1f21910df..83b0eb454 100644 --- a/docs/High-Level-APIs.md +++ b/docs/High-Level-APIs.md @@ -98,38 +98,45 @@ You can also use the Patching APIs to use the kernels for a specific model archi Liger also exposes a patch for the [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) training framework, replacing Megatron's native RMSNorm and both vocab-parallel -cross-entropy paths (fused and unfused) with Liger kernels. Liger also exposes -a hidden-state-to-loss FLCE module for contiguous tensor-parallel vocabulary -shards. +cross-entropy paths (fused and unfused) with Liger kernels. An additional +opt-in patch fuses the GPT output projection with vocab-parallel +cross-entropy, avoiding a separately returned logits tensor between the +projection and loss. | **Framework** | **API** | **Supported Operations** | |---------------|--------------------------------------------------------|--------------------------| -| Megatron-LM | `liger_kernel.megatron.apply_liger_kernel_to_megatron` | RMSNorm, CrossEntropyLoss | +| Megatron-LM | `liger_kernel.megatron.apply_liger_kernel_to_megatron` | RMSNorm, CrossEntropyLoss, fused output projection + CrossEntropyLoss | | Megatron-LM | `liger_kernel.megatron.LigerMegatronFusedLinearCrossEntropy` | Fused output projection + CrossEntropyLoss | -**Scope**: The monkey patch supports `tensor_model_parallel_size=1` only for -cross-entropy. Vocab-parallel cross-entropy patching (TP>1) remains follow-up -work; the patch raises a `RuntimeError` at patch time or call time if TP>1 is -detected. - -The separately wired `LigerMegatronFusedLinearCrossEntropy` module accepts -replicated hidden states, the calling rank's contiguous `[V_local, H]` -output-weight shard, and global target indices. It supports TP1 and TP>1 -through the supplied process group. The default implementation uses Triton -local kernels and NCCL. Set `LIGER_KERNEL_IMPL=cutile` or `cutedsl` before -importing Liger to select a CuTile local path or the SM100 CuTe DSL persistent -projection. +**Scope**: Both cross-entropy patches and FLCE support TP1 and TP>1. The FLCE +patch requires Megatron-Core 0.18 or newer and a native +`ColumnParallelLinear` output layer in BF16 or FP16. It intentionally rejects +gathered logits, sequence parallelism, gradient-accumulation fusion, deferred +embedding wgrad, disabled output dgrad reduction, MTP, MuP output scaling, and +separately returned output bias. A caller-supplied `output_processor` takes +precedence and is left unchanged. + +The patch installs FLCE through Megatron's `GPTModel._postprocess` +`output_processor` hook. For custom model integrations, the +`LigerMegatronFusedLinearCrossEntropy` module accepts replicated hidden +states, the calling rank's contiguous `[V_local, H]` output-weight shard, and +global target indices. The default implementation uses Triton local kernels +and NCCL. Set `LIGER_KERNEL_IMPL=cutile` or `cutedsl` before importing Liger +to select a CuTile local path or the SM100 CuTe DSL persistent projection. **Usage**: ```python from liger_kernel.megatron import apply_liger_kernel_to_megatron -# Call before Megatron's forward pass reaches compute_language_model_loss. -# Defaults match Megatron's native CE behavior; no CE-specific config needed. -apply_liger_kernel_to_megatron(rms_norm=True, cross_entropy=True) +# Call before constructing the GPT model. FLCE is separately opt-in. +apply_liger_kernel_to_megatron( + rms_norm=True, + cross_entropy=True, + fused_linear_cross_entropy=True, +) -# Or wire the hidden-state-to-loss operation into a vocab-sharded output layer. +# Or wire FLCE into a custom vocab-sharded output layer directly. from liger_kernel.megatron import LigerMegatronFusedLinearCrossEntropy loss_fn = LigerMegatronFusedLinearCrossEntropy(ignore_index=-100) @@ -139,7 +146,8 @@ loss = loss_fn(hidden, local_output_weight, global_targets, tp_group=tp_group) Both the fused (`config.cross_entropy_loss_fusion=True`, `cross_entropy_fusion_impl='native'`) and unfused (`config.cross_entropy_loss_fusion=False`) CE paths are patched in a single -call, so Megatron picks up Liger regardless of which path your config selects. +call. When `fused_linear_cross_entropy=True`, labeled standard GPT forwards +bypass both materialized-logit paths and use FLCE directly. For training setups that need explicit kernel configuration (custom `ignore_index`, `label_smoothing`, etc.), instantiate @@ -152,6 +160,12 @@ For training setups that need explicit kernel configuration (custom show_docstring: true show_signature: true +::: liger_kernel.megatron.liger_megatron_fused_linear_cross_entropy_output_processor + options: + extra: + show_docstring: true + show_signature: true + ::: liger_kernel.megatron.apply_liger_kernel_to_megatron options: extra: diff --git a/examples/megatron/README.md b/examples/megatron/README.md index abf06e950..e62208f4c 100644 --- a/examples/megatron/README.md +++ b/examples/megatron/README.md @@ -11,6 +11,7 @@ so you can see which slots picked up Liger. |---|---|---|---|---| | RMSNorm | `rms_norm=True` (on by default) | `LocalSpecProvider.layer_norm`, `transformer_block.LayerNormImpl` | `LigerMegatronRMSNorm` | every norm slot, incl. block-level `final_layernorm` | | Cross-entropy | `cross_entropy=True` (opt-in) | `fused_cross_entropy.fused_vocab_parallel_cross_entropy`, `tensor_parallel.cross_entropy.vocab_parallel_cross_entropy` | `LigerMegatronCrossEntropy` | none — `GPTModel` subclass overriding `compute_language_model_loss` | +| Fused output projection + cross-entropy | `fused_linear_cross_entropy=True` (opt-in) | `GPTModel._postprocess` | `LigerMegatronFusedLinearCrossEntropy` | custom GPT `output_processor` | | SwiGLU | `swiglu=True` (opt-in) | `fusions.fused_bias_swiglu.SwiGLUFunction` | `LigerMegatronSwiGLU` | the `mlp` module slot — an `MLP` subclass | Notes that apply to the table: @@ -23,6 +24,10 @@ Notes that apply to the table: Megatron. - Cross-entropy and SwiGLU are wired through subclasses (no dedicated spec slot). +- FLCE requires Megatron-Core 0.18 or newer, BF16/FP16, the native + `ColumnParallelLinear` output layer, and no sequence parallelism or + gradient-accumulation fusion. The example script keeps its small FP32 + configuration and therefore does not enable this flag. ## Prerequisites diff --git a/src/liger_kernel/megatron/__init__.py b/src/liger_kernel/megatron/__init__.py index 4cb498d13..e54fbb2ab 100644 --- a/src/liger_kernel/megatron/__init__.py +++ b/src/liger_kernel/megatron/__init__.py @@ -12,10 +12,12 @@ offload, and neither touches the bias or MoE-routed variants. LigerMegatronFusedLinearCrossEntropy — hidden-state-to-loss fused output projection for tensor-parallel vocabulary shards. + liger_megatron_fused_linear_cross_entropy_output_processor — adapts FLCE + to Megatron's GPT output-processor hook. apply_liger_kernel_to_megatron — patches Megatron-Core so existing training scripts pick up Liger kernels with one line. Currently supports RMSNorm (via BackendSpecProvider), both the fused and unfused - vocab-parallel cross-entropy paths, and SwiGLU. + vocab-parallel cross-entropy paths, opt-in GPT FLCE, and SwiGLU. The general-purpose ``LigerVocabParallelCrossEntropy`` Module lives under ``liger_kernel.transformers`` alongside the other nn.Module wrappers; the @@ -26,6 +28,7 @@ from liger_kernel.megatron.cross_entropy import LigerMegatronCrossEntropy from liger_kernel.megatron.fused_linear_cross_entropy import LigerMegatronFusedLinearCrossEntropy +from liger_kernel.megatron.fused_linear_cross_entropy import liger_megatron_fused_linear_cross_entropy_output_processor from liger_kernel.megatron.monkey_patch import apply_liger_kernel_to_megatron from liger_kernel.megatron.rms_norm import LigerMegatronRMSNorm from liger_kernel.megatron.swiglu import LigerMegatronSwiGLU @@ -36,4 +39,5 @@ "LigerMegatronRMSNorm", "LigerMegatronSwiGLU", "apply_liger_kernel_to_megatron", + "liger_megatron_fused_linear_cross_entropy_output_processor", ] diff --git a/src/liger_kernel/megatron/fused_linear_cross_entropy.py b/src/liger_kernel/megatron/fused_linear_cross_entropy.py index eb87c627c..0e5507344 100644 --- a/src/liger_kernel/megatron/fused_linear_cross_entropy.py +++ b/src/liger_kernel/megatron/fused_linear_cross_entropy.py @@ -41,3 +41,60 @@ def forward( def extra_repr(self) -> str: return f"ignore_index={self.ignore_index}" + + +def liger_megatron_fused_linear_cross_entropy_output_processor( + *, + hidden_states: torch.Tensor, + output_layer, + output_weight: torch.Tensor | None, + labels: torch.Tensor, + runtime_gather_output: bool | None, + config, + **_, +) -> torch.Tensor: + """Megatron ``GPTModel`` output processor for the native TP output layer.""" + unsupported = [] + if type(output_layer).__name__ != "ColumnParallelLinear": + unsupported.append("the output layer is not Megatron's native ColumnParallelLinear") + if getattr(output_layer, "sequence_parallel", False): + unsupported.append("sequence_parallel=True") + if getattr(output_layer, "gradient_accumulation_fusion", False): + unsupported.append("gradient_accumulation_fusion=True") + if getattr(output_layer, "disable_grad_reduce", False): + unsupported.append("output-layer dgrad reduction is disabled") + if getattr(output_layer, "explicit_expert_comm", False): + unsupported.append("the output layer uses explicit expert communication") + if getattr(output_layer, "skip_bias_add", False): + unsupported.append("the output layer returns bias separately") + if getattr(config, "defer_embedding_wgrad_compute", False): + unsupported.append("defer_embedding_wgrad_compute=True") + if getattr(config, "mtp_num_layers", None): + unsupported.append("MTP is enabled") + if getattr(config, "use_mup", False): + unsupported.append("MuP output scaling is enabled") + + gather_output = ( + getattr(output_layer, "gather_output", False) if runtime_gather_output is None else runtime_gather_output + ) + if gather_output: + unsupported.append("the output layer gathers TP logits") + if unsupported: + raise RuntimeError( + "Liger Megatron FLCE does not support this GPT output configuration: " + "; ".join(unsupported) + ) + + weight = output_weight if output_weight is not None else getattr(output_layer, "weight", None) + if weight is None: + raise RuntimeError("Liger Megatron FLCE requires an output weight tensor.") + + labels_sb = labels.transpose(0, 1).contiguous() + loss_sb = LigerMegatronFusedLinearCrossEntropyFunction.apply( + hidden_states, + weight, + labels_sb, + getattr(output_layer, "bias", None), + getattr(output_layer, "tp_group", None), + -100, + ) + return loss_sb.transpose(0, 1).contiguous() diff --git a/src/liger_kernel/megatron/monkey_patch.py b/src/liger_kernel/megatron/monkey_patch.py index e97567ca3..7b4db807c 100644 --- a/src/liger_kernel/megatron/monkey_patch.py +++ b/src/liger_kernel/megatron/monkey_patch.py @@ -2,16 +2,27 @@ from __future__ import annotations +import functools +import inspect import logging +import sys logger = logging.getLogger(__name__) _PATCH_MARKER = "__liger_patched__" +def _replace_loaded_binding(module_name: str, symbol_name: str, original, replacement) -> None: + """Update a known by-name import without clobbering third-party patches.""" + module = sys.modules.get(module_name) + if module is not None and getattr(module, symbol_name, None) is original: + setattr(module, symbol_name, replacement) + + def apply_liger_kernel_to_megatron( rms_norm: bool = True, cross_entropy: bool = False, + fused_linear_cross_entropy: bool = False, swiglu: bool = False, ) -> None: """Patch Megatron-Core to use Liger Triton kernels. @@ -39,6 +50,11 @@ def apply_liger_kernel_to_megatron( wrapper additionally honors a runtime ``label_smoothing`` argument, matching native's ``(logits, target, label_smoothing=0.0, tp_group=None)``. + fused_linear_cross_entropy: When ``True`` inject Liger's fused local + output projection and vocab-parallel cross-entropy through + ``GPTModel._postprocess`` for supported native + ``ColumnParallelLinear`` training configurations. This is opt-in + and independent of ``cross_entropy``. swiglu: When ``True`` replace ``megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction`` with Liger's Triton SiLU-multiply kernel, covering the dense ``MLP`` and the MoE @@ -68,6 +84,8 @@ def apply_liger_kernel_to_megatron( if cross_entropy: _patch_fused_vocab_parallel_cross_entropy() _patch_vocab_parallel_cross_entropy() + if fused_linear_cross_entropy: + _patch_gpt_fused_linear_cross_entropy() if swiglu: _patch_swiglu_function() @@ -183,7 +201,14 @@ def _patch_fused_vocab_parallel_cross_entropy() -> None: ) if getattr(fused_ce.fused_vocab_parallel_cross_entropy, _PATCH_MARKER, False): - return # already patched + replacement = fused_ce.fused_vocab_parallel_cross_entropy + _replace_loaded_binding( + "megatron.core.models.common.language_module.language_module", + "fused_vocab_parallel_cross_entropy", + replacement.__wrapped__, + replacement, + ) + return original = fused_ce.fused_vocab_parallel_cross_entropy @@ -197,11 +222,15 @@ def liger_fused_vocab_parallel_cross_entropy(vocab_parallel_logits, target, tp_g setattr(liger_fused_vocab_parallel_cross_entropy, _PATCH_MARKER, True) setattr(liger_fused_vocab_parallel_cross_entropy, "__wrapped__", original) fused_ce.fused_vocab_parallel_cross_entropy = liger_fused_vocab_parallel_cross_entropy - - logger.info( - "Patched megatron.core.fusions.fused_cross_entropy.fused_vocab_parallel_cross_entropy with Liger cross-entropy." + _replace_loaded_binding( + "megatron.core.models.common.language_module.language_module", + "fused_vocab_parallel_cross_entropy", + original, + liger_fused_vocab_parallel_cross_entropy, ) + logger.info("Patched Megatron's fused vocab-parallel cross-entropy definition and loaded LanguageModule binding.") + def _patch_vocab_parallel_cross_entropy() -> None: """Replace ``megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy``. @@ -229,7 +258,14 @@ def _patch_vocab_parallel_cross_entropy() -> None: ) if getattr(unfused_ce.vocab_parallel_cross_entropy, _PATCH_MARKER, False): - return # already patched + replacement = unfused_ce.vocab_parallel_cross_entropy + _replace_loaded_binding( + "megatron.core.tensor_parallel", + "vocab_parallel_cross_entropy", + replacement.__wrapped__, + replacement, + ) + return original = unfused_ce.vocab_parallel_cross_entropy @@ -260,11 +296,58 @@ def liger_vocab_parallel_cross_entropy( setattr(liger_vocab_parallel_cross_entropy, _PATCH_MARKER, True) setattr(liger_vocab_parallel_cross_entropy, "__wrapped__", original) unfused_ce.vocab_parallel_cross_entropy = liger_vocab_parallel_cross_entropy + _replace_loaded_binding( + "megatron.core.tensor_parallel", + "vocab_parallel_cross_entropy", + original, + liger_vocab_parallel_cross_entropy, + ) - logger.info( - "Patched megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy with Liger cross-entropy." + logger.info("Patched Megatron's unfused vocab-parallel cross-entropy definition and tensor_parallel export.") + + +def _patch_gpt_fused_linear_cross_entropy() -> None: + """Inject Liger FLCE through Megatron's GPT output-processor hook.""" + try: + import megatron.core.models.gpt.gpt_model as gpt_model + except ImportError as exc: + raise ImportError( + "apply_liger_kernel_to_megatron(fused_linear_cross_entropy=True) requires " + "megatron-core with megatron.core.models.gpt.gpt_model.GPTModel." + ) from exc + + if not hasattr(gpt_model, "GPTModel") or not hasattr(gpt_model.GPTModel, "_postprocess"): + raise ImportError( + "megatron.core.models.gpt.gpt_model.GPTModel._postprocess was not found. " + "The symbol path may have changed in your Megatron-Core version." + ) + + current = gpt_model.GPTModel._postprocess + if getattr(current, _PATCH_MARKER, False): + return + signature = inspect.signature(current) + if not {"labels", "output_processor"} <= signature.parameters.keys(): + raise ImportError( + "Megatron GPTModel._postprocess does not expose the labels and output_processor hooks " + "required by Liger FLCE. Upgrade to Megatron-Core 0.18 or newer." + ) + + from liger_kernel.megatron.fused_linear_cross_entropy import ( + liger_megatron_fused_linear_cross_entropy_output_processor, ) + @functools.wraps(current) + def liger_postprocess(self, *args, **kwargs): + bound = signature.bind(self, *args, **kwargs) + bound.apply_defaults() + if bound.arguments["labels"] is not None and bound.arguments["output_processor"] is None: + bound.arguments["output_processor"] = liger_megatron_fused_linear_cross_entropy_output_processor + return current(*bound.args, **bound.kwargs) + + setattr(liger_postprocess, _PATCH_MARKER, True) + gpt_model.GPTModel._postprocess = liger_postprocess + logger.info("Patched Megatron GPTModel._postprocess to inject Liger fused linear cross-entropy.") + def _patch_swiglu_function() -> None: """Replace ``megatron.core.fusions.fused_bias_swiglu.SwiGLUFunction`` with Liger. diff --git a/test/megatron/test_fused_linear_cross_entropy.py b/test/megatron/test_fused_linear_cross_entropy.py index 0b4fe2c7f..af3aafa43 100644 --- a/test/megatron/test_fused_linear_cross_entropy.py +++ b/test/megatron/test_fused_linear_cross_entropy.py @@ -2,6 +2,8 @@ import tempfile +from types import SimpleNamespace + import pytest import torch import torch.distributed as dist @@ -9,6 +11,7 @@ import torch.nn.functional as F from liger_kernel.megatron import LigerMegatronFusedLinearCrossEntropy +from liger_kernel.megatron import liger_megatron_fused_linear_cross_entropy_output_processor from liger_kernel.ops.megatron_fused_linear_cross_entropy import liger_megatron_fused_linear_cross_entropy @@ -76,6 +79,71 @@ def test_megatron_flce_module_contract(): assert "ignore_index=-1" in repr(module) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Megatron FLCE requires CUDA") +def test_megatron_flce_output_processor_matches_materialized_path(): + class ColumnParallelLinear: + def __init__(self, weight, bias): + self.weight = weight + self.bias = bias + self.tp_group = None + self.gather_output = False + self.sequence_parallel = False + self.gradient_accumulation_fusion = False + self.disable_grad_reduce = False + self.explicit_expert_comm = False + self.skip_bias_add = False + + torch.manual_seed(43) + hidden = torch.randn(3, 2, 8, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(16, 8, device="cuda", dtype=torch.bfloat16) * 0.02 + bias = torch.randn(16, device="cuda", dtype=torch.bfloat16) * 0.02 + labels = torch.randint(16, (2, 3), device="cuda") + output_layer = ColumnParallelLinear(weight, bias) + config = SimpleNamespace( + defer_embedding_wgrad_compute=False, + mtp_num_layers=None, + use_mup=False, + ) + + actual = liger_megatron_fused_linear_cross_entropy_output_processor( + hidden_states=hidden, + output_layer=output_layer, + output_weight=None, + labels=labels, + runtime_gather_output=None, + config=config, + ) + reference = _reference_loss(hidden, weight, labels.t().contiguous(), bias).t().contiguous() + torch.testing.assert_close(actual, reference, atol=5e-3, rtol=5e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Megatron FLCE requires CUDA") +def test_megatron_flce_output_processor_rejects_gathered_logits(): + output_layer = SimpleNamespace( + gather_output=True, + sequence_parallel=False, + gradient_accumulation_fusion=False, + disable_grad_reduce=False, + explicit_expert_comm=False, + skip_bias_add=False, + ) + config = SimpleNamespace( + defer_embedding_wgrad_compute=False, + mtp_num_layers=None, + use_mup=False, + ) + + with pytest.raises(RuntimeError, match="not Megatron's native ColumnParallelLinear.*gathers TP logits"): + liger_megatron_fused_linear_cross_entropy_output_processor( + hidden_states=torch.empty(1, 1, 8, device="cuda", dtype=torch.bfloat16), + output_layer=output_layer, + output_weight=None, + labels=torch.zeros(1, 1, device="cuda", dtype=torch.long), + runtime_gather_output=None, + config=config, + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Megatron FLCE requires CUDA") @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) def test_megatron_flce_triton_split_k_dx_matches_pytorch(dtype): diff --git a/test/megatron/test_monkey_patch.py b/test/megatron/test_monkey_patch.py index 634a4abd5..ad88a078a 100644 --- a/test/megatron/test_monkey_patch.py +++ b/test/megatron/test_monkey_patch.py @@ -8,7 +8,7 @@ - patching is idempotent (calling apply twice doesn't stack wrappers) - the patch is a no-op when the kernel flag is False - missing megatron-core / missing symbol path raise helpful ``ImportError``\\s -- kernel-specific dispatch contracts (e.g. CE TP>1 raises; RMSNorm only displaces the +- kernel-specific dispatch contracts (e.g. FLCE configuration guards; RMSNorm only displaces the ``WrappedTorchNorm`` fallback, not TE / Apex) - end-to-end: the patched symbol invoked with real tensors produces correct output @@ -66,6 +66,10 @@ def _install_fake_megatron_ce( tensor_parallel = types.ModuleType("megatron.core.tensor_parallel") unfused_ce = types.ModuleType("megatron.core.tensor_parallel.cross_entropy") parallel_state = types.ModuleType("megatron.core.parallel_state") + models = sys.modules.get("megatron.core.models") or types.ModuleType("megatron.core.models") + common = types.ModuleType("megatron.core.models.common") + language_module_package = types.ModuleType("megatron.core.models.common.language_module") + language_module = types.ModuleType("megatron.core.models.common.language_module.language_module") if with_fused_symbol: @@ -73,6 +77,7 @@ def original_fused_vocab_parallel_cross_entropy(vocab_parallel_logits, target, t raise AssertionError("original megatron fused kernel called — patch failed") fused_ce.fused_vocab_parallel_cross_entropy = original_fused_vocab_parallel_cross_entropy + language_module.fused_vocab_parallel_cross_entropy = original_fused_vocab_parallel_cross_entropy if with_unfused_symbol: @@ -85,6 +90,7 @@ def original_vocab_parallel_cross_entropy( raise AssertionError("original megatron unfused kernel called — patch failed") unfused_ce.vocab_parallel_cross_entropy = original_vocab_parallel_cross_entropy + tensor_parallel.vocab_parallel_cross_entropy = original_vocab_parallel_cross_entropy parallel_state.get_tensor_model_parallel_world_size = lambda: tp_size @@ -93,12 +99,21 @@ def original_vocab_parallel_cross_entropy( sys.modules["megatron.core.tensor_parallel"] = tensor_parallel sys.modules["megatron.core.tensor_parallel.cross_entropy"] = unfused_ce sys.modules["megatron.core.parallel_state"] = parallel_state + sys.modules["megatron.core.models"] = models + sys.modules["megatron.core.models.common"] = common + sys.modules["megatron.core.models.common.language_module"] = language_module_package + sys.modules["megatron.core.models.common.language_module.language_module"] = language_module megatron_core.fusions = fusions megatron_core.tensor_parallel = tensor_parallel megatron_core.parallel_state = parallel_state + megatron_core.models = models fusions.fused_cross_entropy = fused_ce tensor_parallel.cross_entropy = unfused_ce + models.common = common + common.language_module = language_module_package + language_module_package.language_module = language_module + language_module.tensor_parallel = tensor_parallel return fused_ce, unfused_ce @@ -195,6 +210,11 @@ def _uninstall_fake_megatron(): # CE side "megatron.core.parallel_state", "megatron.core.fusions.fused_cross_entropy", + "megatron.core.models.common.language_module.language_module", + "megatron.core.models.common.language_module", + "megatron.core.models.common", + "megatron.core.models.gpt.gpt_model", + "megatron.core.models.gpt", # RMSNorm side "megatron.core.models.backends", "megatron.core.models", @@ -216,6 +236,43 @@ def _uninstall_fake_megatron(): sys.modules.pop(mod, None) +def _install_fake_megatron_gpt(with_output_processor: bool = True): + _, megatron_core = _ensure_megatron_roots() + models = sys.modules.get("megatron.core.models") or types.ModuleType("megatron.core.models") + gpt = types.ModuleType("megatron.core.models.gpt") + gpt_model = types.ModuleType("megatron.core.models.gpt.gpt_model") + + if with_output_processor: + + class GPTModel: + def _postprocess(self, labels=None, output_processor=None): + return output_processor + + else: + + class GPTModel: + def _postprocess(self, labels=None): + return labels + + gpt_model.GPTModel = GPTModel + sys.modules["megatron.core.models"] = models + sys.modules["megatron.core.models.gpt"] = gpt + sys.modules["megatron.core.models.gpt.gpt_model"] = gpt_model + megatron_core.models = models + models.gpt = gpt + gpt.gpt_model = gpt_model + return gpt_model + + +@pytest.fixture +def fake_megatron_gpt(): + gpt_model = _install_fake_megatron_gpt() + try: + yield gpt_model + finally: + _uninstall_fake_megatron() + + @pytest.fixture def fake_megatron_ce(): fused_ce, unfused_ce = _install_fake_megatron_ce(tp_size=1) @@ -399,6 +456,30 @@ def test_patch_replaces_both_fused_and_unfused_symbols_in_one_call(fake_megatron assert unfused_ce.vocab_parallel_cross_entropy.__name__ == "liger_vocab_parallel_cross_entropy" +def test_patch_rebinds_loaded_language_module_fused_consumer(fake_megatron_ce): + fused_ce, _ = fake_megatron_ce + language_module = sys.modules["megatron.core.models.common.language_module.language_module"] + original = language_module.fused_vocab_parallel_cross_entropy + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=True) + + assert language_module.fused_vocab_parallel_cross_entropy is fused_ce.fused_vocab_parallel_cross_entropy + assert language_module.fused_vocab_parallel_cross_entropy is not original + + +def test_patch_rebinds_loaded_tensor_parallel_export(fake_megatron_ce): + _, unfused_ce = fake_megatron_ce + tensor_parallel = sys.modules["megatron.core.tensor_parallel"] + original = tensor_parallel.vocab_parallel_cross_entropy + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=True) + + assert tensor_parallel.vocab_parallel_cross_entropy is unfused_ce.vocab_parallel_cross_entropy + assert tensor_parallel.vocab_parallel_cross_entropy is not original + + def test_patch_with_cross_entropy_false_leaves_ce_symbols_untouched(fake_megatron_ce): """Default ``cross_entropy=False`` must not touch the CE symbols even if the call runs.""" fused_ce, unfused_ce = fake_megatron_ce @@ -423,10 +504,19 @@ def test_patch_is_idempotent_for_both_symbols(fake_megatron_ce): fused_first = fused_ce.fused_vocab_parallel_cross_entropy unfused_first = unfused_ce.vocab_parallel_cross_entropy + # Model a consumer restoring its import-time binding after the definitions + # were patched; the idempotent path must repair these stale references. + language_module = sys.modules["megatron.core.models.common.language_module.language_module"] + tensor_parallel = sys.modules["megatron.core.tensor_parallel"] + language_module.fused_vocab_parallel_cross_entropy = fused_first.__wrapped__ + tensor_parallel.vocab_parallel_cross_entropy = unfused_first.__wrapped__ + apply_liger_kernel_to_megatron(rms_norm=False, cross_entropy=True) # Same identity → no stacked wrapping. assert fused_ce.fused_vocab_parallel_cross_entropy is fused_first assert unfused_ce.vocab_parallel_cross_entropy is unfused_first + assert language_module.fused_vocab_parallel_cross_entropy is fused_first + assert tensor_parallel.vocab_parallel_cross_entropy is unfused_first # __wrapped__ still references the original Megatron symbol, not the first Liger wrapper. assert fused_first.__wrapped__.__name__ == "original_fused_vocab_parallel_cross_entropy" assert unfused_first.__wrapped__.__name__ == "original_vocab_parallel_cross_entropy" @@ -920,9 +1010,11 @@ def test_import_from_root(): accidental __init__.py removals so the docs' import snippets keep working.""" try: from liger_kernel.megatron import LigerMegatronCrossEntropy # noqa: F401 + from liger_kernel.megatron import LigerMegatronFusedLinearCrossEntropy # noqa: F401 from liger_kernel.megatron import LigerMegatronRMSNorm # noqa: F401 from liger_kernel.megatron import LigerMegatronSwiGLU # noqa: F401 from liger_kernel.megatron import apply_liger_kernel_to_megatron # noqa: F401 + from liger_kernel.megatron import liger_megatron_fused_linear_cross_entropy_output_processor # noqa: F401 except Exception: pytest.fail("Importing public Megatron symbols from liger_kernel.megatron failed.") @@ -943,6 +1035,59 @@ def test_public_apply_function_has_no_ce_specific_kwargs(): ) +def test_flce_patch_injects_output_processor_for_labeled_gpt_calls(fake_megatron_gpt): + from liger_kernel.megatron import apply_liger_kernel_to_megatron + from liger_kernel.megatron import liger_megatron_fused_linear_cross_entropy_output_processor + + original = fake_megatron_gpt.GPTModel._postprocess + apply_liger_kernel_to_megatron(rms_norm=False, fused_linear_cross_entropy=True) + patched = fake_megatron_gpt.GPTModel._postprocess + + assert patched is not original + assert patched.__wrapped__ is original + assert patched(object(), labels=object()) is liger_megatron_fused_linear_cross_entropy_output_processor + assert patched(object(), labels=None) is None + + +def test_flce_patch_is_opt_in(fake_megatron_gpt): + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + original = fake_megatron_gpt.GPTModel._postprocess + apply_liger_kernel_to_megatron(rms_norm=False) + + assert fake_megatron_gpt.GPTModel._postprocess is original + + +def test_flce_patch_preserves_custom_output_processor(fake_megatron_gpt): + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + custom = object() + apply_liger_kernel_to_megatron(rms_norm=False, fused_linear_cross_entropy=True) + + assert fake_megatron_gpt.GPTModel()._postprocess(labels=object(), output_processor=custom) is custom + + +def test_flce_patch_is_idempotent(fake_megatron_gpt): + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + apply_liger_kernel_to_megatron(rms_norm=False, fused_linear_cross_entropy=True) + first = fake_megatron_gpt.GPTModel._postprocess + apply_liger_kernel_to_megatron(rms_norm=False, fused_linear_cross_entropy=True) + + assert fake_megatron_gpt.GPTModel._postprocess is first + + +def test_flce_patch_requires_megatron_output_processor_hook(): + _install_fake_megatron_gpt(with_output_processor=False) + try: + from liger_kernel.megatron import apply_liger_kernel_to_megatron + + with pytest.raises(ImportError, match="Megatron-Core 0.18 or newer"): + apply_liger_kernel_to_megatron(rms_norm=False, fused_linear_cross_entropy=True) + finally: + _uninstall_fake_megatron() + + # =========================================================================== # 5. End-to-end integration through the patched CE symbols # =========================================================================== From 84835f121748d40e31e2d29a089acaf3bc4fd563 Mon Sep 17 00:00:00 2001 From: Justin Hu Date: Tue, 11 Aug 2026 22:08:56 +0000 Subject: [PATCH 7/7] docs: trim Megatron FLCE documentation Remove benchmark- and example-specific README additions and keep the high-level API documentation concise. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- benchmark/README.md | 28 +--------------------------- docs/High-Level-APIs.md | 37 ++++++++++--------------------------- examples/megatron/README.md | 5 ----- 3 files changed, 11 insertions(+), 59 deletions(-) diff --git a/benchmark/README.md b/benchmark/README.md index c8f048e50..1e33bcd2e 100644 --- a/benchmark/README.md +++ b/benchmark/README.md @@ -186,30 +186,4 @@ visualizer at it with `--data-file`: python ../benchmarks_visualizer.py \ --kernel-name cross_entropy --metric-name speed \ --data-file data/all_benchmark_data_cutedsl.csv -``` - -### Megatron fused linear cross-entropy - -The Megatron FLCE benchmark checks loss and gradient parity before measuring -the default Triton implementation and the opt-in CuTile and CuTe DSL backends: - -```bash -cd benchmark/scripts -python benchmark_megatron_fused_linear_cross_entropy.py \ - --tp-size 4 --token-counts 512 2048 --vocab-sizes 32000 128256 \ - --providers megatron-compatible liger liger-cutile - -# 7B/Llama-3-scale output layer -python benchmark_megatron_fused_linear_cross_entropy.py \ - --tp-size 4 --token-counts 16384 --hidden-size 4096 \ - --vocab-sizes 128256 \ - --providers megatron-compatible liger liger-cutile - -# SM100 with nvidia-cutlass-dsl installed -python benchmark_megatron_fused_linear_cross_entropy.py \ - --tp-size 4 --providers megatron-compatible liger-cutedsl -``` - -Targets are global vocabulary indices; each rank owns a contiguous vocabulary -shard. Results include forward, backward, full-step latency, and peak Torch -memory. +``` \ No newline at end of file diff --git a/docs/High-Level-APIs.md b/docs/High-Level-APIs.md index 83b0eb454..64029f60b 100644 --- a/docs/High-Level-APIs.md +++ b/docs/High-Level-APIs.md @@ -98,56 +98,39 @@ You can also use the Patching APIs to use the kernels for a specific model archi Liger also exposes a patch for the [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) training framework, replacing Megatron's native RMSNorm and both vocab-parallel -cross-entropy paths (fused and unfused) with Liger kernels. An additional -opt-in patch fuses the GPT output projection with vocab-parallel -cross-entropy, avoiding a separately returned logits tensor between the -projection and loss. +cross-entropy paths (fused and unfused) with Liger kernels. It can also fuse +the GPT output projection with cross-entropy. | **Framework** | **API** | **Supported Operations** | |---------------|--------------------------------------------------------|--------------------------| | Megatron-LM | `liger_kernel.megatron.apply_liger_kernel_to_megatron` | RMSNorm, CrossEntropyLoss, fused output projection + CrossEntropyLoss | | Megatron-LM | `liger_kernel.megatron.LigerMegatronFusedLinearCrossEntropy` | Fused output projection + CrossEntropyLoss | -**Scope**: Both cross-entropy patches and FLCE support TP1 and TP>1. The FLCE -patch requires Megatron-Core 0.18 or newer and a native -`ColumnParallelLinear` output layer in BF16 or FP16. It intentionally rejects -gathered logits, sequence parallelism, gradient-accumulation fusion, deferred -embedding wgrad, disabled output dgrad reduction, MTP, MuP output scaling, and -separately returned output bias. A caller-supplied `output_processor` takes -precedence and is left unchanged. - -The patch installs FLCE through Megatron's `GPTModel._postprocess` -`output_processor` hook. For custom model integrations, the -`LigerMegatronFusedLinearCrossEntropy` module accepts replicated hidden -states, the calling rank's contiguous `[V_local, H]` output-weight shard, and -global target indices. The default implementation uses Triton local kernels -and NCCL. Set `LIGER_KERNEL_IMPL=cutile` or `cutedsl` before importing Liger -to select a CuTile local path or the SM100 CuTe DSL persistent projection. +Both cross-entropy patches and FLCE support TP1 and TP>1. Automatic FLCE +patching requires Megatron-Core 0.18 or newer, BF16 or FP16, a native +`ColumnParallelLinear` output layer, and no sequence parallelism, +gradient-accumulation fusion, or gathered logits. Unsupported configurations +raise an explicit error. **Usage**: ```python from liger_kernel.megatron import apply_liger_kernel_to_megatron -# Call before constructing the GPT model. FLCE is separately opt-in. apply_liger_kernel_to_megatron( rms_norm=True, cross_entropy=True, fused_linear_cross_entropy=True, ) - -# Or wire FLCE into a custom vocab-sharded output layer directly. -from liger_kernel.megatron import LigerMegatronFusedLinearCrossEntropy - -loss_fn = LigerMegatronFusedLinearCrossEntropy(ignore_index=-100) -loss = loss_fn(hidden, local_output_weight, global_targets, tp_group=tp_group) ``` Both the fused (`config.cross_entropy_loss_fusion=True`, `cross_entropy_fusion_impl='native'`) and unfused (`config.cross_entropy_loss_fusion=False`) CE paths are patched in a single call. When `fused_linear_cross_entropy=True`, labeled standard GPT forwards -bypass both materialized-logit paths and use FLCE directly. +bypass both paths and use FLCE directly. For custom integrations, use +`LigerMegatronFusedLinearCrossEntropy` with replicated hidden states, a +contiguous local vocabulary shard, and global target indices. For training setups that need explicit kernel configuration (custom `ignore_index`, `label_smoothing`, etc.), instantiate diff --git a/examples/megatron/README.md b/examples/megatron/README.md index e62208f4c..abf06e950 100644 --- a/examples/megatron/README.md +++ b/examples/megatron/README.md @@ -11,7 +11,6 @@ so you can see which slots picked up Liger. |---|---|---|---|---| | RMSNorm | `rms_norm=True` (on by default) | `LocalSpecProvider.layer_norm`, `transformer_block.LayerNormImpl` | `LigerMegatronRMSNorm` | every norm slot, incl. block-level `final_layernorm` | | Cross-entropy | `cross_entropy=True` (opt-in) | `fused_cross_entropy.fused_vocab_parallel_cross_entropy`, `tensor_parallel.cross_entropy.vocab_parallel_cross_entropy` | `LigerMegatronCrossEntropy` | none — `GPTModel` subclass overriding `compute_language_model_loss` | -| Fused output projection + cross-entropy | `fused_linear_cross_entropy=True` (opt-in) | `GPTModel._postprocess` | `LigerMegatronFusedLinearCrossEntropy` | custom GPT `output_processor` | | SwiGLU | `swiglu=True` (opt-in) | `fusions.fused_bias_swiglu.SwiGLUFunction` | `LigerMegatronSwiGLU` | the `mlp` module slot — an `MLP` subclass | Notes that apply to the table: @@ -24,10 +23,6 @@ Notes that apply to the table: Megatron. - Cross-entropy and SwiGLU are wired through subclasses (no dedicated spec slot). -- FLCE requires Megatron-Core 0.18 or newer, BF16/FP16, the native - `ColumnParallelLinear` output layer, and no sequence parallelism or - gradient-accumulation fusion. The example script keeps its small FP32 - configuration and therefore does not enable this flag. ## Prerequisites