linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add FusedLinearKLDivLoss (fused linear + KL divergence for distillation) (#1423)

Browse Source 
## Summary

Add `FusedLinearKLDivLoss`, which fuses the vocabulary-head linear
projection with
KL(Q || P) loss against an explicit target distribution, for knowledge
distillation
workflows where teacher probabilities are available directly (e.g.
cached / offline
soft labels), complementing `LigerKLDIVLoss` (element-wise, no fusion)
and
`LigerFusedLinearJSD` (teacher side needs its own projection).

The kernel follows the chunked gradient-in-forward structure of
`fused_linear_jsd.py`,
so the full `BT x V` logits tensor is never materialized, and adopts the
`_CHUNK_MEM_CONST = 16` token-chunk geometry of
`fused_linear_cross_entropy.py`.

## Details

- Supports `reduction` (`"batchmean"` / `"mean"` / `"sum"`, same
semantics as
`torch.nn.KLDivLoss`), `ignore_index` via `shift_labels`, softmax
`temperature`,
  and an `eps` clamp so that `0 * log(0)` contributes 0.
- Supports `accum_dtype` for fp32 weight-gradient accumulation across
chunked
  low-precision GEMMs (same semantics as `LigerFusedLinearJSD` /
`LigerFusedLinearCrossEntropyLoss`), including the
`torch.addmm(out_dtype=fp32)`
  fast path on torch >= 2.8 / sm80+ and a fallback elsewhere.
- The log-softmax backward is computed in place on the saved log-probs
and the
`chunk_size x V` fp32 temporaries are freed eagerly, so the peak during
the
  trailing GEMMs stays low.
- `reduction="none"` is intentionally unsupported (gradient-in-forward
requires a
  scalar loss); a `ValueError` points users to `LigerKLDIVLoss` instead.
- bf16 + `reduction="sum"` test tolerances follow the existing
convention in
  `test_fused_linear_cross_entropy.py`.
- Peak-memory improvement over the unfused PyTorch baseline (full
fwd+bwd,
`max_memory_allocated`, bf16, llama-3 head shape H=4096 / V=128256, RTX
4060
Laptop GPU): BT=2048 6537 -> 4310 MB (1.52x), BT=4096 12072 -> 6594 MB
(1.83x);
at BT=1024 the benchmark harness reports 4793 -> 4244 MB. Runtime is at
parity
  with the baseline at BT=1024 (0.89x).

## Testing Done

- Hardware Type: NVIDIA GeForce RTX 4060 Laptop GPU (sm89); NVIDIA Tesla
T4 (sm75)

- [x] run `make test` to ensure correctness (32 pre-existing failures
unrelated to
this change: 18 `test_mlp.py` shared-memory OOR on sm89 and 14
`test_grpo_loss.py`
  LUSPO numerics; both reproduce on `main` without this change)
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence (does not
exercise the new
  standalone loss)

<details>
<summary>Test logs (RTX 4060 Laptop, torch 2.11.0+cu128, triton
3.7.1)</summary>

```
make test: 32 failed, 4152 passed, 914 skipped, 15 xfailed in 1050.50s (0:17:30)
  test/transformers/test_fused_linear_kl_div.py: 172 passed, 0 failed, 0 skipped
  failures (all pre-existing on this GPU, reproduced on main without this change):
    18 x test/transformers/test_mlp.py
        triton.runtime.errors.OutOfResources: shared memory
        Required: 131072, Hardware limit: 101376 (sm89 has 100 KB usable smem)
    14 x test/chunked_loss/test_grpo_loss.py
        chunked GRPO LUSPO numerics (also seen by other contributors, see #1397)
make checkstyle: ruff check and ruff format --check pass on all files in this PR
```

</details>

<details>
<summary>Test logs (Kaggle Tesla T4, torch 2.10.0+cu128, triton
3.6.0)</summary>

```
test/transformers/test_fused_linear_kl_div.py: 51 passed, 25 skipped in 10.29s
(bf16 cases skipped by design: T4 is sm75 without native bf16 support)
```

</details>
This commit is contained in:
Yulong Cao
2026-09-14 02:34:33 +00:00
committed by GitHub
parent c64878355b
commit f6385e909b
9 changed files with 1241 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1
View File
@@ -471,6 +471,7 @@ loss.backward()
| KLDivergence | `liger_kernel.transformers.LigerKLDIVLoss` |
| JSD | `liger_kernel.transformers.LigerJSD` |
| Fused Linear JSD | `liger_kernel.transformers.LigerFusedLinearJSD` |
| Fused Linear KL Divergence | `liger_kernel.transformers.LigerFusedLinearKLDivLoss` |
| TVD | `liger_kernel.transformers.LigerTVDLoss` |
### Experimental Kernels
benchmark/scripts/benchmark_fused_linear_kl_div.py
+163
View File
@@ -0,0 +1,163 @@
import torch
from benchmark_model_configs import MODEL_REGISTRY
from benchmark_model_configs import build_model_config_sweep
from benchmark_model_configs import build_token_length_sweep
from benchmark_model_configs import get_benchmark_model_config
from utils import SingleBenchmarkRunInput
from utils import build_memory_bench_fn
from utils import build_speed_bench_fn
from utils import parse_benchmark_script_args
from utils import run_benchmarks
from liger_kernel.transformers.fused_linear_kl_div import LigerFusedLinearKLDivLoss
from liger_kernel.utils import infer_device
device = infer_device()
class TorchLMHeadKLDiv(torch.nn.Module):
"""Ground truth implementation of the linear fused with torch based KL divergence loss.
:param H: hidden size
:param V: vocab size
:param temperature: softmax temperature
"""
def __init__(
self,
H: int,
V: int,
dtype: torch.dtype,
device: torch.device,
ignore_index: int = -100,
temperature: float = 1.0,
eps: float = 1e-10,
):
super(TorchLMHeadKLDiv, self).__init__()
self.student_lin = torch.nn.Linear(in_features=H, out_features=V, bias=False, dtype=dtype, device=device)
self.ignore_index = ignore_index
self.temperature = temperature
self.eps = eps
self.dtype = dtype
def forward(self, student_input, target, label=None):
logits = self.student_lin(student_input).to(torch.float32) / self.temperature
log_p = torch.log_softmax(logits, dim=-1)
q = target.to(torch.float32)
# KL(q || p) = sum(q * (log q - log p)); 0 * log 0 is treated as 0 via the clamp
loss_mat = q * (torch.log(q.clamp_min(self.eps)) - log_p)
if label is not None:
keep = (label != self.ignore_index).to(torch.float32).unsqueeze(-1)
loss_mat = loss_mat * keep
n_non_ignore = int(keep.sum().item())
else:
n_non_ignore = student_input.shape[0]
total = loss_mat.sum()
if n_non_ignore == 0:
return (total * 0.0).to(self.dtype)
return (total / n_non_ignore).to(self.dtype)
class LigerLMHeadKLDiv(torch.nn.Module):
def __init__(
self,
H: int,
V: int,
dtype: torch.dtype,
device: torch.device,
ignore_index: int = -100,
temperature: float = 1.0,
):
super().__init__()
self.student_lin = torch.nn.Linear(in_features=H, out_features=V, bias=False, dtype=dtype, device=device)
self.fused_kl = LigerFusedLinearKLDivLoss(
ignore_index=ignore_index,
temperature=temperature,
)
def forward(self, student_input, target, label=None):
return self.fused_kl(student_input, self.student_lin.weight, target, label)
def setup_fused_linear_kl_div(input: SingleBenchmarkRunInput):
"""Create input tensors and fused linear KL divergence from benchmark config."""
cfg = input.extra_benchmark_config
if isinstance(input.x, str):
model_cfg = MODEL_REGISTRY[input.x]
BT = cfg["seq_len"] * cfg["bsz"]
V = model_cfg.vocab_size
H = model_cfg.hidden_size
dtype = model_cfg.dtype
else:
BT = input.x
V = cfg["vocab_size"]
H = cfg["hidden_size"]
dtype = cfg["dtype"]
torch_lm_head_kl = TorchLMHeadKLDiv(H=H, V=V, dtype=dtype, device=device).to(device)
liger_lm_head_kl = LigerLMHeadKLDiv(H=H, V=V, dtype=dtype, device=device).to(device)
# init the linear in all FusedLinearKLDivs with the same weights
torch_lm_head_kl.student_lin.weight.data = liger_lm_head_kl.student_lin.weight.data = torch.rand(
V, H, device=device, dtype=dtype
)
student_input = torch.rand(BT, H, requires_grad=True, dtype=dtype, device=device)
target = torch.rand(BT, V, device=device, dtype=torch.float32).softmax(dim=-1).to(dtype)
if input.kernel_provider == "liger":
lm_head = liger_lm_head_kl
elif input.kernel_provider == "torch":
lm_head = torch_lm_head_kl
else:
raise ValueError(f"Invalid provider: {input.kernel_provider} for FusedLinearKLDiv")
return student_input, lambda _: lm_head(student_input, target)
if __name__ == "__main__":
args = parse_benchmark_script_args()
if args.sweep_mode == "model_config":
common_configs = build_model_config_sweep(
kernel_name="fused_linear_kl_div",
setup_fn=setup_fused_linear_kl_div,
model_keys=["hidden_size", "vocab_size", "dtype"],
probe_dim="BT",
probe_provider="torch",
bt=args.bt,
overwrite=args.overwrite,
)
else:
model = get_benchmark_model_config(args.model)
common_configs = build_token_length_sweep(
kernel_name="fused_linear_kl_div",
probe_x=1024,
model=model,
setup_fn=setup_fused_linear_kl_div,
model_keys=["hidden_size", "vocab_size", "dtype"],
scale_dim="BT",
x_label="B * T",
probe_provider="torch",
overwrite=args.overwrite,
)
common_configs["kernel_providers"] = ["torch", "liger"]
run_benchmarks(
bench_test_fn=build_speed_bench_fn(setup_fused_linear_kl_div),
kernel_operation_modes=["forward", "backward", "full"],
metric_name="speed",
metric_unit="ms",
**common_configs,
)
run_benchmarks(
bench_test_fn=build_memory_bench_fn(setup_fused_linear_kl_div),
kernel_operation_modes=["full"],
metric_name="memory",
metric_unit="MB",
**common_configs,
)
docs/Low-Level-APIs.md
+1
View File
@@ -151,6 +151,7 @@ Functional APIs are also available:
| KLDivergence | `liger_kernel.transformers.LigerKLDIVLoss` |
| JSD | `liger_kernel.transformers.LigerJSD` |
| Fused Linear JSD | `liger_kernel.transformers.LigerFusedLinearJSD` |
| Fused Linear KL Divergence | `liger_kernel.transformers.LigerFusedLinearKLDivLoss` |
## Experimental Kernels
src/liger_kernel/ops/__init__.py
+3
View File
@@ -49,6 +49,9 @@ from liger_kernel.ops.fused_linear_cross_entropy import fused_linear_cross_entro
from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction # noqa: F401
from liger_kernel.ops.fused_linear_jsd import fused_linear_jsd_backward # noqa: F401
from liger_kernel.ops.fused_linear_jsd import fused_linear_jsd_forward # noqa: F401
from liger_kernel.ops.fused_linear_kl_div import LigerFusedLinearKLDivFunction # noqa: F401
from liger_kernel.ops.fused_linear_kl_div import fused_linear_kl_div_backward # noqa: F401
from liger_kernel.ops.fused_linear_kl_div import fused_linear_kl_div_forward # noqa: F401
from liger_kernel.ops.fused_linear_scaled_cross_entropy import LigerFusedLinearScaledCrossEntropyFunction # noqa: F401
from liger_kernel.ops.fused_linear_scaled_cross_entropy import (
LigerFusedLinearScaledCrossEntropyTPFunction, # noqa: F401
src/liger_kernel/ops/fused_linear_kl_div.py
+391
View File
@@ -0,0 +1,391 @@
from typing import Literal
from typing import Optional
import torch
import triton
import triton.language as tl
from packaging.version import Version
from liger_kernel.ops.utils import amp_custom_bwd
from liger_kernel.ops.utils import amp_custom_fwd
from liger_kernel.ops.utils import element_mul_kernel
from liger_kernel.ops.utils import is_hip
from liger_kernel.utils import infer_device
MAX_FUSED_SIZE = 4096 if infer_device() == "xpu" else 65536 // 2
# Memory-budget constant for the token-chunk geometry below. The historical
# formula (inc_factor = cdiv(V, H), i.e. budget C=1) sizes chunks to a memory
# FLOOR of ~BT x H transient logits, which at LLM vocab shapes forces many tiny
# chunks and a launch-bound wall. Same constant and rationale as
# fused_linear_cross_entropy.py; see the measured numbers in its comment.
_CHUNK_MEM_CONST = 16
_TORCH_VERSION = Version(torch.__version__.split("+")[0])
_ADDMM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0")
REDUCTION_LITERAL = Literal["batchmean", "mean", "sum"]
def get_num_warps(BLOCK_SIZE):
num_warps = 4
if BLOCK_SIZE >= 32768:
num_warps = 32 if not is_hip() else 16
elif BLOCK_SIZE >= 8192:
num_warps = 16
elif BLOCK_SIZE >= 2048:
num_warps = 8
return num_warps
@triton.jit
def _kl_div_kernel(
X_ptr, # [BT, V] student logits (x @ W^T), overwritten in place with dL/dlogits
X_stride,
Q_ptr, # [BT, V] target probabilities
Q_stride,
loss_ptr, # [BT], per-row loss (already scaled by reduction)
loss_stride,
label_ptr,
ignore_index: tl.constexpr,
n_cols,
temperature,
eps,
scale, # pre-computed reduction scale fused into loss and gradients
BLOCK_SIZE: tl.constexpr,
HAS_LABEL: tl.constexpr,
):
pid = tl.program_id(0).to(tl.int64)
X_ptr += pid * X_stride
Q_ptr += pid * Q_stride
loss_ptr += pid * loss_stride
label_ptr += pid
if HAS_LABEL:
label = tl.load(label_ptr)
if label == ignore_index:
for i in range(0, n_cols, BLOCK_SIZE):
offsets = i + tl.arange(0, BLOCK_SIZE)
tl.store(X_ptr + offsets, 0.0, mask=offsets < n_cols)
return
# Everything below runs on fp32 values upcast from X on load (registers),
# so HBM only ever sees the native-precision logits buffer. Y, the student
# log-probabilities, is never materialized: Y = x / T - logsumexp(x / T).
# Pass 1: row max of x / T (masked lanes forced to -inf, sign-of-T safe).
m = float("-inf")
for i in range(0, n_cols, BLOCK_SIZE):
offsets = i + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_cols
x = tl.load(X_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
x = tl.where(mask, x / temperature, float("-inf"))
m = tl.maximum(m, tl.max(x, axis=0))
# Pass 2: softmax denominator, plus the loss pieces that don't need lse yet.
# KL(Q || P) = sum(Q * (log(Q) - Y))
# = sum(Q * (log(max(Q, eps)) - x / T)) + lse * sum(Q) (= a1 + lse * a2)
# 0 * log 0 is treated as 0 by clamping Q to eps before the log.
a1 = 0.0
a2 = 0.0
s = 0.0
for i in range(0, n_cols, BLOCK_SIZE):
offsets = i + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_cols
x = tl.load(X_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
x = tl.where(mask, x / temperature, float("-inf"))
q = tl.load(Q_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
s += tl.sum(tl.exp(x - m), axis=0)
a1 += tl.sum(tl.where(mask, q * (tl.log(tl.maximum(q, eps)) - x), 0.0), axis=0)
a2 += tl.sum(q, axis=0)
lse = m + tl.log(s)
tl.store(loss_ptr, (a1 + lse * a2) * scale)
# Pass 3: gradient w.r.t. the logits, stored in place over X. Backprop
# through log-softmax with g = dL/dY = -Q * scale:
# dL/dlogits = (g - softmax(x/T) * sum(g)) / T = -scale * (Q - softmax * a2) / T
for i in range(0, n_cols, BLOCK_SIZE):
offsets = i + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_cols
x = tl.load(X_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
x = tl.where(mask, x / temperature, float("-inf"))
q = tl.load(Q_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
softmax = tl.exp(x - m) / s
d = -scale * (q - softmax * a2) / temperature
tl.store(X_ptr + offsets, d, mask=mask)
def fused_linear_kl_div_forward(
student_input,
student_weight,
target,
shift_labels,
reduction,
ignore_index,
has_label,
temperature,
eps,
accum_dtype=None,
):
device = student_input.device
# inputs have shape: BT x H
# the only chunk-sized transient is the native-precision logits/grad buffer of
# shape chunk_size x V (the FP32 math happens in-kernel on upcast values in
# registers); partitioning the number of tokens BT into chunks keeps it within
# the transient budget of _CHUNK_MEM_CONST x BT x H (see fused_linear_cross_entropy.py):
# inc_factor = (V + _CHUNK_MEM_CONST*H - 1) // (_CHUNK_MEM_CONST*H), chunk_size = (BT + inc_factor - 1) // inc_factor
BT, H = student_input.shape
V = student_weight.shape[0]
BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V))
inc_factor = triton.cdiv(V, _CHUNK_MEM_CONST * H)
chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor))
chunk_size = min(chunk_size, BT) # a single chunk covers BT when the budget allows; never exceed BT
num_chunks = triton.cdiv(BT, chunk_size) # (BT + chunk_size - 1) // chunk_size
if accum_dtype is None:
grad_weight = torch.zeros_like(student_weight, device=device) if student_weight.requires_grad else None
else:
grad_weight = (
torch.zeros_like(student_weight, dtype=accum_dtype, device=device) if student_weight.requires_grad else None
)
grad_input = torch.zeros_like(student_input)
# we use fp32 for the per-row loss accumulator
loss_1d = torch.zeros((BT,), dtype=torch.float32, device=device)
if has_label:
n_non_ignore = (shift_labels != ignore_index).sum().item()
else:
n_non_ignore = BT
# Pre-compute reduction scale (fused into the kernel; matches torch.nn.KLDivLoss semantics)
if n_non_ignore == 0:
scale = 0.0
elif reduction == "batchmean":
scale = 1.0 / n_non_ignore
elif reduction == "mean":
scale = 1.0 / (n_non_ignore * V)
else: # "sum"
scale = 1.0
for chunk_id in range(num_chunks):
start_idx = chunk_id * chunk_size
end_idx = min((chunk_id + 1) * chunk_size, BT)
# chunk the inputs, shape: chunk_size x H / chunk_size x V
input_chunk = student_input[start_idx:end_idx]
target_chunk = target[start_idx:end_idx]
# shape: chunk_size x V. The GEMM output keeps its native precision in
# HBM; the FP32 upcast, temperature scaling, log-softmax and the
# log-softmax backward all happen inside _kl_div_kernel in registers,
# so this is the only chunk-sized transient and the kernel overwrites
# it in place with the gradient w.r.t. the logits.
logits_chunk = input_chunk @ student_weight.t()
chunk_n_rows = logits_chunk.shape[0]
_kl_div_kernel[(chunk_n_rows,)](
X_ptr=logits_chunk,
X_stride=logits_chunk.stride(-2),
Q_ptr=target_chunk,
Q_stride=target_chunk.stride(-2),
loss_ptr=loss_1d[start_idx:end_idx],
loss_stride=loss_1d[start_idx:end_idx].stride(0),
label_ptr=(
shift_labels[start_idx:end_idx] if has_label else torch.empty(1, device=device)
), # dummy ptr if no label
ignore_index=ignore_index,
n_cols=V,
temperature=temperature,
eps=eps,
scale=scale,
BLOCK_SIZE=BLOCK_SIZE,
HAS_LABEL=has_label,
num_warps=get_num_warps(BLOCK_SIZE),
)
# logits_chunk now holds dL/dlogits, shape: chunk_size x V
grad_input[start_idx:end_idx] = logits_chunk @ student_weight
if grad_weight is not None:
grad_logits_t = logits_chunk.t()
if (
_ADDMM_SUPPORTS_OUT_DTYPE
and grad_weight.device.type == "cuda"
and torch.cuda.get_device_capability(grad_weight.device)[0] >= 8
and grad_weight.dtype == torch.float32
and grad_logits_t.dtype in (torch.float16, torch.bfloat16)
):
# FP32 accumulator (accum_dtype=torch.float32). Unlike torch.mm,
# torch.addmm's out_dtype path does not participate in autocast operand
# casting, so the accumulator can stay fp32 while grad_logits is the
# autocast dtype. addmm requires mat1 and mat2 to share a dtype, so align
# input_chunk before accumulating.
chunk_input = input_chunk
if chunk_input.dtype != grad_logits_t.dtype:
chunk_input = chunk_input.to(grad_logits_t.dtype)
torch.addmm(
grad_weight,
grad_logits_t,
chunk_input,
out_dtype=torch.float32,
out=grad_weight,
)
elif (
grad_weight.device.type == "cuda"
and torch.cuda.get_device_capability(grad_weight.device)[0] >= 8
and grad_logits_t.dtype in (torch.float16, torch.bfloat16)
and grad_weight.dtype == grad_logits_t.dtype
):
# Low-precision accumulator (accum_dtype=None with bf16/fp16 params) whose
# dtype already matches grad_logits. Accumulate straight into grad_weight
# with addmm(out=grad_weight), mirroring the fused_linear_cross_entropy
# accumulation; avoids the legacy path's parameter-sized low-precision
# temporary + cast per chunk.
chunk_input = input_chunk
if chunk_input.dtype != grad_logits_t.dtype:
chunk_input = chunk_input.to(grad_logits_t.dtype)
torch.addmm(
grad_weight,
grad_logits_t,
chunk_input,
out=grad_weight,
)
else:
# Legacy fallback: unsupported torch/device (no out_dtype), fp64
# accumulators, or a dtype mismatch (e.g. fp32 grad_logits promoted under
# AMP). Correct but allocates a parameter-sized fp32 temporary before
# summing into grad_weight.
grad_weight.add_(torch.mm(grad_logits_t, input_chunk).float())
del logits_chunk # free the chunk transient before the next chunk's GEMM
loss = loss_1d.sum()
grad_weight = (
grad_weight.to(student_weight.dtype) if grad_weight is not None and accum_dtype is not None else grad_weight
)
return loss, grad_input, grad_weight
def fused_linear_kl_div_backward(grad_output, grad_input, grad_weight):
# If KL is the last layer, grad_output is 1.0. Skip the mul to save time
if torch.ne(grad_output, torch.tensor(1.0, device=grad_output.device)):
# We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place
# for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton.
BT, H = grad_input.shape
n_rows = BT
BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(H))
element_mul_kernel[(n_rows,)](
grad_input,
grad_input.stride(-2),
grad_output,
H,
BLOCK_SIZE=BLOCK_SIZE,
num_warps=32 if not is_hip() else 16,
)
# handle grad_weight
if grad_weight is not None:
V, H = grad_weight.shape
n_rows = V
element_mul_kernel[(n_rows,)](
grad_weight,
grad_weight.stride(-2),
grad_output,
H,
BLOCK_SIZE=BLOCK_SIZE,
num_warps=32 if not is_hip() else 16,
)
return grad_input, grad_weight
class LigerFusedLinearKLDivFunction(torch.autograd.Function):
"""
Fusing the last linear layer with KL divergence against an explicit target distribution.
Handle the forward and backward pass of the final linear layer via KL by avoiding
the materialization of the large logits tensor. Since the KL loss is the last layer, we can
compute the gradient at the forward pass.
The Triton kernel consumes the native-precision GEMM output directly: the FP32
upcast, temperature scaling, log-softmax, the KL loss and the log-softmax
backward all happen inside the kernel on values in registers, and the logits
buffer is overwritten in place with the gradient w.r.t. the logits.
"""
@staticmethod
@amp_custom_fwd
def forward(
ctx,
student_input: torch.Tensor,
student_weight: torch.Tensor,
target: torch.Tensor,
shift_labels: Optional[torch.LongTensor] = None,
reduction: REDUCTION_LITERAL = "batchmean",
ignore_index: int = -100,
temperature: float = 1.0,
eps: float = 1e-10,
accum_dtype: Optional[torch.dtype] = None,
):
"""
Args:
student_input (torch.tensor): input of the last projection layer in the student model, with shape (BT, H), where B is batch size, T is sequence length, H is hidden dimension.
student_weight (torch.tensor): the last projection layer in the student model, with shape (V, H), where V is vocab size.
target (torch.tensor): target distribution with shape (BT, V), expected to be probabilities. Computes KL(target || softmax(student_input @ student_weight^T / temperature)).
shift_labels (Optional[torch.LongTensor]): indicator of next predicted vocab with shape (BT,). Rows whose label equals ignore_index are excluded from the loss and receive zero gradients. Default: `None`
reduction (REDUCTION_LITERAL, optional): reduction to be applied over the per-row KL losses: `"batchmean"`, `"mean"` or `"sum"` (same semantics as torch.nn.KLDivLoss). `"none"` is not supported because the gradient-in-forward scheme requires a scalar loss; use :class:`liger_kernel.transformers.kl_div.LigerKLDIVLoss` for element-wise losses. Default: `"batchmean"`
ignore_index (int, optional): the index to ignore. Default: `-100`
temperature (float, optional): temperature applied to the student logits before softmax. Default: `1.0`
eps (float, optional): small value clamping the target before log so entries with q = 0 contribute 0 (0 * log 0 = 0). Default: `1e-10`
accum_dtype (torch.dtype, optional): dtype of the weight gradient accumulator. Accumulating in `torch.float32` avoids precision loss from the chunked bf16/fp16 GEMM accumulation, at the cost of an extra fp32 weight-sized buffer. Default: `None` (same dtype as `student_weight`)
Returns:
loss (torch.Tensor): KL divergence, scaled by the reduction
"""
if reduction not in ("batchmean", "mean", "sum"):
raise ValueError(
f"reduction must be one of ('batchmean', 'mean', 'sum'), got {reduction!r}. "
"reduction='none' is not supported by the fused kernel; "
"use LigerKLDIVLoss for element-wise losses."
)
assert temperature != 0, "temperature cannot be 0."
if not target.is_contiguous():
target = target.contiguous()
has_label = False
if shift_labels is not None:
assert shift_labels.shape == (student_input.shape[0],), (
f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
)
shift_labels = shift_labels.contiguous()
has_label = True
loss, grad_input, grad_weight = fused_linear_kl_div_forward(
student_input,
student_weight,
target,
shift_labels,
reduction,
ignore_index,
has_label,
temperature,
eps,
accum_dtype,
)
# downcast to dtype and store for backward
ctx.save_for_backward(
grad_input.detach(),
grad_weight.detach() if grad_weight is not None else None,
)
return loss
@staticmethod
@amp_custom_bwd
def backward(ctx, grad_output):
(grad_input, grad_weight) = ctx.saved_tensors
grad_input, grad_weight = fused_linear_kl_div_backward(grad_output, grad_input, grad_weight)
return (grad_input, grad_weight, None, None, None, None, None, None, None)
src/liger_kernel/transformers/__init__.py
+2
View File
@@ -8,6 +8,7 @@ from liger_kernel.transformers.dyt import LigerDyT # noqa: F401
from liger_kernel.transformers.fused_add_rms_norm import LigerFusedAddRMSNorm # noqa: F401
from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss # noqa: F401
from liger_kernel.transformers.fused_linear_jsd import LigerFusedLinearJSD # noqa: F401
from liger_kernel.transformers.fused_linear_kl_div import LigerFusedLinearKLDivLoss # noqa: F401
from liger_kernel.transformers.geglu import LigerGEGLUMLP # noqa: F401
from liger_kernel.transformers.jsd import LigerJSD # noqa: F401
from liger_kernel.transformers.kl_div import LigerKLDIVLoss # noqa: F401
@@ -183,6 +184,7 @@ __all__ = [
"LigerDyT",
"LigerFusedLinearCrossEntropyLoss",
"LigerFusedLinearJSD",
"LigerFusedLinearKLDivLoss",
"LigerGEGLUMLP",
"LigerJSD",
"LigerLayerNorm",
src/liger_kernel/transformers/functional.py
+25
View File
@@ -9,6 +9,7 @@ from liger_kernel.ops import LigerDyTFunction
from liger_kernel.ops import LigerFusedAddRMSNormFunction
from liger_kernel.ops import LigerFusedLinearCrossEntropyFunction
from liger_kernel.ops import LigerFusedLinearJSDFunction
from liger_kernel.ops import LigerFusedLinearKLDivFunction
from liger_kernel.ops import LigerFusedLinearScaledCrossEntropyTPFunction
from liger_kernel.ops import LigerFusedNeighborhoodAttentionFunction
from liger_kernel.ops import LigerGELUMulFunction
@@ -174,6 +175,30 @@ def liger_fused_linear_jsd(
)
def liger_fused_linear_kl_div(
student_input,
student_weight,
target,
shift_labels=None,
reduction: str = "batchmean",
ignore_index: int = -100,
temperature: float = 1.0,
eps: float = 1e-10,
accum_dtype: Optional[torch.dtype] = None,
):
return LigerFusedLinearKLDivFunction.apply(
student_input,
student_weight,
target,
shift_labels,
reduction,
ignore_index,
temperature,
eps,
accum_dtype,
)
def liger_geglu(a, b):
return LigerGELUMulFunction.apply(a, b)
src/liger_kernel/transformers/fused_linear_kl_div.py
+87
View File
@@ -0,0 +1,87 @@
from typing import Optional
import torch
from liger_kernel.ops import LigerFusedLinearKLDivFunction
class LigerFusedLinearKLDivLoss(torch.nn.Module):
r"""Fusing the last linear layer with KL divergence against an explicit target distribution.
Handle the forward and backward pass of the final linear layer via KL by avoiding
the materialization of the large logits tensor. Unlike :class:`LigerFusedLinearJSD`,
the teacher distribution is given directly as a probability tensor, so no teacher
projection is fused: this is the fused-linear counterpart of
:class:`liger_kernel.transformers.kl_div.LigerKLDIVLoss` (and ``torch.nn.KLDivLoss``).
Useful when the teacher distribution is available as data, e.g. cached soft labels
from offline distillation datasets.
Args:
reduction (str): reduction to be applied over the per-row KL losses: `"batchmean"`, `"mean"` or `"sum"` (same semantics as torch.nn.KLDivLoss). `"none"` is not supported. Default: `"batchmean"`
ignore_index (int): The index to ignore in the target. Default: `-100`
temperature (float): temperature applied to the student logits before softmax. Default: `1.0`
eps (float): small value clamping the target before log so entries with q = 0 contribute 0 (0 * log 0 = 0). Default: `1e-10`
accum_dtype (torch.dtype): the dtype of intermediate result buffers for weight gradient accumulations. Default: `None`
Shape:
- student_input: :math:`(BT, H)`, where B is batch size, T is sequence length, H is hidden dimension.
- student_weight: :math:`(V, H)`, where V is vocab size.
- target: :math:`(BT, V)`, target probabilities.
- shift_labels: :math:`(BT,)`
- Output: a scalar.
Examples:
```python
>>> (B, T, H, V) = (2, 2, 3, 10)
>>> fused_kl = LigerFusedLinearKLDivLoss(temperature=2.0)
>>> # generate inputs and weights
>>> student_input = torch.rand(B * T, H, device="cuda", requires_grad=True)
>>> student_lin = torch.nn.Linear(H, V, bias=False, device="cuda")
>>> # target distribution given directly, e.g. cached teacher probabilities
>>> target = torch.rand(B * T, V, device="cuda").softmax(dim=-1)
>>> output = fused_kl(student_input, student_lin.weight, target)
>>> output.backward()
```
"""
def __init__(
self,
reduction: Optional[str] = "batchmean",
ignore_index: int = -100,
temperature: float = 1.0,
eps: float = 1e-10,
accum_dtype: Optional[torch.dtype] = None,
):
super().__init__()
assert temperature != 0, "temperature cannot be 0."
self.reduction = reduction
self.temperature = temperature
self.ignore_index = ignore_index
self.eps = eps
self.accum_dtype = accum_dtype
def forward(
self,
student_input: torch.Tensor,
student_weight: torch.Tensor,
target: torch.Tensor,
shift_labels: Optional[torch.LongTensor] = None,
):
return LigerFusedLinearKLDivFunction.apply(
student_input,
student_weight,
target,
shift_labels,
self.reduction,
self.ignore_index,
self.temperature,
self.eps,
self.accum_dtype,
)
def extra_repr(self):
return (
f"reduction={self.reduction}, ignore_index={self.ignore_index}, "
f"temperature={self.temperature}, eps={self.eps}, accum_dtype={self.accum_dtype}"
)
test/transformers/test_fused_linear_kl_div.py
+568
View File
@@ -0,0 +1,568 @@
import pytest
import torch
from test.utils import assert_verbose_allclose
from test.utils import set_seed
from test.utils import supports_bfloat16
from liger_kernel.ops import LigerFusedLinearKLDivFunction
from liger_kernel.transformers.functional import liger_fused_linear_kl_div
from liger_kernel.transformers.fused_linear_kl_div import LigerFusedLinearKLDivLoss
from liger_kernel.utils import infer_device
device = infer_device()
set_seed(42)
class TorchLMHeadKLDiv(torch.nn.Module):
"""Ground truth implementation of the linear fused with torch based KL divergence loss.
:param H: hidden size
:param V: vocab size
:param temperature: softmax temperature
:param reduction: reduction over the per-row KL losses
"""
def __init__(
self,
H: int,
V: int,
dtype: torch.dtype,
device: torch.device,
reduction: str = "batchmean",
ignore_index: int = -100,
temperature: float = 1.0,
eps: float = 1e-10,
):
super().__init__()
self.student_lin = torch.nn.Linear(in_features=H, out_features=V, bias=False, dtype=dtype, device=device)
self.reduction = reduction
self.ignore_index = ignore_index
self.temperature = temperature
self.eps = eps
def forward(self, student_input, target, label=None):
logits = self.student_lin(student_input).to(torch.float32) / self.temperature
log_p = torch.log_softmax(logits, dim=-1)
q = target.to(torch.float32)
# KL(q || p) = sum(q * (log q - log p)); 0 * log 0 is treated as 0 via the clamp
loss_mat = q * (torch.log(q.clamp_min(self.eps)) - log_p)
if label is not None:
keep = (label != self.ignore_index).to(torch.float32).unsqueeze(-1)
loss_mat = loss_mat * keep
n_non_ignore = int(keep.sum().item())
else:
n_non_ignore = student_input.shape[0]
total = loss_mat.sum()
if self.reduction == "sum":
return total
if n_non_ignore == 0:
return total * 0.0
if self.reduction == "batchmean":
return total / n_non_ignore
return total / (n_non_ignore * self.student_lin.out_features) # "mean"
class LigerLMHeadKLDiv(torch.nn.Module):
def __init__(
self,
H: int,
V: int,
dtype: torch.dtype,
device: torch.device,
reduction: str = "batchmean",
ignore_index: int = -100,
temperature: float = 1.0,
eps: float = 1e-10,
accum_dtype=None,
):
super().__init__()
self.student_lin = torch.nn.Linear(in_features=H, out_features=V, bias=False, dtype=dtype, device=device)
self.fused_kl = LigerFusedLinearKLDivLoss(
reduction=reduction,
ignore_index=ignore_index,
temperature=temperature,
eps=eps,
accum_dtype=accum_dtype,
)
def forward(self, student_input, target, label=None):
return self.fused_kl(student_input, self.student_lin.weight, target, label)
#############################################################################
# Test the correctness of the fused linear KL divergence
#############################################################################
@pytest.mark.parametrize(
"B, T, H, V",
[
(8, 128, 1024, 4096),
(4, 423, 167, 1423), # random shape
],
)
@pytest.mark.parametrize(
"temperature, reduction, accum_dtype, scalar, dtype, atol, rtol",
[
# with reduction="sum" the gradients are O(BT) larger than with "batchmean",
# so the absolute bf16 rounding noise of the chunked GEMMs scales up accordingly
pytest.param(
1.0,
"batchmean",
None,
1.0,
torch.bfloat16,
5e-3,
5e-2,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
(1.0, "batchmean", None, 1.0, torch.float32, 1e-5, 5e-4),
(1.0, "batchmean", None, 1.0, torch.float16, 5e-3, 5e-2),
pytest.param(
2.0,
"sum",
None,
1.0,
torch.bfloat16,
5e-0,
5e1,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
# pass non-default accum_dtype once to ensure it works along
(2.0, "sum", torch.float32, 1.0, torch.float32, 1e-3, 5e-2),
(2.0, "sum", None, 1.0, torch.float16, 5e-3, 5e-2),
pytest.param(
1.0,
"mean",
None,
1.0,
torch.bfloat16,
5e-3,
5e-2,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
(1.0, "mean", None, 1.0, torch.float32, 1e-5, 5e-4),
(1.0, "mean", None, 1.0, torch.float16, 5e-3, 5e-2),
],
)
def test_correctness(B, T, H, V, scalar, dtype, reduction, temperature, accum_dtype, atol, rtol):
torch_lm_head_kl = TorchLMHeadKLDiv(
H=H,
V=V,
dtype=dtype,
device=device,
temperature=temperature,
reduction=reduction,
).to(device)
liger_lm_head_kl = LigerLMHeadKLDiv(
H=H,
V=V,
dtype=dtype,
device=device,
temperature=temperature,
reduction=reduction,
accum_dtype=accum_dtype,
).to(device)
# init the linear in all FusedLinearKLDivs with the same weights
torch_lm_head_kl.student_lin.weight.data = liger_lm_head_kl.student_lin.weight.data = torch.rand(
V, H, device=device, dtype=dtype
)
_tensor = torch.randn(B * T, H, device=device, dtype=dtype) * scalar
_input1 = _tensor.detach().clone().requires_grad_(True)
_input2 = _tensor.detach().clone().requires_grad_(True)
target = torch.rand(B * T, V, device=device, dtype=torch.float32).softmax(dim=-1).to(dtype)
with torch.autograd.detect_anomaly():
output1 = torch_lm_head_kl(_input1, target)
output2 = liger_lm_head_kl(_input2, target)
assert_verbose_allclose(output1, output2, atol=atol, rtol=rtol)
output1.backward()
output2.backward()
assert_verbose_allclose(_input1.grad, _input2.grad, atol=atol, rtol=rtol)
assert_verbose_allclose(
torch_lm_head_kl.student_lin.weight.grad,
liger_lm_head_kl.student_lin.weight.grad,
atol=atol,
rtol=rtol,
)
@pytest.mark.parametrize(
"B, T, H, V",
[
(8, 128, 1024, 4096),
(4, 423, 167, 1423), # random shape
],
)
@pytest.mark.parametrize(
"temperature, reduction, ignore_index, scalar, dtype, atol, rtol",
[
pytest.param(
1.0,
"batchmean",
2,
1.0,
torch.bfloat16,
5e-3,
5e-2,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
(1.0, "batchmean", 2, 1.0, torch.float32, 1e-5, 5e-4),
(1.0, "batchmean", 2, 1.0, torch.float16, 5e-3, 5e-2),
# see the comment in test_correctness for the relaxed bf16 tolerance of reduction="sum"
pytest.param(
2.0,
"sum",
42,
1.0,
torch.bfloat16,
5e-0,
5e1,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
(2.0, "sum", 42, 1.0, torch.float32, 1e-3, 5e-2),
(2.0, "sum", 42, 1.0, torch.float16, 5e-3, 5e-2),
pytest.param(
1.0,
"mean",
-100,
1.0,
torch.bfloat16,
5e-3,
5e-2,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
(1.0, "mean", -100, 1.0, torch.float32, 1e-5, 5e-4),
(1.0, "mean", -100, 1.0, torch.float16, 5e-3, 5e-2),
],
)
def test_correctness_with_ignore_index(B, T, H, V, scalar, dtype, reduction, ignore_index, temperature, atol, rtol):
torch_lm_head_kl = TorchLMHeadKLDiv(
H=H,
V=V,
dtype=dtype,
device=device,
temperature=temperature,
ignore_index=ignore_index,
reduction=reduction,
).to(device)
liger_lm_head_kl = LigerLMHeadKLDiv(
H=H,
V=V,
dtype=dtype,
device=device,
temperature=temperature,
ignore_index=ignore_index,
reduction=reduction,
).to(device)
# init the linear in all FusedLinearKLDivs with the same weights
torch_lm_head_kl.student_lin.weight.data = liger_lm_head_kl.student_lin.weight.data = torch.rand(
V, H, device=device, dtype=dtype
)
_tensor = torch.randn(B * T, H, device=device, dtype=dtype) * scalar
_input1 = _tensor.detach().clone().requires_grad_(True)
_input2 = _tensor.detach().clone().requires_grad_(True)
target = torch.rand(B * T, V, device=device, dtype=torch.float32).softmax(dim=-1).to(dtype)
label = torch.randint(0, V, (B * T,), device=device, dtype=torch.long)
# Assign some random number of elements as ignore_index
num_elements_to_assign = torch.randint(
1, B * T // 2, (1,)
).item() # Random number of elements to set to ignore_index
indices_to_assign = torch.randperm(B * T)[:num_elements_to_assign] # Randomly select indices
label[indices_to_assign] = ignore_index
output1 = torch_lm_head_kl(_input1, target, label)
output2 = liger_lm_head_kl(_input2, target, label)
assert_verbose_allclose(output1, output2, atol=atol, rtol=rtol)
output1.backward()
output2.backward()
assert_verbose_allclose(_input1.grad, _input2.grad, atol=atol, rtol=rtol)
assert_verbose_allclose(
torch_lm_head_kl.student_lin.weight.grad,
liger_lm_head_kl.student_lin.weight.grad,
atol=atol,
rtol=rtol,
)
@pytest.mark.parametrize(
"B, T, H, V",
[
(2, 2, 8, 8),
# weird shapes
(9, 7, 41, 41),
],
)
@pytest.mark.parametrize(
"scalar, dtype, atol, rtol",
[
pytest.param(
0.5,
torch.bfloat16,
5e-3,
5e-2,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
(0.5, torch.float32, 1e-5, 5e-4),
(0.5, torch.float16, 5e-3, 5e-2),
],
)
@pytest.mark.parametrize("temperature, reduction, ignore_index", [(1.0, "batchmean", -100), (2.0, "sum", 42)])
@pytest.mark.parametrize("accum_dtype", [None, torch.float32])
def test_correctness_functional(
B, T, H, V, scalar, dtype, reduction, ignore_index, temperature, accum_dtype, atol, rtol
):
# init the linear in all FusedLinearKLDivs with the same weights
_weight = torch.rand(V, H, device=device, dtype=dtype)
_weight1 = _weight.detach().clone().requires_grad_(True)
_weight2 = _weight.detach().clone().requires_grad_(True)
_tensor = torch.randn(B * T, H, device=device, dtype=dtype) * scalar
_input1 = _tensor.detach().clone().requires_grad_(True)
_input2 = _tensor.detach().clone().requires_grad_(True)
target = torch.rand(B * T, V, device=device, dtype=torch.float32).softmax(dim=-1).to(dtype)
label = torch.randint(0, V, (B * T,), device=device, dtype=torch.long)
# Assign some random number of elements as ignore_index
num_elements_to_assign = torch.randint(
1, B * T // 2, (1,)
).item() # Random number of elements to set to ignore_index
indices_to_assign = torch.randperm(B * T)[:num_elements_to_assign] # Randomly select indices
label[indices_to_assign] = ignore_index
output1 = liger_fused_linear_kl_div(
student_input=_input1,
student_weight=_weight1,
target=target,
shift_labels=label,
reduction=reduction,
ignore_index=ignore_index,
temperature=temperature,
accum_dtype=accum_dtype,
)
output2 = LigerFusedLinearKLDivFunction.apply(
_input2,
_weight2,
target,
label,
reduction,
ignore_index,
temperature,
1e-10,
accum_dtype,
)
assert_verbose_allclose(output1, output2, atol=atol, rtol=rtol)
output1.backward()
output2.backward()
assert_verbose_allclose(_input1.grad, _input2.grad, atol=atol, rtol=rtol)
assert_verbose_allclose(_weight1.grad, _weight2.grad, atol=atol, rtol=rtol)
@pytest.mark.parametrize(
"B, T, H, V",
[
(8, 128, 1024, 4096),
(4, 423, 167, 1423), # random shape
],
)
@pytest.mark.parametrize(
"scalar, dtype, atol, rtol",
[
pytest.param(
1.0,
torch.bfloat16,
5e-3,
5e-2,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
(1.0, torch.float32, 1e-5, 5e-4),
(1.0, torch.float16, 5e-3, 5e-2),
],
)
@pytest.mark.parametrize(
"temperature, reduction, ignore_index",
[
(1.0, "batchmean", 2),
(2.0, "sum", 42),
],
)
def test_correctness_all_ignored(B, T, H, V, scalar, dtype, reduction, ignore_index, temperature, atol, rtol):
torch_lm_head_kl = TorchLMHeadKLDiv(
H=H,
V=V,
dtype=dtype,
device=device,
temperature=temperature,
ignore_index=ignore_index,
reduction=reduction,
).to(device)
liger_lm_head_kl = LigerLMHeadKLDiv(
H=H,
V=V,
dtype=dtype,
device=device,
temperature=temperature,
ignore_index=ignore_index,
reduction=reduction,
).to(device)
# init the linear in all FusedLinearKLDivs with the same weights
torch_lm_head_kl.student_lin.weight.data = liger_lm_head_kl.student_lin.weight.data = torch.rand(
V, H, device=device, dtype=dtype
)
_tensor = torch.randn(B * T, H, device=device, dtype=dtype) * scalar
_input1 = _tensor.detach().clone().requires_grad_(True)
_input2 = _tensor.detach().clone().requires_grad_(True)
target = torch.rand(B * T, V, device=device, dtype=torch.float32).softmax(dim=-1).to(dtype)
label = torch.full((B * T,), ignore_index, device=device, dtype=torch.long)
output1 = torch_lm_head_kl(_input1, target, label)
output2 = liger_lm_head_kl(_input2, target, label)
assert_verbose_allclose(output1, output2, atol=atol, rtol=rtol)
assert_verbose_allclose(output2, torch.zeros_like(output2), atol=atol, rtol=rtol)
output2.backward()
assert_verbose_allclose(torch.zeros_like(_input2.grad), _input2.grad, atol=atol, rtol=rtol)
@pytest.mark.parametrize(
"autocast_dtype, atol, rtol",
[
pytest.param(
torch.bfloat16,
5e-3,
5e-2,
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
),
(torch.float16, 5e-3, 5e-2),
],
)
def test_amp(autocast_dtype, atol, rtol):
B = 2
T = 4
H = 2048
V = 3200
scalar = 1.0
ignore_index = -100
temperature = 1.0
reduction = "batchmean"
dtype = torch.float32
torch_lm_head_kl = TorchLMHeadKLDiv(
H=H,
V=V,
dtype=dtype,
device=device,
temperature=temperature,
ignore_index=ignore_index,
reduction=reduction,
).to(device)
liger_lm_head_kl = LigerLMHeadKLDiv(
H=H,
V=V,
dtype=dtype,
device=device,
temperature=temperature,
ignore_index=ignore_index,
reduction=reduction,
).to(device)
# init the linear in all FusedLinearKLDivs with the same weights
torch_lm_head_kl.student_lin.weight.data = liger_lm_head_kl.student_lin.weight.data = torch.rand(
V, H, device=device, dtype=dtype
)
_tensor = torch.randn(B * T, H, device=device, dtype=autocast_dtype) * scalar
_input1 = _tensor.detach().clone().requires_grad_(True)
_input2 = _tensor.detach().clone().requires_grad_(True)
target = torch.rand(B * T, V, device=device, dtype=torch.float32).softmax(dim=-1).to(autocast_dtype)
label = torch.randint(0, V, (B * T,), device=device, dtype=torch.long)
# Assign some random number of elements as ignore_index
num_elements_to_assign = torch.randint(
1, B * T // 2, (1,)
).item() # Random number of elements to set to ignore_index
indices_to_assign = torch.randperm(B * T)[:num_elements_to_assign] # Randomly select indices
label[indices_to_assign] = ignore_index
with torch.autocast(device_type=device, dtype=autocast_dtype):
output1 = torch_lm_head_kl(_input1, target, label)
output2 = liger_lm_head_kl(_input2, target, label)
assert_verbose_allclose(output1, output2, atol=atol, rtol=rtol)
output1.backward()
output2.backward()
assert_verbose_allclose(_input1.grad, _input2.grad, atol=atol, rtol=rtol)
assert_verbose_allclose(
torch_lm_head_kl.student_lin.weight.grad,
liger_lm_head_kl.student_lin.weight.grad,
atol=atol,
rtol=rtol,
)
def test_torch_kl_div_parity():
# cross-check the fused kernel against torch.nn.functional.kl_div in fp32
B = 2
T = 8
H = 128
V = 1024
torch.manual_seed(0)
_input = torch.randn(B * T, H, device=device, dtype=torch.float32, requires_grad=True)
_weight = torch.randn(V, H, device=device, dtype=torch.float32, requires_grad=True)
target = torch.rand(B * T, V, device=device, dtype=torch.float32).softmax(dim=-1)
expected = torch.nn.functional.kl_div(
torch.log_softmax(_input @ _weight.t(), dim=-1),
target,
reduction="batchmean",
)
got = liger_fused_linear_kl_div(_input, _weight, target)
assert_verbose_allclose(expected, got, atol=1e-6, rtol=1e-5)
def test_invalid_reduction():
B = 1
T = 4
H = 8
V = 16
_input = torch.randn(B * T, H, device=device, dtype=torch.float32, requires_grad=True)
_weight = torch.randn(V, H, device=device, dtype=torch.float32)
target = torch.rand(B * T, V, device=device, dtype=torch.float32).softmax(dim=-1)
with pytest.raises(ValueError, match="reduction must be one of"):
liger_fused_linear_kl_div(_input, _weight, target, reduction="none")