mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
feat(cutile): chunk FLCE with Triton sizing and cuBLAS GEMMs
Reuse Triton default token-chunk geometry through a private helper without adding a public chunk-size parameter. Bound cuTile logits/statistics workspace, use cuBLAS for all GEMMs, retain FP32 dW until upstream scaling, and preserve legacy low-level helpers. Enable SM100 alongside SM90 and cover numerical, memory, gradient-mask, and older-Torch fallback behavior. Co-authored-by: Copilot <[email protected]>
This commit is contained in:
Vaibhav Jindal
co-authored by
Copilot
Copilot
parent
95b01e9402
commit
ecd0780285
@@ -209,6 +209,28 @@ LIGER_KERNEL_IMPL=cutile python your_script.py
|
||||
|
||||
`LIGER_KERNEL_IMPL` selects an opt-in implementation registered with Liger (currently `cutile` and `cutedsl`). Selecting one on an unsupported device, or without the required dependencies installed, raises an error.
|
||||
|
||||
#### cuTile Fused Linear Cross Entropy
|
||||
|
||||
With `LIGER_KERNEL_IMPL=cutile`, `LigerFusedLinearCrossEntropyLoss` routes to a cuTile Fused Linear Cross Entropy kernel for BF16 `input`/`weight` on Hopper (SM90) and Blackwell (SM100). No API change is needed — the default route just works:
|
||||
|
||||
```bash
|
||||
LIGER_KERNEL_IMPL=cutile python your_script.py
|
||||
```
|
||||
|
||||
```python
|
||||
from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss
|
||||
|
||||
# Same call as the default backend — there is no chunk_size argument.
|
||||
loss_fn = LigerFusedLinearCrossEntropyLoss(reduction="mean")
|
||||
loss = loss_fn(lm_head_weight, hidden_states, labels)
|
||||
```
|
||||
|
||||
Implementation notes:
|
||||
|
||||
- Token chunking uses the **exact same default geometry as the Triton backend** (`inc = ceil(V / (16·H))`, `chunk = min(N, next_pow2(ceil(N / inc)))`); there is no public `chunk_size` knob.
|
||||
- All three GEMMs — projection, input-gradient (`dX`), and weight-gradient (`dW`) — run through PyTorch/cuBLAS. `dW` is accumulated across token chunks in FP32 (`mm` then `addmm`) and the combined upstream/mean scale is applied in FP32 before a single BF16 cast, so small gradients are preserved.
|
||||
- Supports BF16 `input` and `weight` with `mean`/`sum` reductions; FP32 accumulation is the default and only accumulation policy. It is **not** a full feature/precision match for the Triton kernel (e.g. no bias, class weights, z-loss, label smoothing, softcap, or token scaling).
|
||||
|
||||
### Enable CuTe DSL Backend
|
||||
|
||||
CuTe DSL is the optional, CUDA-only Python DSL shipped with NVIDIA CUTLASS (`import cutlass.cute`), targeting Hopper (SM90) and Blackwell (SM100/SM110). After installing the `cutedsl` extra, enable it explicitly:
|
||||
|
||||
@@ -2,94 +2,84 @@
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""Hopper-optimized cuTile fused linear cross entropy.
|
||||
"""cuTile fused linear cross entropy for Hopper (SM90) and Blackwell (SM100).
|
||||
|
||||
The projection and input-gradient GEMMs use PyTorch's Hopper-tuned BLAS path.
|
||||
cuTile computes partitioned CE statistics, overwrites logits with dZ, and
|
||||
computes dW with an FP32-accumulating 128x256x64 MMA schedule.
|
||||
All three GEMMs -- the projection, the input-gradient (dX), and the
|
||||
weight-gradient (dW) -- use PyTorch's tuned cuBLAS path. cuTile computes
|
||||
partitioned CE statistics, overwrites a reusable ``[min(N, C), V]`` logits
|
||||
buffer with dZ, and applies the final FP32 dW scale-and-cast. dW is accumulated
|
||||
across token chunks in FP32 by cuBLAS ``mm`` (first chunk) then ``addmm``
|
||||
(subsequent chunks), so the transient logits workspace is
|
||||
``O(min(N, C) * V)`` instead of the full ``O(N * V)`` materialization; the
|
||||
retained gradients (raw BF16 dX and the FP32 dW accumulator) add ``O(N*H + V*H)``.
|
||||
|
||||
The upstream/mean scale is applied out-of-place in backward: the FP32 dW
|
||||
accumulator is retained until backward and cast to the weight dtype exactly once,
|
||||
after multiplying by the combined ``upstream * mean_normalizer`` factor. Casting
|
||||
after (not before) that multiply preserves the numeric range of small gradients
|
||||
that would otherwise flush to zero in BF16.
|
||||
|
||||
Two low-level interfaces are exported. ``fused_linear_cross_entropy_forward`` /
|
||||
``fused_linear_cross_entropy_backward`` retain the original unchunked contract
|
||||
(the forward returns the full ``[N, V]`` dZ and the 7-argument backward runs the
|
||||
dW matmul). The token-chunked path used by the autograd Function lives in
|
||||
``chunked_fused_linear_cross_entropy_forward`` /
|
||||
``chunked_fused_linear_cross_entropy_backward``.
|
||||
"""
|
||||
|
||||
import operator
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import cuda.tile as ct
|
||||
import torch
|
||||
|
||||
from liger_kernel.ops.cutile.ops.utils import _next_power_of_2
|
||||
from liger_kernel.ops.fused_linear_cross_entropy import _get_chunk_size
|
||||
from liger_kernel.ops.utils import amp_custom_bwd
|
||||
from liger_kernel.ops.utils import amp_custom_fwd
|
||||
from liger_kernel.ops.utils import compare_version
|
||||
|
||||
ConstBool = ct.Constant[bool]
|
||||
ConstInt = ct.Constant[int]
|
||||
|
||||
LOG2E = 1.4426950408889634
|
||||
DW_TILE_M = 128
|
||||
DW_TILE_N = 128
|
||||
DW_TILE_K = 64
|
||||
SCALE_CAST_TILE_M = 128
|
||||
SCALE_CAST_TILE_N = 128
|
||||
CE_BLOCK_SIZE = 2048
|
||||
LOGITS_STATS_BLOCK_SIZE = 4096
|
||||
MAX_STATS_BLOCK_SIZE = 1024
|
||||
|
||||
# The FP32 dW accumulation uses the ``torch.mm`` / ``torch.addmm``
|
||||
# ``out_dtype=torch.float32`` overload (BF16 operands -> FP32 accumulator) added
|
||||
# in torch 2.8. On older torch the operands are upcast to FP32 first; see
|
||||
# ``_accumulate_dw_fp32``. This is a runtime capability check only -- it does not
|
||||
# raise the package floor.
|
||||
_ADDMM_SUPPORTS_OUT_DTYPE = compare_version("torch", operator.ge, "2.8.0")
|
||||
|
||||
@ct.kernel(num_worker_warps=ct.ByTarget(sm_90=8), opt_level=3)
|
||||
def _matmul_tn_kernel(
|
||||
a,
|
||||
b,
|
||||
output,
|
||||
grad_output,
|
||||
|
||||
@ct.kernel(num_worker_warps=ct.ByTarget(sm_90=8, default=8), opt_level=3)
|
||||
def _scale_cast_kernel(
|
||||
src,
|
||||
dst,
|
||||
scale,
|
||||
TILE_M: ConstInt,
|
||||
TILE_N: ConstInt,
|
||||
TILE_K: ConstInt,
|
||||
):
|
||||
"""Compute ``output = (a.T @ b) * grad_output``."""
|
||||
"""Scale an FP32 ``src`` tile by a device scalar and cast into ``dst.dtype``.
|
||||
|
||||
Used to turn the retained FP32 dW accumulator into the final weight-dtype
|
||||
gradient without allocating an O(V*H) FP32 temporary: the multiply by the
|
||||
combined ``upstream * mean_normalizer`` scalar happens in FP32 and the BF16
|
||||
cast is applied exactly once, per element, on store.
|
||||
"""
|
||||
pid_m = ct.bid(0)
|
||||
pid_n = ct.bid(1)
|
||||
a_view = a.tiled_view((TILE_K, TILE_M), padding_mode=ct.PaddingMode.ZERO)
|
||||
b_view = b.tiled_view((TILE_K, TILE_N), padding_mode=ct.PaddingMode.ZERO)
|
||||
acc = ct.zeros((TILE_M, TILE_N), dtype=ct.float32)
|
||||
|
||||
for pid_k in range(a_view.num_tiles(0)):
|
||||
a_tile = a_view.load((pid_k, pid_m), latency=10)
|
||||
b_tile = b_view.load((pid_k, pid_n), latency=10)
|
||||
acc = ct.mma(ct.transpose(a_tile), b_tile, acc)
|
||||
|
||||
scale = ct.astype(ct.gather(grad_output, (), check_bounds=False), ct.float32)
|
||||
ct.store(output, (pid_m, pid_n), ct.astype(acc * scale, output.dtype), latency=1)
|
||||
|
||||
|
||||
@ct.kernel(num_worker_warps=ct.ByTarget(sm_90=8), opt_level=3)
|
||||
def _matmul_tn_n4_kernel(
|
||||
a,
|
||||
b,
|
||||
output,
|
||||
grad_output,
|
||||
TILE_M: ConstInt,
|
||||
TILE_N: ConstInt,
|
||||
TILE_K: ConstInt,
|
||||
LOAD_LATENCY: ConstInt,
|
||||
):
|
||||
"""Compute four adjacent dW column tiles while sharing each dZ load."""
|
||||
pid_n = ct.bid(0) * 4
|
||||
pid_m = ct.bid(1)
|
||||
a_view = a.tiled_view((TILE_K, TILE_M), padding_mode=ct.PaddingMode.ZERO)
|
||||
b_view = b.tiled_view((TILE_K, TILE_N), padding_mode=ct.PaddingMode.ZERO)
|
||||
output_view = output.tiled_view((TILE_M, TILE_N), padding_mode=ct.PaddingMode.ZERO)
|
||||
acc0 = ct.zeros((TILE_M, TILE_N), dtype=ct.float32)
|
||||
acc1 = ct.zeros((TILE_M, TILE_N), dtype=ct.float32)
|
||||
acc2 = ct.zeros((TILE_M, TILE_N), dtype=ct.float32)
|
||||
acc3 = ct.zeros((TILE_M, TILE_N), dtype=ct.float32)
|
||||
|
||||
for pid_k in range(a_view.num_tiles(0)):
|
||||
a_tile = ct.transpose(a_view.load((pid_k, pid_m), latency=LOAD_LATENCY))
|
||||
acc0 = ct.mma(a_tile, b_view.load((pid_k, pid_n), latency=LOAD_LATENCY), acc0)
|
||||
acc1 = ct.mma(a_tile, b_view.load((pid_k, pid_n + 1), latency=LOAD_LATENCY), acc1)
|
||||
acc2 = ct.mma(a_tile, b_view.load((pid_k, pid_n + 2), latency=LOAD_LATENCY), acc2)
|
||||
acc3 = ct.mma(a_tile, b_view.load((pid_k, pid_n + 3), latency=LOAD_LATENCY), acc3)
|
||||
|
||||
scale = ct.astype(ct.gather(grad_output, (), check_bounds=False), ct.float32)
|
||||
output_view.store((pid_m, pid_n), ct.astype(acc0 * scale, output.dtype), latency=1)
|
||||
output_view.store((pid_m, pid_n + 1), ct.astype(acc1 * scale, output.dtype), latency=1)
|
||||
output_view.store((pid_m, pid_n + 2), ct.astype(acc2 * scale, output.dtype), latency=1)
|
||||
output_view.store((pid_m, pid_n + 3), ct.astype(acc3 * scale, output.dtype), latency=1)
|
||||
src_view = src.tiled_view((TILE_M, TILE_N), padding_mode=ct.PaddingMode.ZERO)
|
||||
dst_view = dst.tiled_view((TILE_M, TILE_N), padding_mode=ct.PaddingMode.ZERO)
|
||||
tile = ct.astype(src_view.load((pid_m, pid_n), latency=2), ct.float32)
|
||||
tile_scale = ct.astype(ct.gather(scale, (), check_bounds=False), ct.float32)
|
||||
dst_view.store((pid_m, pid_n), ct.astype(tile * tile_scale, dst.dtype), latency=1)
|
||||
|
||||
|
||||
@ct.kernel(occupancy=8)
|
||||
@@ -177,36 +167,36 @@ def _fused_cross_entropy_dz_kernel(
|
||||
ct.scatter(logits, (row, cols), ct.astype(gradient, logits.dtype), check_bounds=True)
|
||||
|
||||
|
||||
def _launch_matmul_tn(
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
grad_output: torch.Tensor,
|
||||
def _accumulate_dw_fp32(
|
||||
dweight_accum: torch.Tensor,
|
||||
dz: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
accumulate: bool,
|
||||
) -> None:
|
||||
if output.shape[1] % 256 == 0:
|
||||
_launch_cutile(
|
||||
output.device,
|
||||
(
|
||||
output.shape[1] // 256,
|
||||
(output.shape[0] + DW_TILE_M - 1) // DW_TILE_M,
|
||||
1,
|
||||
),
|
||||
_matmul_tn_n4_kernel,
|
||||
(a, b, output, grad_output, DW_TILE_M, 64, DW_TILE_K, 5),
|
||||
)
|
||||
return
|
||||
"""Accumulate ``dZ.T @ X`` into the FP32 ``dweight_accum`` buffer via cuBLAS.
|
||||
|
||||
tile_n = min(256, _next_power_of_2(output.shape[1]))
|
||||
_launch_cutile(
|
||||
output.device,
|
||||
(
|
||||
(output.shape[0] + DW_TILE_M - 1) // DW_TILE_M,
|
||||
(output.shape[1] + tile_n - 1) // tile_n,
|
||||
1,
|
||||
),
|
||||
_matmul_tn_kernel,
|
||||
(a, b, output, grad_output, DW_TILE_M, tile_n, DW_TILE_K),
|
||||
)
|
||||
``dz`` is ``[rows, V]`` and ``x`` is ``[rows, H]``; the contribution
|
||||
``dZ.T @ X`` is ``[V, H]``. The first token chunk (``accumulate`` is
|
||||
``False``) overwrites the buffer with ``torch.mm``; subsequent chunks add
|
||||
with ``torch.addmm``. Accumulation is always FP32: on torch>=2.8 the BF16
|
||||
operands drive the ``out_dtype=torch.float32`` overload directly; on older
|
||||
torch (no overload) only the bounded BF16 operands are upcast to FP32 before
|
||||
the same mm/addmm into the FP32 buffer. This single helper is shared by the
|
||||
chunked and legacy paths so their dW math is identical.
|
||||
"""
|
||||
dz_t = dz.t() # [V, rows] view
|
||||
if _ADDMM_SUPPORTS_OUT_DTYPE:
|
||||
if accumulate:
|
||||
torch.addmm(dweight_accum, dz_t, x, out_dtype=torch.float32, out=dweight_accum)
|
||||
else:
|
||||
torch.mm(dz_t, x, out_dtype=torch.float32, out=dweight_accum)
|
||||
else:
|
||||
dz_t_fp32 = dz_t.to(torch.float32)
|
||||
x_fp32 = x.to(torch.float32)
|
||||
if accumulate:
|
||||
torch.addmm(dweight_accum, dz_t_fp32, x_fp32, out=dweight_accum)
|
||||
else:
|
||||
torch.mm(dz_t_fp32, x_fp32, out=dweight_accum)
|
||||
|
||||
|
||||
def _launch_cutile(device: torch.device, grid, kernel, args) -> None:
|
||||
@@ -214,6 +204,20 @@ def _launch_cutile(device: torch.device, grid, kernel, args) -> None:
|
||||
ct.launch(torch.cuda.current_stream(device), grid, kernel, args)
|
||||
|
||||
|
||||
def _launch_scale_cast(src: torch.Tensor, dst: torch.Tensor, scale: torch.Tensor) -> None:
|
||||
"""Compute ``dst = (src * scale).astype(dst.dtype)`` without an FP32 temporary."""
|
||||
_launch_cutile(
|
||||
dst.device,
|
||||
(
|
||||
(src.shape[0] + SCALE_CAST_TILE_M - 1) // SCALE_CAST_TILE_M,
|
||||
(src.shape[1] + SCALE_CAST_TILE_N - 1) // SCALE_CAST_TILE_N,
|
||||
1,
|
||||
),
|
||||
_scale_cast_kernel,
|
||||
(src, dst, scale, SCALE_CAST_TILE_M, SCALE_CAST_TILE_N),
|
||||
)
|
||||
|
||||
|
||||
def _reject_unsupported(
|
||||
bias,
|
||||
ce_weight,
|
||||
@@ -233,14 +237,34 @@ def _reject_unsupported(
|
||||
"label smoothing": bool(label_smoothing),
|
||||
"softcap": softcap is not None,
|
||||
"return_z_loss": return_z_loss,
|
||||
"accum_dtype": accum_dtype is not None,
|
||||
"token scaling": use_token_scaling,
|
||||
"return_token_accuracy": return_token_accuracy,
|
||||
"return_predicted_tokens": return_predicted_tokens,
|
||||
}
|
||||
enabled = [name for name, value in unsupported.items() if value]
|
||||
if enabled:
|
||||
raise NotImplementedError(f"cuTile SM90 FLCE does not support: {', '.join(enabled)}")
|
||||
raise NotImplementedError(f"cuTile FLCE does not support: {', '.join(enabled)}")
|
||||
# FP32 accumulation is the default and only supported accumulation dtype.
|
||||
if accum_dtype is not None and accum_dtype != torch.float32:
|
||||
raise NotImplementedError(
|
||||
f"cuTile FLCE only supports FP32 accumulation (accum_dtype=torch.float32 or None), got {accum_dtype!r}"
|
||||
)
|
||||
|
||||
|
||||
# Impl names that identify this cuTile kernel itself; accepted so callers that
|
||||
# route by name can target it explicitly. These are self-identity only -- this
|
||||
# Function does not dispatch to alternative CE implementations.
|
||||
_SELF_CE_IMPLS = (None, "cutile", "nvidia-cutile")
|
||||
|
||||
|
||||
def _validate_ce_dispatch(ce_impl, ce_mode) -> None:
|
||||
if ce_impl not in _SELF_CE_IMPLS:
|
||||
raise NotImplementedError(
|
||||
f"cuTile FLCE does not dispatch to ce_impl={ce_impl!r}; "
|
||||
f"expected one of {tuple(name for name in _SELF_CE_IMPLS)}"
|
||||
)
|
||||
if ce_mode is not None:
|
||||
raise NotImplementedError(f"cuTile FLCE only supports the default CE mode (ce_mode=None), got {ce_mode!r}")
|
||||
|
||||
|
||||
def _validate_inputs(
|
||||
@@ -249,14 +273,17 @@ def _validate_inputs(
|
||||
target: torch.Tensor,
|
||||
reduction: str,
|
||||
) -> None:
|
||||
if not torch.cuda.is_available() or torch.cuda.get_device_capability(_input.device) != (9, 0):
|
||||
raise RuntimeError("cuTile FLCE requires a Hopper (compute capability 9.0) GPU")
|
||||
if _input.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16:
|
||||
raise TypeError("cuTile FLCE supports BF16 input and weight only")
|
||||
if target.dtype != torch.int64:
|
||||
raise TypeError("target must be an int64 tensor")
|
||||
if _input.ndim != 2 or weight.ndim != 2 or target.ndim != 1:
|
||||
raise ValueError("expected input[M, H], weight[V, H], and target[M]")
|
||||
if 0 in _input.shape or 0 in weight.shape or target.shape[0] == 0:
|
||||
raise ValueError(
|
||||
f"cuTile FLCE does not support empty tensors: input {tuple(_input.shape)}, "
|
||||
f"weight {tuple(weight.shape)}, target {tuple(target.shape)}"
|
||||
)
|
||||
if _input.shape[0] != target.shape[0] or _input.shape[1] != weight.shape[1]:
|
||||
raise ValueError(
|
||||
f"incompatible input {tuple(_input.shape)}, weight {tuple(weight.shape)}, and target {tuple(target.shape)}"
|
||||
@@ -265,8 +292,18 @@ def _validate_inputs(
|
||||
raise ValueError(f"reduction must be 'mean' or 'sum', got {reduction!r}")
|
||||
if _input.device != weight.device or _input.device != target.device:
|
||||
raise ValueError("input, weight, and target must be on the same CUDA device")
|
||||
# Reject CPU tensors before touching CUDA device-capability queries.
|
||||
if _input.device.type != "cuda" or weight.device.type != "cuda" or target.device.type != "cuda":
|
||||
raise RuntimeError("cuTile FLCE requires CUDA (Hopper 9.0 or Blackwell SM100 10.0) tensors")
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("cuTile FLCE requires an available CUDA device")
|
||||
if not _input.is_contiguous() or not weight.is_contiguous() or not target.is_contiguous():
|
||||
raise ValueError("cuTile FLCE requires contiguous input, weight, and target tensors")
|
||||
capability = torch.cuda.get_device_capability(_input.device)
|
||||
if capability not in ((9, 0), (10, 0)):
|
||||
raise RuntimeError(
|
||||
f"cuTile FLCE requires a Hopper (compute capability 9.0) or Blackwell SM100 (10.0) GPU, got {capability}"
|
||||
)
|
||||
num_stats_partitions = (weight.shape[0] + LOGITS_STATS_BLOCK_SIZE - 1) // LOGITS_STATS_BLOCK_SIZE
|
||||
if num_stats_partitions > MAX_STATS_BLOCK_SIZE:
|
||||
raise NotImplementedError(
|
||||
@@ -274,6 +311,95 @@ def _validate_inputs(
|
||||
)
|
||||
|
||||
|
||||
def _ce_setup(_input, weight, target, reduction, ignore_index):
|
||||
"""Shared prologue: derive shapes, validate targets, and build the scales.
|
||||
|
||||
Returns the geometry the CE kernel needs plus ``loss_scale`` (passed to the
|
||||
kernel for the mean-reduced per-row loss) and ``gradient_scale`` (the global
|
||||
mean normalizer, ``1`` for ``sum``) applied to the gradients in backward.
|
||||
"""
|
||||
tokens, hidden_size = _input.shape
|
||||
vocab_size = weight.shape[0]
|
||||
num_vocab_tiles = (vocab_size + LOGITS_STATS_BLOCK_SIZE - 1) // LOGITS_STATS_BLOCK_SIZE
|
||||
stats_block_size = _next_power_of_2(num_vocab_tiles)
|
||||
needs_dz = _input.requires_grad or weight.requires_grad
|
||||
reduction_mean = reduction == "mean"
|
||||
|
||||
target_mask = target != ignore_index
|
||||
valid_target = ~target_mask | ((target >= 0) & (target < vocab_size))
|
||||
torch._assert_async(valid_target.all(), f"target values must be in [0, {vocab_size}) or equal ignore_index")
|
||||
# CE normalization uses the GLOBAL count of non-ignored tokens across all chunks.
|
||||
loss_scale = (
|
||||
target_mask.sum().clamp_min(1).to(torch.float32).reciprocal()
|
||||
if reduction_mean
|
||||
else torch.empty((), dtype=torch.float32, device=_input.device)
|
||||
)
|
||||
gradient_scale = loss_scale if reduction_mean else torch.ones((), dtype=torch.float32, device=_input.device)
|
||||
return (
|
||||
tokens,
|
||||
hidden_size,
|
||||
vocab_size,
|
||||
num_vocab_tiles,
|
||||
stats_block_size,
|
||||
needs_dz,
|
||||
reduction_mean,
|
||||
loss_scale,
|
||||
gradient_scale,
|
||||
)
|
||||
|
||||
|
||||
def _project_and_dz(
|
||||
input_chunk,
|
||||
weight,
|
||||
logits_chunk,
|
||||
target_chunk,
|
||||
loss_chunk,
|
||||
loss_scale,
|
||||
partial_max,
|
||||
partial_sum,
|
||||
completion_count,
|
||||
vocab_size,
|
||||
num_vocab_tiles,
|
||||
stats_block_size,
|
||||
ignore_index,
|
||||
needs_dz,
|
||||
reduction_mean,
|
||||
):
|
||||
"""Project a token chunk to logits and overwrite it in place with dZ.
|
||||
|
||||
``torch.mm`` fills ``logits_chunk`` with the BF16 projection; the fused CE
|
||||
kernel then writes the per-token loss and (when ``needs_dz``) replaces the
|
||||
logits with the softmax-minus-one-hot gradient dZ. Shared by the legacy
|
||||
unchunked path (one full-batch call) and the chunked path (per chunk).
|
||||
"""
|
||||
rows = logits_chunk.shape[0]
|
||||
torch.mm(input_chunk, weight.t(), out=logits_chunk)
|
||||
# Counters must be reset before every partition reduction.
|
||||
completion_count.zero_()
|
||||
_launch_cutile(
|
||||
logits_chunk.device,
|
||||
(num_vocab_tiles * rows, 1, 1),
|
||||
_fused_cross_entropy_dz_kernel,
|
||||
(
|
||||
logits_chunk,
|
||||
target_chunk,
|
||||
loss_chunk,
|
||||
loss_scale,
|
||||
partial_max,
|
||||
partial_sum,
|
||||
completion_count,
|
||||
vocab_size,
|
||||
num_vocab_tiles,
|
||||
ignore_index,
|
||||
LOGITS_STATS_BLOCK_SIZE,
|
||||
CE_BLOCK_SIZE,
|
||||
stats_block_size,
|
||||
needs_dz,
|
||||
reduction_mean,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def fused_linear_cross_entropy_forward(
|
||||
_input,
|
||||
weight,
|
||||
@@ -291,6 +417,15 @@ def fused_linear_cross_entropy_forward(
|
||||
return_token_accuracy=False,
|
||||
return_predicted_tokens=False,
|
||||
):
|
||||
"""Legacy unchunked forward.
|
||||
|
||||
Materializes the full ``[N, V]`` logits buffer, overwrites it with dZ, and
|
||||
returns that dZ as the third element (``grad_logits``). The paired
|
||||
:func:`fused_linear_cross_entropy_backward` consumes ``grad_logits`` to run
|
||||
the dW matmul. This is the original low-level contract, retained for
|
||||
compatibility; the autograd Function uses the bounded
|
||||
:func:`chunked_fused_linear_cross_entropy_forward` instead.
|
||||
"""
|
||||
_reject_unsupported(
|
||||
bias,
|
||||
ce_weight,
|
||||
@@ -305,47 +440,40 @@ def fused_linear_cross_entropy_forward(
|
||||
)
|
||||
_validate_inputs(_input, weight, target, reduction)
|
||||
|
||||
tokens, hidden_size = _input.shape
|
||||
vocab_size = weight.shape[0]
|
||||
num_vocab_tiles = (vocab_size + LOGITS_STATS_BLOCK_SIZE - 1) // LOGITS_STATS_BLOCK_SIZE
|
||||
stats_block_size = _next_power_of_2(num_vocab_tiles)
|
||||
needs_dz = _input.requires_grad or weight.requires_grad
|
||||
(
|
||||
tokens,
|
||||
hidden_size,
|
||||
vocab_size,
|
||||
num_vocab_tiles,
|
||||
stats_block_size,
|
||||
needs_dz,
|
||||
reduction_mean,
|
||||
loss_scale,
|
||||
gradient_scale,
|
||||
) = _ce_setup(_input, weight, target, reduction, ignore_index)
|
||||
|
||||
logits = torch.empty((tokens, vocab_size), dtype=torch.bfloat16, device=_input.device)
|
||||
loss_1d = torch.empty(tokens, dtype=torch.float32, device=_input.device)
|
||||
partial_max = torch.empty((tokens, num_vocab_tiles), dtype=torch.float32, device=_input.device)
|
||||
partial_sum = torch.empty_like(partial_max)
|
||||
completion_count = torch.zeros(tokens, dtype=torch.int32, device=_input.device)
|
||||
target_mask = target != ignore_index
|
||||
valid_target = ~target_mask | ((target >= 0) & (target < vocab_size))
|
||||
torch._assert_async(valid_target.all(), f"target values must be in [0, {vocab_size}) or equal ignore_index")
|
||||
loss_scale = (
|
||||
target_mask.sum().clamp_min(1).to(torch.float32).reciprocal()
|
||||
if reduction == "mean"
|
||||
else torch.empty((), dtype=torch.float32, device=_input.device)
|
||||
)
|
||||
|
||||
torch.mm(_input, weight.t(), out=logits)
|
||||
_launch_cutile(
|
||||
logits.device,
|
||||
(num_vocab_tiles * tokens, 1, 1),
|
||||
_fused_cross_entropy_dz_kernel,
|
||||
(
|
||||
logits,
|
||||
target,
|
||||
loss_1d,
|
||||
loss_scale,
|
||||
partial_max,
|
||||
partial_sum,
|
||||
completion_count,
|
||||
vocab_size,
|
||||
num_vocab_tiles,
|
||||
ignore_index,
|
||||
LOGITS_STATS_BLOCK_SIZE,
|
||||
CE_BLOCK_SIZE,
|
||||
stats_block_size,
|
||||
needs_dz,
|
||||
reduction == "mean",
|
||||
),
|
||||
_project_and_dz(
|
||||
_input,
|
||||
weight,
|
||||
logits,
|
||||
target,
|
||||
loss_1d,
|
||||
loss_scale,
|
||||
partial_max,
|
||||
partial_sum,
|
||||
completion_count,
|
||||
vocab_size,
|
||||
num_vocab_tiles,
|
||||
stats_block_size,
|
||||
ignore_index,
|
||||
needs_dz,
|
||||
reduction_mean,
|
||||
)
|
||||
|
||||
grad_input = None
|
||||
@@ -354,7 +482,8 @@ def fused_linear_cross_entropy_forward(
|
||||
torch.mm(logits, weight, out=grad_input)
|
||||
|
||||
loss = loss_1d.sum()
|
||||
gradient_scale = loss_scale if reduction == "mean" else torch.ones((), dtype=torch.float32, device=_input.device)
|
||||
# Third element is the full [N, V] dZ (legacy contract): the paired backward
|
||||
# runs the dW matmul from it. Scaling is deferred to backward.
|
||||
return loss, grad_input, logits if weight.requires_grad else None, gradient_scale
|
||||
|
||||
|
||||
@@ -367,19 +496,191 @@ def fused_linear_cross_entropy_backward(
|
||||
weight_shape: torch.Size,
|
||||
weight_dtype: torch.dtype,
|
||||
):
|
||||
"""Legacy 7-argument backward paired with the unchunked forward.
|
||||
|
||||
Applies ``upstream * gradient_scale`` to dX. dW is computed in FP32 by cuBLAS
|
||||
(``dZ.T @ X`` accumulated into an FP32 buffer), then the combined
|
||||
``upstream * gradient_scale`` factor is applied in FP32 and the result cast to
|
||||
``weight_dtype`` exactly once by the fused scale-and-cast kernel. Casting
|
||||
after (not before) the upstream multiply preserves small gradients that would
|
||||
flush to zero if rounded to BF16 first. All operations are out-of-place, so
|
||||
the saved tensors are not mutated across repeated backward calls.
|
||||
"""
|
||||
total_scale = grad_output * gradient_scale
|
||||
grad_input_out = None
|
||||
if grad_input is not None:
|
||||
grad_input = grad_input * total_scale
|
||||
grad_input_out = grad_input * total_scale
|
||||
|
||||
grad_weight = None
|
||||
if grad_logits is not None:
|
||||
dweight_accum = torch.empty(weight_shape, dtype=torch.float32, device=grad_logits.device)
|
||||
_accumulate_dw_fp32(dweight_accum, grad_logits, _input, accumulate=False)
|
||||
combined_scale = total_scale.to(torch.float32)
|
||||
grad_weight = torch.empty(weight_shape, dtype=weight_dtype, device=grad_logits.device)
|
||||
_launch_matmul_tn(grad_logits, _input, grad_weight, total_scale)
|
||||
_launch_scale_cast(dweight_accum, grad_weight, combined_scale)
|
||||
|
||||
return grad_input, grad_weight
|
||||
return grad_input_out, grad_weight
|
||||
|
||||
|
||||
def chunked_fused_linear_cross_entropy_forward(
|
||||
_input,
|
||||
weight,
|
||||
target,
|
||||
ce_weight=None,
|
||||
bias=None,
|
||||
ignore_index=-100,
|
||||
lse_square_scale=0.0,
|
||||
label_smoothing=0.0,
|
||||
reduction="mean",
|
||||
softcap=None,
|
||||
return_z_loss=False,
|
||||
accum_dtype=None,
|
||||
use_token_scaling=False,
|
||||
return_token_accuracy=False,
|
||||
return_predicted_tokens=False,
|
||||
):
|
||||
"""Token-chunked forward used by the autograd Function.
|
||||
|
||||
Processes tokens in row chunks so the transient logits/dZ workspace is
|
||||
``O(min(N, C) * V)``, where the token-chunk size ``C`` comes from the shared
|
||||
:func:`liger_kernel.ops.fused_linear_cross_entropy._get_chunk_size` helper --
|
||||
the EXACT same default geometry as the Triton FLCE backend, with no public
|
||||
``chunk_size`` knob. dX is written as the raw BF16 projection and dW is
|
||||
accumulated in an FP32 buffer; NEITHER is scaled here. The returned third
|
||||
element is the raw FP32 dW accumulator (``O(V*H)``), which
|
||||
:func:`chunked_fused_linear_cross_entropy_backward` scales and casts to the
|
||||
weight dtype exactly once -- deferring the BF16 cast preserves small
|
||||
gradients that would flush to zero if cast per-chunk in the forward.
|
||||
"""
|
||||
_reject_unsupported(
|
||||
bias,
|
||||
ce_weight,
|
||||
lse_square_scale,
|
||||
label_smoothing,
|
||||
softcap,
|
||||
return_z_loss,
|
||||
accum_dtype,
|
||||
use_token_scaling,
|
||||
return_token_accuracy,
|
||||
return_predicted_tokens,
|
||||
)
|
||||
_validate_inputs(_input, weight, target, reduction)
|
||||
|
||||
(
|
||||
tokens,
|
||||
hidden_size,
|
||||
vocab_size,
|
||||
num_vocab_tiles,
|
||||
stats_block_size,
|
||||
needs_dz,
|
||||
reduction_mean,
|
||||
loss_scale,
|
||||
gradient_scale,
|
||||
) = _ce_setup(_input, weight, target, reduction, ignore_index)
|
||||
|
||||
chunk = _get_chunk_size(tokens, hidden_size, vocab_size)
|
||||
buffer_rows = min(tokens, chunk)
|
||||
|
||||
# Reusable per-chunk logits/dZ buffer of shape [min(N, C), V] and the
|
||||
# partitioned CE statistics/counter scratch of shape [min(N, C), num_partitions].
|
||||
logits = torch.empty((buffer_rows, vocab_size), dtype=torch.bfloat16, device=_input.device)
|
||||
loss_1d = torch.empty(tokens, dtype=torch.float32, device=_input.device)
|
||||
partial_max = torch.empty((buffer_rows, num_vocab_tiles), dtype=torch.float32, device=_input.device)
|
||||
partial_sum = torch.empty_like(partial_max)
|
||||
completion_count = torch.zeros(buffer_rows, dtype=torch.int32, device=_input.device)
|
||||
|
||||
grad_input = None
|
||||
if _input.requires_grad:
|
||||
grad_input = torch.empty((tokens, hidden_size), dtype=_input.dtype, device=_input.device)
|
||||
|
||||
# FP32 dW accumulator reused across chunks. It is retained UNSCALED for
|
||||
# backward; the mean normalizer and upstream scalar are combined and the
|
||||
# single BF16 cast is applied there, never per chunk.
|
||||
dweight_accum = None
|
||||
if weight.requires_grad:
|
||||
dweight_accum = torch.empty((vocab_size, hidden_size), dtype=torch.float32, device=_input.device)
|
||||
|
||||
for chunk_index, start in enumerate(range(0, tokens, chunk)):
|
||||
stop = min(start + chunk, tokens)
|
||||
rows = stop - start
|
||||
logits_chunk = logits[:rows]
|
||||
input_chunk = _input[start:stop]
|
||||
|
||||
_project_and_dz(
|
||||
input_chunk,
|
||||
weight,
|
||||
logits_chunk,
|
||||
target[start:stop],
|
||||
loss_1d[start:stop],
|
||||
loss_scale,
|
||||
partial_max[:rows],
|
||||
partial_sum[:rows],
|
||||
completion_count[:rows],
|
||||
vocab_size,
|
||||
num_vocab_tiles,
|
||||
stats_block_size,
|
||||
ignore_index,
|
||||
needs_dz,
|
||||
reduction_mean,
|
||||
)
|
||||
|
||||
# logits_chunk now holds dZ. dX is scaled AFTER the low-precision GEMM
|
||||
# (in backward), matching the baseline; store the raw projection here.
|
||||
if grad_input is not None:
|
||||
torch.mm(logits_chunk, weight, out=grad_input[start:stop])
|
||||
|
||||
# Accumulate raw dZ.T @ X into the FP32 dW buffer via cuBLAS: overwrite
|
||||
# (mm) on the first chunk, add (addmm) on subsequent chunks.
|
||||
if dweight_accum is not None:
|
||||
_accumulate_dw_fp32(dweight_accum, logits_chunk, input_chunk, chunk_index > 0)
|
||||
|
||||
loss = loss_1d.sum()
|
||||
# Retain the RAW FP32 dW accumulator (scaled + cast in backward) so small
|
||||
# gradients are not lost to an early BF16 cast.
|
||||
return loss, grad_input, dweight_accum, gradient_scale
|
||||
|
||||
|
||||
def chunked_fused_linear_cross_entropy_backward(
|
||||
grad_output: torch.Tensor,
|
||||
grad_input: Optional[torch.Tensor],
|
||||
dweight_accum: Optional[torch.Tensor],
|
||||
gradient_scale: torch.Tensor,
|
||||
weight_dtype: torch.dtype,
|
||||
):
|
||||
"""Backward for the chunked forward.
|
||||
|
||||
dX is the raw BF16 projection, scaled out-of-place by
|
||||
``upstream * gradient_scale``. dW is the retained FP32 accumulator: the
|
||||
combined ``upstream * gradient_scale`` factor is applied in FP32 and the
|
||||
result cast to ``weight_dtype`` exactly once by a fused scale-and-cast
|
||||
kernel (no O(V*H) FP32 temporary, and small gradients survive the cast).
|
||||
Every output is a fresh tensor, so the saved gradients are never mutated --
|
||||
repeated backward with a different upstream (including ``1``) stays correct.
|
||||
"""
|
||||
grad_input_out = None
|
||||
if grad_input is not None:
|
||||
grad_input_out = grad_input * (grad_output * gradient_scale)
|
||||
|
||||
grad_weight_out = None
|
||||
if dweight_accum is not None:
|
||||
combined_scale = (grad_output * gradient_scale).to(torch.float32)
|
||||
grad_weight_out = torch.empty(dweight_accum.shape, dtype=weight_dtype, device=dweight_accum.device)
|
||||
_launch_scale_cast(dweight_accum, grad_weight_out, combined_scale)
|
||||
|
||||
return grad_input_out, grad_weight_out
|
||||
|
||||
|
||||
class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
|
||||
# This Function computes dX/dW during forward with token chunking. The
|
||||
# transient logits/dZ workspace is O(min(N, C) * V) instead of the full
|
||||
# O(N * V); the retained gradients add O(N*H + V*H). The token-chunk size C
|
||||
# is the shared default Triton geometry (_get_chunk_size); there is no public
|
||||
# chunk_size knob.
|
||||
# ce_impl / ce_mode are accepted only as self-identity placeholders (see
|
||||
# _validate_ce_dispatch); this Function does NOT dispatch to alternative CE
|
||||
# implementations, so it does not advertise inner-impl dispatch.
|
||||
supports_inner_impl_dispatch = False
|
||||
|
||||
@staticmethod
|
||||
@amp_custom_fwd
|
||||
def forward(
|
||||
@@ -399,8 +700,11 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
|
||||
use_token_scaling=False,
|
||||
return_token_accuracy=False,
|
||||
return_predicted_tokens=False,
|
||||
ce_impl=None,
|
||||
ce_mode=None,
|
||||
):
|
||||
loss, grad_input, grad_logits, gradient_scale = fused_linear_cross_entropy_forward(
|
||||
_validate_ce_dispatch(ce_impl, ce_mode)
|
||||
loss, grad_input, dweight_accum, gradient_scale = chunked_fused_linear_cross_entropy_forward(
|
||||
_input=_input,
|
||||
weight=weight,
|
||||
target=target,
|
||||
@@ -418,13 +722,10 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
|
||||
return_predicted_tokens=return_predicted_tokens,
|
||||
)
|
||||
|
||||
ctx.save_for_backward(
|
||||
grad_input,
|
||||
grad_logits,
|
||||
_input.detach() if weight.requires_grad else None,
|
||||
gradient_scale,
|
||||
)
|
||||
ctx.weight_shape = weight.shape
|
||||
# Retain the raw BF16 dX, the raw FP32 dW accumulator, and the scale.
|
||||
# Neither the full [N, V] dZ nor the input activations are saved. The
|
||||
# target weight dtype is tracked so backward can cast dW once.
|
||||
ctx.save_for_backward(grad_input, dweight_accum, gradient_scale)
|
||||
ctx.weight_dtype = weight.dtype
|
||||
return loss, None, None, None
|
||||
|
||||
@@ -432,37 +733,39 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
|
||||
@amp_custom_bwd
|
||||
def backward(ctx, grad_output, grad_output2, grad_output3, grad_output4):
|
||||
del grad_output2, grad_output3, grad_output4
|
||||
grad_input, grad_logits, _input, gradient_scale = ctx.saved_tensors
|
||||
grad_input, grad_weight = fused_linear_cross_entropy_backward(
|
||||
grad_input, dweight_accum, gradient_scale = ctx.saved_tensors
|
||||
grad_input, grad_weight = chunked_fused_linear_cross_entropy_backward(
|
||||
grad_output,
|
||||
grad_input,
|
||||
grad_logits,
|
||||
_input,
|
||||
dweight_accum,
|
||||
gradient_scale,
|
||||
ctx.weight_shape,
|
||||
ctx.weight_dtype,
|
||||
)
|
||||
return (
|
||||
grad_input,
|
||||
grad_weight,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None, # target
|
||||
None, # bias
|
||||
None, # ce_weight
|
||||
None, # ignore_index
|
||||
None, # lse_square_scale
|
||||
None, # label_smoothing
|
||||
None, # reduction
|
||||
None, # softcap
|
||||
None, # return_z_loss
|
||||
None, # accum_dtype
|
||||
None, # use_token_scaling
|
||||
None, # return_token_accuracy
|
||||
None, # return_predicted_tokens
|
||||
None, # ce_impl
|
||||
None, # ce_mode
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LigerFusedLinearCrossEntropyFunction",
|
||||
"chunked_fused_linear_cross_entropy_backward",
|
||||
"chunked_fused_linear_cross_entropy_forward",
|
||||
"fused_linear_cross_entropy_backward",
|
||||
"fused_linear_cross_entropy_forward",
|
||||
]
|
||||
|
||||
@@ -40,6 +40,19 @@ _TORCH_VERSION = Version(torch.__version__.split("+")[0])
|
||||
_ADDMM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0")
|
||||
|
||||
|
||||
def _get_chunk_size(BT, H, V):
|
||||
"""Token-chunk size shared by the Triton and cuTile FLCE backends.
|
||||
|
||||
Widens the transient logits budget to ``_CHUNK_MEM_CONST x BT x H`` (C=1 is a
|
||||
memory floor, not a perf target): ``inc_factor = cdiv(V, C*H)``, then the
|
||||
chunk is the next power of two of ``cdiv(BT, inc_factor)``, clamped to ``BT``
|
||||
so a single chunk covers ``BT`` when the budget allows (and when ``V <= C*H``).
|
||||
"""
|
||||
inc_factor = triton.cdiv(V, _CHUNK_MEM_CONST * H)
|
||||
chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor)) # (BT + inc_factor - 1) // inc_factor
|
||||
return min(chunk_size, BT) # a single chunk covers BT when the budget allows; never exceed BT
|
||||
|
||||
|
||||
def fused_linear_cross_entropy_forward(
|
||||
_input,
|
||||
weight,
|
||||
@@ -89,9 +102,7 @@ def fused_linear_cross_entropy_forward(
|
||||
V = weight.shape[0]
|
||||
|
||||
# widen the transient logits budget to C x BT x H (C=1 is a memory floor, not a perf target)
|
||||
inc_factor = triton.cdiv(V, _CHUNK_MEM_CONST * H)
|
||||
chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor)) # (BT + inc_factor - 1) // inc_factor
|
||||
chunk_size = min(chunk_size, BT) # a single chunk covers BT when the budget allows; never exceed BT
|
||||
chunk_size = _get_chunk_size(BT, H, V)
|
||||
num_chunks = triton.cdiv(BT, chunk_size) # (BT + chunk_size - 1) // chunk_size
|
||||
|
||||
grad_input = torch.zeros_like(_input, device=device)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user