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

perf(flce): reuse logits buffers and write input gradients directly (#1468)

Browse Source 
## Summary

Reduce peak allocated memory in fused linear cross entropy by reusing
chunk-logits storage and writing input gradients directly into their
destination.

On B300 with BF16/FP16, the reference 8,192-token workload uses **33.3%
less incremental peak memory**; the 65,536-token workload uses **46.5%
less**. Latency is effectively unchanged. This is a buffer-management
optimization, not a new kernel or chunk-size policy.

## Details

- Allocate one reusable logits buffer on eligible NVIDIA CUDA paths, and
use `torch.mm(..., out=...)` to fill each chunk.
- Write dX directly into its preallocated slice instead of allocating a
GEMM result and copying it.
- Release old logits, gradient/transpose views, CE argument tuples, and
token-scaling temporaries once consumed. Release reusable workspace
before the final gradient-dtype conversion.
- Preserve the original allocating matmul and bias-promotion behavior
when direct output would change dtype or autocast semantics.

The direct-output path requires no-grad execution, matching
FP16/BF16/FP32 operands, and disabled or dtype-matching CUDA autocast.
Logits reuse additionally requires absent/same-dtype bias; direct dX
requires a matching destination dtype. Other cases retain the original
matmul behavior. The older autocast-query API remains supported.

## B300 measurements

Full forward + backward, mean reduction, no bias, default accumulation,
H=4096 and V=128256. Memory is incremental peak allocated MiB above the
resident input/parameter baseline, not reserved or total process memory.

| Dtype | Tokens | Before ms | After ms | Before peak MiB | After peak
MiB | Memory reduction |
|---|---:|---:|---:|---:|---:|---:|
| BF16 | 8,192 | 19.193 | 19.067 | 3,102.04 | 2,068.04 | 33.3% |
| BF16 | 16,384 | 38.301 | 38.420 | 5,202.08 | 3,134.08 | 39.8% |
| BF16 | 65,536 | 146.720 | 148.472 | 17,802.31 | 9,530.31 | 46.5% |
| FP16 | 8,192 | 16.120 | 16.164 | 3,102.04 | 2,068.04 | 33.3% |
| FP16 | 16,384 | 31.688 | 31.556 | 5,202.08 | 3,134.08 | 39.8% |
| FP16 | 65,536 | 124.956 | 123.641 | 17,802.31 | 9,530.31 | 46.5% |

Across 23 cases covering BF16, FP16, FP32, FP32 accumulation, AMP, and
optional features, observed latency changes ranged from **2.2% faster to
1.4% slower**. No material speedup is claimed. The measured losses and
gradients matched the baseline with zero numerical difference in these
cases.

Methodology:

- NVIDIA B300 SM103; PyTorch 2.11.0.1+cu130; Triton 3.6.0.
- Baseline: FLCE source at `f2ac965f619aa4aa090c34252f87349a968b7597`,
unchanged by the two subsequent main commits used as this PR's base.
- Baseline and candidate loaded in the same process; identical seeded
tensors, chunk sizes, precision, and upstream gradient 1.0.
- Fresh forward/backward graph per iteration; both leaf gradients
cleared. No optimizer step.
- Compilation excluded; eager CUDA-event `triton.testing.do_bench`, 100
ms warmup, 400 ms repetition budget, four alternating-order rounds;
median of round medians.
- Memory: median of three separate warmed peak-allocation probes.

The improvement removes approximately one overlapping logits slab plus a
temporary dX chunk. It does **not** make memory independent of token
count: the unchanged C=16 policy still grows the chunk size with total
tokens.

## Hopper and other GPUs

There are no new SM100/SM103-specific instructions or kernels. The
optimization uses ordinary PyTorch/cuBLAS output-buffer operations, so
the same allocation-lifetime benefit is expected on eligible Hopper and
other NVIDIA CUDA workloads.

**Only B300 has been measured.** H100/H200/Ampere latency, exact peak
allocations, and real older-PyTorch behavior still need hardware
coverage; identical performance percentages are not claimed. ROCm and
non-CUDA devices do not enter the new direct-output path. The HIP
eligibility simulation below is not a ROCm hardware run.

## Testing Done

- **322 passed:** transformer FLCE suite, ops FLCE suite, and CuTe
shape-fallback contracts.
- **170 passed:** ops FLCE suite with `LIGER_KERNEL_STRICT=1` (overlaps
the preceding coverage).
- **58 passed / 19 NVIDIA-only skips:** targeted tests with only the HIP
eligibility flag simulated on B300.
- Regression coverage includes actual aten dispatch, storage lifetime
through views, ragged/noncontiguous chunks, AMP and bias promotion, FP32
accumulation, token scaling, forward-only execution, and deferred
`reduction="none"` backward.
- Ruff lint, format checks, and `git diff --check` pass.
- Full-repository `make test` and convergence suites were not run.

Only the FLCE implementation and its existing regression-test file are
included. Local dependency pins, benchmark artifacts, and unrelated
files are excluded.

Co-authored-by: Copilot <[email protected]>
This commit is contained in:
Vaibhav Jindal
2026-09-15 23:50:07 +00:00
committed by GitHub
co-authored by Copilot
parent dc1512630d
commit de77856a01
2 changed files with 390 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/liger_kernel/ops/fused_linear_cross_entropy.py
+61 -4
View File
@@ -40,6 +40,25 @@ _TORCH_VERSION = Version(torch.__version__.split("+")[0])
_ADDMM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0")
def _can_use_mm_out(_input, weight):
# out= bypasses autocast. Only use it when @ would use these exact operands
# and return their dtype; keep AMP casts/rounding on the original path.
if (
_input.device.type != "cuda"
or is_hip()
or torch.is_grad_enabled()
or _input.dtype != weight.dtype
or _input.dtype not in (torch.float16, torch.bfloat16, torch.float32)
):
return False
if torch.is_autocast_enabled():
autocast_dtype = (
torch.get_autocast_dtype("cuda") if _TORCH_VERSION >= Version("2.4.0") else torch.get_autocast_gpu_dtype()
)
return _input.dtype == autocast_dtype
return True
def fused_linear_cross_entropy_forward(
_input,
weight,
@@ -133,15 +152,30 @@ def fused_linear_cross_entropy_forward(
if ce_weight.stride(-1) != 1:
ce_weight = ce_weight.contiguous()
use_mm_out = _can_use_mm_out(_input, weight)
# A promoted bias needs the original out-of-place add. Do not reserve an
# unused low-precision buffer alongside those promoted logits.
logits_storage = (
torch.empty((chunk_size, V), dtype=_input.dtype, device=device)
if use_mm_out and (bias is None or bias.dtype == _input.dtype)
else None
)
for chunk_id in range(num_chunks):
start_idx = chunk_id * chunk_size
end_idx = min((chunk_id + 1) * chunk_size, BT)
_input_chunk = _input[start_idx:end_idx] # chunk_size x H
# when doing matmul, use the original precision
logits_chunk = _input_chunk @ weight.t() # chunk_size x V
if bias is not None:
logits_chunk = logits_chunk + bias
if logits_storage is not None:
logits_chunk = logits_storage[: end_idx - start_idx]
torch.mm(_input_chunk, weight.t(), out=logits_chunk)
if bias is not None:
logits_chunk.add_(bias)
else:
logits_chunk = _input_chunk @ weight.t() # chunk_size x V
if bias is not None:
logits_chunk = logits_chunk + bias
target_chunk = target[start_idx:end_idx] # chunk_size,
@@ -155,6 +189,7 @@ def fused_linear_cross_entropy_forward(
# Compute softmax to get predicted probabilities
probs = torch.softmax(logits_for_softmax, dim=-1)
del logits_for_softmax
# Get predicted probabilities for token scaling, handling ignored targets
valid_target_mask = target_chunk != ignore_index
@@ -164,16 +199,19 @@ def fused_linear_cross_entropy_forward(
# Gather probabilities only for valid targets
valid_probs = probs[valid_target_mask]
pred_probs_valid = torch.gather(valid_probs, -1, valid_targets.unsqueeze(-1)).squeeze(-1)
del valid_probs
# Create full tensor with zeros for ignored targets
pred_probs = torch.zeros_like(target_chunk, dtype=probs.dtype, device=probs.device)
pred_probs[valid_target_mask] = pred_probs_valid
del pred_probs_valid
else:
# All targets are ignored
pred_probs = torch.zeros_like(target_chunk, dtype=probs.dtype, device=probs.device)
# Store the scaling factors
scaling_factors = pred_probs.detach() # Detach to ensure no gradient flow
del probs, pred_probs, valid_target_mask, valid_targets
# ensure _input and target are contiguous
logits_chunk = logits_chunk.contiguous()
@@ -248,6 +286,11 @@ def fused_linear_cross_entropy_forward(
impl=ce_impl,
mode=ce_mode,
)
del ce_args
# CE may return an alias of logits. Keep only the gradient reference so
# replacing it with scaled gradients also releases fallback logits.
del logits_chunk
# Apply token scaling if requested
if use_token_scaling:
@@ -268,6 +311,7 @@ def fused_linear_cross_entropy_forward(
# Expand scaling factors to match gradient dimensions
scaling_factors_expanded = scaling_factors.unsqueeze(-1) # chunk_size x 1
grad_logits_chunk = grad_logits_chunk * scaling_factors_expanded
del scaling_factors, scaling_factors_expanded
# reduction="none": fold the per-token upstream gradient in HERE, before the
# projections below sum over the token dimension. grad_weight is
@@ -280,7 +324,10 @@ def fused_linear_cross_entropy_forward(
)
if input_requires_grad:
grad_input[start_idx:end_idx] = grad_logits_chunk @ weight
if use_mm_out and grad_logits_chunk.dtype == grad_input.dtype:
torch.mm(grad_logits_chunk, weight, out=grad_input[start_idx:end_idx])
else:
grad_input[start_idx:end_idx] = grad_logits_chunk @ weight
if grad_weight is not None and input_requires_grad:
grad_logits_t = grad_logits_chunk.t()
@@ -306,6 +353,7 @@ def fused_linear_cross_entropy_forward(
out_dtype=torch.float32,
out=grad_weight,
)
del input_chunk
elif (
grad_weight.device.type == "cuda"
and torch.cuda.get_device_capability(grad_weight.device)[0] >= 8
@@ -327,11 +375,13 @@ def fused_linear_cross_entropy_forward(
input_chunk,
out=grad_weight,
)
del input_chunk
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 += torch.mm(grad_logits_chunk.t(), _input_chunk).float()
del grad_logits_t
if bias is not None and input_requires_grad:
torch.add(
@@ -341,6 +391,13 @@ def fused_linear_cross_entropy_forward(
alpha=1.0,
)
# In particular, no transposed/scaled gradient or CE argument tuple may
# keep the previous chunk alive when the fallback allocates new logits.
del grad_logits_chunk, _input_chunk, target_chunk
# Release even the persistent buffer before allocating a dtype-converted dW.
del logits_storage
# Need extra calculations for backward if reduction=='none'. Not supporting reduction='none' now.
# if reduction == "none":
# loss = loss_1d
test/ops/test_fused_linear_cross_entropy.py
+329 -6
View File
@@ -19,6 +19,7 @@ import torch
import torch.nn.functional as F
import triton
from torch.multiprocessing.reductions import StorageWeakRef
from torch.utils._python_dispatch import TorchDispatchMode
import liger_kernel
@@ -317,6 +318,11 @@ requires_addmm_out_dtype = pytest.mark.skipif(
reason="addmm out_dtype overload (fp32 accumulator fast path) requires torch >= 2.8",
)
requires_nvidia_cuda = pytest.mark.skipif(
not torch.cuda.is_available() or flce_ops.is_hip(),
reason="reusable mm(out=) buffers require NVIDIA CUDA",
)
def _production_chunk_size(BT: int, V: int, H: int) -> int:
"""The chunk size the production loop will pick for these dims, using the
@@ -337,13 +343,17 @@ _ADDMM_OPS = ("aten.addmm.out", "aten.addmm.dtype_out")
class _OpRecorder(TorchDispatchMode):
"""Records the aten ops we care about: addmm (with the out-tensor shape and
whether the out_dtype overload was used) and ``_to_copy`` (with in/out dtype
and shape so we can detect parameter-sized upcasts)."""
whether the out_dtype overload was used), mm outputs/copies, and dtype casts.
Store metadata only, so the recorder cannot extend a buffer's lifetime."""
def __init__(self):
self.addmm_out_shapes = [] # list of torch.Size for aten.addmm.out
self.addmm_dtype_out_shapes = [] # list of torch.Size for aten.addmm.dtype_out
self.mm_default = 0 # count of aten.mm.default
self.mm_default_shapes = []
self.mm_out = [] # (shape, storage pointer, data pointer)
self.copies = [] # (destination shape, destination storage pointer)
self.empty_shapes = []
self.to_copies = [] # list of (in_dtype, out_dtype, shape)
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
@@ -358,6 +368,15 @@ class _OpRecorder(TorchDispatchMode):
self.addmm_dtype_out_shapes.append(shape)
elif name == "aten.mm.default":
self.mm_default += 1
self.mm_default_shapes.append((args[0].shape[0], args[1].shape[1]))
elif name == "aten.mm.out":
out = kwargs["out"]
self.mm_out.append((tuple(out.shape), out.untyped_storage().data_ptr(), out.data_ptr()))
elif name == "aten.copy_.default":
dst = args[0]
self.copies.append((tuple(dst.shape), dst.untyped_storage().data_ptr()))
elif name == "aten.empty.memory_format":
self.empty_shapes.append(tuple(args[0]))
elif name == "aten._to_copy.default":
src = args[0] if args else None
if isinstance(src, torch.Tensor):
@@ -383,6 +402,55 @@ class _OpRecorder(TorchDispatchMode):
return any(src in low_dtypes and dst == torch.float32 and sh == tgt for (src, dst, sh) in self.to_copies)
class _BufferRecorder(_OpRecorder):
"""Check storage lifetime, including views that outlive their base tensor.
StorageWeakRef observes the C++ storage without keeping it alive. Python
tensor weakrefs alone would miss aliases such as grad_logits_t and ce_args.
"""
def __init__(self, BT, V, H):
super().__init__()
self.V, self.H = V, H
self.chunk_size = _production_chunk_size(BT, V, H)
self.chunk_storages = []
self.logits_gemms = 0
self.final_casts = 0
def assert_released(self, reusable_ptr=None):
assert all(ref.expired() or ptr == reusable_ptr for ptr, ref in self.chunk_storages), (
"a logits/probability/scaled-gradient storage is still live"
)
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
kwargs = kwargs or {}
name = str(func)
if name in ("aten.mm.default", "aten.mm.out") and tuple(args[1].shape) == (self.H, self.V):
out = kwargs.get("out")
reusable_ptr = out.untyped_storage().data_ptr() if out is not None else None
self.assert_released(reusable_ptr)
self.logits_gemms += 1
if (
name == "aten._to_copy.default"
and tuple(args[0].shape) == (self.V, self.H)
and args[0].dtype == torch.float32
and kwargs.get("dtype") in (torch.bfloat16, torch.float16)
):
self.assert_released()
self.final_casts += 1
result = super().__torch_dispatch__(func, types, args, kwargs)
if (
isinstance(result, torch.Tensor)
and result.ndim == 2
and result.shape[1] == self.V
and 0 < result.shape[0] <= self.chunk_size
):
storage = result.untyped_storage()
self.chunk_storages.append((storage.data_ptr(), StorageWeakRef(storage)))
return result
# ---------------------------------------------------------------------------
# Reference + driver helpers
# ---------------------------------------------------------------------------
@@ -447,6 +515,7 @@ def _run_flce(
token_grad_output=None,
recorder: Optional[_OpRecorder] = None,
autocast_dtype=None,
ce_impl=None,
):
"""Invoke the forward the same way the autograd Function does: under
``no_grad`` (Function.forward runs with grad disabled), optionally under
@@ -466,6 +535,7 @@ def _run_flce(
accum_dtype=accum_dtype,
use_token_scaling=use_token_scaling,
token_grad_output=token_grad_output,
ce_impl=ce_impl,
)
import contextlib
@@ -487,6 +557,8 @@ _ADDMM_TOLS = {
torch.float32: dict(atol=1e-4, rtol=1e-4),
}
_SHARED_CE_IMPLS = [None] + [b for b in _REGISTERED_BACKENDS if b in ("nvidia-triton", "nvidia-cutedsl")]
def _mk_inputs(
BT,
@@ -525,7 +597,7 @@ def _mk_inputs(
class TestFusedLinearCrossEntropyAddmm:
"""Regression suite for the low-precision ``grad_weight`` addmm path.
"""Regression suite for chunk buffers and the ``grad_weight`` addmm path.
Grouped in a class so the ``_force_multichunk`` autouse fixture (which pins
``_CHUNK_MEM_CONST=1``) is scoped to *only* these migrated tests and never
@@ -763,7 +835,7 @@ class TestFusedLinearCrossEntropyAddmm:
if accum_dtype is torch.float32:
# fp32 accumulator falls back to legacy torch.mm(...).float(): mm + a
# parameter-sized fp32 upcast, and no addmm at all into the weight buffer.
assert rec.mm_default >= 1
assert rec.mm_default_shapes.count(tuple(weight.shape)) == triton.cdiv(BT, _production_chunk_size(BT, V, H))
assert rec.has_param_upcast(weight.shape)
assert not rec.has_addmm_out_into(weight.shape)
else:
@@ -775,6 +847,257 @@ class TestFusedLinearCrossEntropyAddmm:
torch.testing.assert_close(gw.float(), ref_gw, **_ADDMM_TOLS[torch.bfloat16])
torch.testing.assert_close(gi.float(), ref_gi, **_ADDMM_TOLS[torch.bfloat16])
# -----------------------------------------------------------------------
# Logits storage reuse and direct dX output, including ragged/strided views
# -----------------------------------------------------------------------
@requires_nvidia_cuda
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@pytest.mark.parametrize("bias", [False, True])
@pytest.mark.parametrize("layout", ["contiguous", "transposed", "strided"])
def test_reuses_logits_and_writes_dx_out(self, monkeypatch, dtype, bias, layout):
BT, V, H = 40, 128, 32
_input, weight, b, _, target = _mk_inputs(BT, V, H, dtype, bias=bias, ce_weight=False, seed=15)
if layout == "transposed":
_input = _input.t().contiguous().t().detach().requires_grad_(True)
weight = weight.t().contiguous().t().detach().requires_grad_(True)
elif layout == "strided":
_input = _input.repeat_interleave(2, dim=1)[:, ::2].detach().requires_grad_(True)
weight = weight.repeat_interleave(2, dim=1)[:, ::2].detach().requires_grad_(True)
if layout != "contiguous":
assert not _input.is_contiguous() and not weight.is_contiguous()
originals = tuple(t.detach().clone() if t is not None else None for t in (_input, weight, b))
rec = _BufferRecorder(BT, V, H)
actual = _run_flce(_input, weight, target, bias=b, recorder=rec)
gi = actual[1]
cs = _production_chunk_size(BT, V, H)
rows = [min(cs, BT - start) for start in range(0, BT, cs)]
assert len(rows) > 1 and rows[-1] < cs
logits_out = [entry for entry in rec.mm_out if entry[0][1] == V]
dx_out = [entry for entry in rec.mm_out if entry[0][1] == H]
assert [shape for shape, _, _ in logits_out] == [(n, V) for n in rows]
assert len({ptr for _, ptr, _ in logits_out}) == 1
assert len({ptr for _, _, ptr in logits_out}) == 1
assert rec.empty_shapes.count((cs, V)) == 1
assert [shape for shape, _, _ in dx_out] == [(n, H) for n in rows]
assert {ptr for _, ptr, _ in dx_out} == {gi.untyped_storage().data_ptr()}
assert [ptr for _, _, ptr in dx_out] == [
gi.data_ptr() + start * gi.stride(0) * gi.element_size() for start in range(0, BT, cs)
]
assert all(ptr != gi.untyped_storage().data_ptr() for _, ptr in rec.copies)
# Any remaining allocating mm must be the unchanged dW fallback, never logits/dX.
assert all(shape == (V, H) for shape in rec.mm_default_shapes)
assert rec.logits_gemms == len(rows)
rec.assert_released()
torch.testing.assert_close((_input, weight, b), originals, atol=0, rtol=0)
monkeypatch.setattr(flce_ops, "_can_use_mm_out", lambda *args: False)
expected = _run_flce(_input, weight, target, bias=b)
torch.testing.assert_close(actual, expected)
@pytest.mark.parametrize("ce_impl", _SHARED_CE_IMPLS)
@pytest.mark.parametrize("buffer_path", ["reuse", "forced-fallback", "promoted-bias"])
@pytest.mark.parametrize("use_token_scaling", [False, True])
@pytest.mark.parametrize("forward_only", [False, True])
def test_chunk_storage_released_before_next_chunk_and_final_cast(
self, monkeypatch, ce_impl, buffer_path, use_token_scaling, forward_only
):
BT, V, H = 40, 128, 32
cs = _production_chunk_size(BT, V, H)
_input, weight, b, cw, target = _mk_inputs(
BT,
V,
H,
torch.bfloat16,
bias=True,
ce_weight=True,
seed=16,
first_chunk_size=cs,
)
reuse = buffer_path == "reuse" and not flce_ops.is_hip()
promoted_bias = buffer_path == "promoted-bias"
if buffer_path == "forced-fallback":
monkeypatch.setattr(flce_ops, "_can_use_mm_out", lambda *args: False)
if promoted_bias:
b = b.detach().float().requires_grad_(True)
kwargs = dict(
bias=b,
ce_weight=cw,
accum_dtype=torch.float32,
use_token_scaling=use_token_scaling,
compute_gradients=not forward_only,
softcap=1.5,
label_smoothing=0.1,
lse_square_scale=1e-4,
return_z_loss=True,
return_token_accuracy=True,
return_predicted_tokens=True,
ce_impl=ce_impl,
)
rec = _BufferRecorder(BT, V, H)
with torch.no_grad(), torch.autocast("cuda", enabled=promoted_bias, dtype=torch.bfloat16), rec:
actual = fused_linear_cross_entropy_forward(_input, weight, target, **kwargs)
assert rec.logits_gemms == triton.cdiv(BT, cs)
assert rec.final_casts == (0 if forward_only else 1)
assert bool(rec.mm_out) == reuse
rec.assert_released()
if forward_only:
assert actual[5] is None and actual[6] is None
assert torch.count_nonzero(actual[4]) == 0
assert len(rec.mm_out) == (triton.cdiv(BT, cs) if reuse else 0)
monkeypatch.setattr(flce_ops, "_can_use_mm_out", lambda *args: False)
with torch.no_grad(), torch.autocast("cuda", enabled=promoted_bias, dtype=torch.bfloat16):
expected = fused_linear_cross_entropy_forward(_input, weight, target, **kwargs)
torch.testing.assert_close(actual, expected)
@pytest.mark.parametrize(
"input_dtype,weight_dtype,bias_dtype,amp_dtype,reuse,dx_out",
[
(torch.bfloat16, torch.bfloat16, None, torch.bfloat16, True, True),
(torch.float16, torch.float16, torch.float16, torch.float16, True, True),
(torch.float32, torch.float32, None, torch.bfloat16, False, False),
(torch.float32, torch.bfloat16, None, torch.bfloat16, False, False),
(torch.bfloat16, torch.float32, None, torch.bfloat16, False, False),
(torch.bfloat16, torch.bfloat16, torch.float32, torch.bfloat16, False, False),
(torch.float32, torch.float32, torch.float32, torch.bfloat16, False, False),
(torch.float16, torch.float16, None, torch.bfloat16, False, False),
(torch.float32, torch.float32, torch.bfloat16, None, False, True),
],
)
def test_buffer_amp_gates_preserve_matmul_and_bias_semantics(
self, monkeypatch, input_dtype, weight_dtype, bias_dtype, amp_dtype, reuse, dx_out
):
reuse = reuse and not flce_ops.is_hip()
dx_out = dx_out and not flce_ops.is_hip()
BT, V, H = 40, 128, 32
cs = _production_chunk_size(BT, V, H)
_input, weight, b, _, target = _mk_inputs(
BT, V, H, weight_dtype, input_dtype=input_dtype, bias=bias_dtype is not None, ce_weight=False, seed=17
)
if b is not None:
b = b.detach().to(bias_dtype).requires_grad_(True)
expected_logits = []
with torch.no_grad(), torch.autocast("cuda", enabled=amp_dtype is not None, dtype=amp_dtype):
for start in range(0, BT, cs):
logits = _input[start : start + cs] @ weight.t()
expected_logits.append(logits if b is None else logits + b)
grad_logits = []
real_dispatch = flce_ops.dispatch
def check_ce_logits(op, *args, **kwargs):
torch.testing.assert_close(args[0], expected_logits[len(grad_logits)], atol=0, rtol=0)
result = real_dispatch(op, *args, **kwargs)
grad_logits.append(result[-1].clone())
return result
monkeypatch.setattr(flce_ops, "dispatch", check_ce_logits)
rec = _OpRecorder()
actual = _run_flce(
_input, weight, target, bias=b, autocast_dtype=amp_dtype, ce_impl="nvidia-triton", recorder=rec
)
n_chunks = triton.cdiv(BT, cs)
assert sum(shape[1] == V for shape, _, _ in rec.mm_out) == (n_chunks if reuse else 0)
assert sum(shape[1] == H for shape, _, _ in rec.mm_out) == (n_chunks if dx_out else 0)
assert rec.empty_shapes.count((cs, V)) == (1 if reuse else 0)
assert len(grad_logits) == n_chunks
with torch.no_grad(), torch.autocast("cuda", enabled=amp_dtype is not None, dtype=amp_dtype):
expected_dx = torch.cat([g @ weight for g in grad_logits]).to(_input.dtype)
torch.testing.assert_close(actual[1], expected_dx, atol=0, rtol=0)
monkeypatch.setattr(flce_ops, "dispatch", real_dispatch)
monkeypatch.setattr(flce_ops, "_can_use_mm_out", lambda *args: False)
expected = _run_flce(_input, weight, target, bias=b, autocast_dtype=amp_dtype, ce_impl="nvidia-triton")
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
@requires_nvidia_cuda
def test_mm_out_gate_uses_legacy_autocast_query_on_old_torch(self, monkeypatch):
_input = torch.empty(1, 1, device="cuda", dtype=torch.bfloat16)
monkeypatch.setattr(flce_ops, "_TORCH_VERSION", flce_ops.Version("2.1.2"))
monkeypatch.setattr(torch, "is_autocast_enabled", lambda: True)
monkeypatch.setattr(torch, "get_autocast_gpu_dtype", lambda: torch.bfloat16)
def unexpected_modern_query(*args):
pytest.fail("torch < 2.4 must not use get_autocast_dtype(device_type)")
monkeypatch.setattr(torch, "get_autocast_dtype", unexpected_modern_query, raising=False)
with torch.no_grad():
assert flce_ops._can_use_mm_out(_input, _input)
monkeypatch.setattr(torch, "get_autocast_gpu_dtype", lambda: torch.float16)
assert not flce_ops._can_use_mm_out(_input, _input)
assert not flce_ops._can_use_mm_out(_input, _input)
def test_mm_out_gate_preserves_hip_fallback(self, monkeypatch):
_input = torch.empty(1, 1, device="cuda", dtype=torch.float32)
monkeypatch.setattr(flce_ops, "is_hip", lambda: True)
with torch.no_grad():
assert not flce_ops._can_use_mm_out(_input, _input)
@pytest.mark.parametrize("ce_impl", _SHARED_CE_IMPLS)
@pytest.mark.parametrize("use_token_scaling", [False, True])
@pytest.mark.parametrize("weight_requires_grad", [False, True])
def test_reduction_none_reenters_with_reusable_buffers(
self, monkeypatch, ce_impl, use_token_scaling, weight_requires_grad
):
BT, V, H = 40, 128, 32
cs = _production_chunk_size(BT, V, H)
_input, weight, b, cw, target = _mk_inputs(
BT,
V,
H,
torch.bfloat16,
bias=True,
ce_weight=True,
seed=18,
first_chunk_size=cs,
weight_requires_grad=weight_requires_grad,
)
upstream = torch.linspace(-0.5, 1.5, BT, device="cuda")
def run():
outputs = flce_ops.LigerFusedLinearCrossEntropyFunction.apply(
_input,
weight,
target,
b,
cw,
-100,
1e-4,
0.1,
"none",
1.5,
True,
torch.float32,
use_token_scaling,
True,
True,
ce_impl,
)
outputs[0].backward(upstream)
return outputs, _input.grad, weight.grad, b.grad
rec = _BufferRecorder(BT, V, H)
with rec:
actual = run()
assert rec.logits_gemms == 2 * triton.cdiv(BT, cs)
assert rec.final_casts == int(weight_requires_grad)
expected_mm_out = 0 if flce_ops.is_hip() else 3 * triton.cdiv(BT, cs)
assert len(rec.mm_out) == expected_mm_out # logits twice, dX only during backward
rec.assert_released()
assert torch.count_nonzero(actual[1][:cs]) == 0
if not weight_requires_grad:
assert actual[2] is None
assert not rec.has_addmm_out_into(weight.shape)
assert not rec.has_addmm_dtype_out_into(weight.shape)
_input.grad = weight.grad = b.grad = None
monkeypatch.setattr(flce_ops, "_can_use_mm_out", lambda *args: False)
expected = run()
torch.testing.assert_close(actual, expected)
# -----------------------------------------------------------------------
# 8. AMP: in-place out ops are not autocast, so _input_chunk must be aligned
# -----------------------------------------------------------------------
@@ -839,7 +1162,7 @@ class TestFusedLinearCrossEntropyAddmm:
# grad_logits promoted to fp32 -> no fp32-out fast path, no low-prec out path.
assert not rec.has_addmm_dtype_out_into(weight.shape)
assert not rec.has_addmm_out_into(weight.shape)
assert rec.mm_default >= 1
assert rec.mm_default_shapes.count(tuple(weight.shape)) == triton.cdiv(BT, _production_chunk_size(BT, V, H))
_, ref_gi, ref_gw, ref_gb = _addmm_reference_grads(_input, weight, target, bias=bias)
torch.testing.assert_close(gw.float(), ref_gw, atol=1e-2, rtol=1e-2)
torch.testing.assert_close(gi.float(), ref_gi, atol=1e-2, rtol=1e-2)
@@ -877,7 +1200,7 @@ class TestFusedLinearCrossEntropyAddmm:
# Sub-SM80 capability -> neither addmm fast path; legacy mm + fp32 upcast instead.
assert not rec.has_addmm_out_into(weight.shape)
assert not rec.has_addmm_dtype_out_into(weight.shape)
assert rec.mm_default >= 1
assert rec.mm_default_shapes.count(tuple(weight.shape)) == triton.cdiv(BT, _production_chunk_size(BT, V, H))
assert rec.has_param_upcast(weight.shape)
_, ref_gi, ref_gw, _ = _addmm_reference_grads(_input, weight, target)
torch.testing.assert_close(gw.float(), ref_gw, **_ADDMM_TOLS[torch.bfloat16])