mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
perf(flce): accumulate low-precision Triton dW with addmm
Accumulate matching BF16/FP16 weight gradients directly into their existing buffer on SM80+, avoiding the per-chunk parameter-sized FP32 conversion. Preserve the guarded FP32 fast path, AMP dtype alignment, and legacy fallbacks. Add capability-aware regression coverage and explicit benchmark accumulator-policy controls. Co-authored-by: Copilot <[email protected]>
This commit is contained in:
Vaibhav Jindal
co-authored by
Copilot
Copilot
parent
e4aad6c0c0
commit
3eddb174e0
@@ -263,9 +263,16 @@ This harness imports the backend `autograd.Function` classes **directly** (no di
|
||||
--chunk-sizes 256 1024 4096
|
||||
```
|
||||
|
||||
Pass `--accum-dtype default` to compare each backend's *native* accumulator policy instead of the FP32-matched default. **Warning:** the `default` panel is **not** precision-matched — with the BF16 weights used here `triton` and `cutedsl` accumulate the cross-chunk `dW` in **BF16** buffers while `cutile` **always** keeps an **FP32** accumulator:
|
||||
|
||||
```bash
|
||||
.venv/bin/python benchmark/scripts/benchmark_flce_backends.py \
|
||||
--backends all3 --accum-dtype default
|
||||
```
|
||||
|
||||
- **Backends & arch policy.** `cutedsl` requires exactly SM100 (Blackwell, e.g. B200); `cutile` runs on SM90 (Hopper) **and** SM100. So on a B200 the default — and `--backends all3` — is all three (`triton`, `cutedsl`, `cutile`) in one process; on SM90 the default is `triton` + `cutile`; elsewhere `triton` only. Explicitly requesting a backend the GPU cannot run fails loudly rather than degrading, and a missing SDK surfaces its own `ImportError`.
|
||||
- **GEMM ownership differs per backend** (this is why the numbers are not apples-to-apples kernel isolation): `triton` runs the projection, `dX` and `dW` all through PyTorch/cuBLAS; `cutedsl`'s projection uses CuTe DSL SM100 GEMM; `dX`/`dW` use PyTorch/cuBLAS; `cutile` uses PyTorch BLAS for the projection and `dX` and only a cuTile MMA schedule for `dW`.
|
||||
- **Numerics.** Inputs are BF16 with `accum_dtype=torch.float32`, so the `dW` accumulation precision matches across backends. The `cutile` `dW` retains raw FP32 through accumulation and is cast to BF16 only after the combined upstream × global-mean scaling in backward, preserving the tiny-value range of the original path.
|
||||
- **Numerics & `--accum-dtype`.** Inputs are BF16. `--accum-dtype fp32` (the default, backward-compatible) forces `accum_dtype=torch.float32` on every backend so the cross-chunk `dW` accumulation precision matches, and the table's `eff_accum` column reads `fp32` for every row. `--accum-dtype default` instead passes `None` and lets each backend apply its own policy: with BF16 weights `triton` and the native `cutedsl` path store the cross-chunk `dW` in **BF16**, while `cutile` **always** retains an **FP32** accumulator — so that panel is **not** precision-matched (`eff_accum` reads `bf16` for `triton`/`cutedsl` and `fp32` for `cutile`). The requested mode is echoed in the header. The `cutile` `dW` retains raw FP32 through accumulation and is cast to BF16 only after the combined upstream × global-mean scaling in backward, preserving the tiny-value range of the original path.
|
||||
- **`chunk` vs `eff_chunk`.** A chunk larger than the token count is clamped to `B*T`; the table prints both the requested `chunk` and the `eff_chunk` actually used, so a clamp is never reported as the request.
|
||||
- **Peak memory** is the *incremental* peak **allocated** bytes (not reserved) above an input-only baseline: both leaf grads are cleared before the baseline is snapshotted, and a dedicated memory-probe iteration is kept separate from the timed loop.
|
||||
|
||||
|
||||
@@ -48,6 +48,16 @@ Requested vs. effective chunk size: a chunk larger than the token count ``N`` is
|
||||
clamped to ``N``; the table reports the ``chunk`` you asked for and the
|
||||
``eff_chunk`` actually used, so a clamp is never mistaken for the request.
|
||||
|
||||
Cross-chunk ``dW`` accumulator (``--accum-dtype``): ``fp32`` (the default, kept for
|
||||
backward compatibility) passes ``torch.float32`` to every backend so the ``dW``
|
||||
accumulation precision is matched across ``triton`` / ``cutedsl`` / ``cutile``.
|
||||
``default`` passes ``None`` and lets each backend apply its own policy: with the
|
||||
BF16 weights used here ``triton`` and the native ``cutedsl`` path store the
|
||||
cross-chunk ``dW`` in **BF16**, while ``cutile`` **always** retains an **FP32**
|
||||
accumulator even when ``None`` is requested. The header states the requested mode
|
||||
and the table's ``eff_accum`` column reports the effective accumulator per backend,
|
||||
so the ``default`` panel is never mislabelled as precision-matched.
|
||||
|
||||
Results are printed to stdout only — no CSV is written into the repository.
|
||||
|
||||
Example::
|
||||
@@ -55,6 +65,10 @@ Example::
|
||||
.venv/bin/python benchmark/scripts/benchmark_flce_backends.py \\
|
||||
--backends all3 --tokens 4096 --hidden-size 2048 --vocab-size 32000 \\
|
||||
--chunk-sizes 256 1024 4096
|
||||
|
||||
# per-backend default accumulator policy (triton/cutedsl BF16, cutile FP32):
|
||||
.venv/bin/python benchmark/scripts/benchmark_flce_backends.py \\
|
||||
--backends all3 --accum-dtype default
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -68,6 +82,10 @@ import torch
|
||||
_ACCUM_DTYPE = torch.float32
|
||||
_DTYPE = torch.bfloat16
|
||||
|
||||
# ``--accum-dtype`` maps a requested mode onto the value passed at Function slot 11:
|
||||
# ``fp32`` -> torch.float32 (matched precision), ``default`` -> None (per-backend policy).
|
||||
_ACCUM_DTYPE_CHOICES = {"fp32": torch.float32, "default": None}
|
||||
|
||||
_MODULE_PATHS = {
|
||||
"triton": "liger_kernel.ops.fused_linear_cross_entropy",
|
||||
"cutedsl": "liger_kernel.ops.cutedsl.ops.fused_linear_cross_entropy",
|
||||
@@ -144,6 +162,21 @@ def _resolve_backends(requested, cc):
|
||||
return list(requested)
|
||||
|
||||
|
||||
def _effective_accum(name, requested):
|
||||
"""Human label of the *effective* cross-chunk ``dW`` accumulator for BF16 weights.
|
||||
|
||||
``fp32`` forces ``torch.float32`` on every backend. Under ``default`` (``None``
|
||||
passed through) the native ``triton`` / ``cutedsl`` paths store the cross-chunk
|
||||
``dW`` in BF16 to match the BF16 weights, while ``cutile`` *always* retains an
|
||||
FP32 accumulator regardless — so the two are not precision-matched.
|
||||
"""
|
||||
if requested == "fp32":
|
||||
return "fp32"
|
||||
if name == "cutile":
|
||||
return "fp32"
|
||||
return "bf16"
|
||||
|
||||
|
||||
def _build_inputs(tokens, hidden, vocab, device, seed=0):
|
||||
"""Deterministic BF16 inputs — identical values reused for every backend."""
|
||||
generator = torch.Generator(device=device).manual_seed(seed)
|
||||
@@ -153,9 +186,11 @@ def _build_inputs(tokens, hidden, vocab, device, seed=0):
|
||||
return x, w, target
|
||||
|
||||
|
||||
def _apply(fn, x, w, target, chunk_size):
|
||||
# 15 base positional args + (ce_impl, ce_mode, chunk_size); accum_dtype=float32
|
||||
# so the dW accumulation precision matches across triton / cutedsl / cutile.
|
||||
def _apply(fn, x, w, target, chunk_size, accum_dtype=_ACCUM_DTYPE):
|
||||
# 15 base positional args + (ce_impl, ce_mode, chunk_size). ``accum_dtype`` is the
|
||||
# cross-chunk dW accumulator at slot 11: torch.float32 (matched precision) or None
|
||||
# (per-backend policy — see _effective_accum). Threaded as a parameter, never a
|
||||
# mutable global, so one process can compare modes without state leak.
|
||||
loss, _, _, _ = fn.apply(
|
||||
x,
|
||||
w,
|
||||
@@ -168,7 +203,7 @@ def _apply(fn, x, w, target, chunk_size):
|
||||
"mean", # reduction
|
||||
None, # softcap
|
||||
False, # return_z_loss
|
||||
_ACCUM_DTYPE, # accum_dtype
|
||||
accum_dtype, # accum_dtype
|
||||
False, # use_token_scaling
|
||||
False, # return_token_accuracy
|
||||
False, # return_predicted_tokens
|
||||
@@ -179,16 +214,16 @@ def _apply(fn, x, w, target, chunk_size):
|
||||
return loss
|
||||
|
||||
|
||||
def _one_iteration(fn, x, w, target, chunk_size, upstream):
|
||||
def _one_iteration(fn, x, w, target, chunk_size, upstream, accum_dtype=_ACCUM_DTYPE):
|
||||
x.grad = None
|
||||
w.grad = None
|
||||
loss = _apply(fn, x, w, target, chunk_size)
|
||||
loss = _apply(fn, x, w, target, chunk_size, accum_dtype=accum_dtype)
|
||||
loss.backward(upstream)
|
||||
|
||||
|
||||
def _loss_value(fn, x, w, target, chunk_size):
|
||||
def _loss_value(fn, x, w, target, chunk_size, accum_dtype=_ACCUM_DTYPE):
|
||||
"""Return a representative loss as a Python float, keeping no autograd graph alive."""
|
||||
loss = _apply(fn, x, w, target, chunk_size)
|
||||
loss = _apply(fn, x, w, target, chunk_size, accum_dtype=accum_dtype)
|
||||
value = float(loss.detach())
|
||||
del loss
|
||||
x.grad = None
|
||||
@@ -196,13 +231,13 @@ def _loss_value(fn, x, w, target, chunk_size):
|
||||
return value
|
||||
|
||||
|
||||
def _time_backend(fn, x, w, target, chunk_size, warmup, iters):
|
||||
def _time_backend(fn, x, w, target, chunk_size, warmup, iters, accum_dtype=_ACCUM_DTYPE):
|
||||
upstream = torch.ones((), device=x.device, dtype=torch.float32)
|
||||
|
||||
# One unconditional compile pass (excluded) so warmup may legitimately be 0.
|
||||
_one_iteration(fn, x, w, target, chunk_size, upstream)
|
||||
_one_iteration(fn, x, w, target, chunk_size, upstream, accum_dtype=accum_dtype)
|
||||
for _ in range(warmup):
|
||||
_one_iteration(fn, x, w, target, chunk_size, upstream)
|
||||
_one_iteration(fn, x, w, target, chunk_size, upstream, accum_dtype=accum_dtype)
|
||||
|
||||
# Baseline must be input-only: drop both leaf grads before snapshotting.
|
||||
x.grad = None
|
||||
@@ -212,7 +247,7 @@ def _time_backend(fn, x, w, target, chunk_size, warmup, iters):
|
||||
# Dedicated memory probe (separate from the timed loop for a precise peak).
|
||||
torch.cuda.reset_peak_memory_stats(x.device)
|
||||
baseline = torch.cuda.memory_allocated(x.device)
|
||||
_one_iteration(fn, x, w, target, chunk_size, upstream)
|
||||
_one_iteration(fn, x, w, target, chunk_size, upstream, accum_dtype=accum_dtype)
|
||||
torch.cuda.synchronize()
|
||||
peak = torch.cuda.max_memory_allocated(x.device)
|
||||
incremental_peak_mb = max(0.0, (peak - baseline) / (1024**2))
|
||||
@@ -223,7 +258,7 @@ def _time_backend(fn, x, w, target, chunk_size, warmup, iters):
|
||||
ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
|
||||
for i in range(iters):
|
||||
starts[i].record()
|
||||
_one_iteration(fn, x, w, target, chunk_size, upstream)
|
||||
_one_iteration(fn, x, w, target, chunk_size, upstream, accum_dtype=accum_dtype)
|
||||
ends[i].record()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
@@ -240,26 +275,45 @@ def _run(args):
|
||||
for name in backends:
|
||||
_check_backend_supported(name, cc)
|
||||
|
||||
requested = args.accum_dtype
|
||||
accum_value = _ACCUM_DTYPE_CHOICES[requested]
|
||||
if requested == "fp32":
|
||||
accum_note = "cross-chunk dW accumulator forced to FP32 on every backend (precision matched)"
|
||||
else:
|
||||
accum_note = (
|
||||
"per-backend accumulator policy for BF16 weights: triton/cutedsl store cross-chunk dW in "
|
||||
"BF16, cutile always retains an FP32 accumulator (NOT precision matched)"
|
||||
)
|
||||
|
||||
print(
|
||||
f"device={torch.cuda.get_device_name(device)} cc={cc} "
|
||||
f"tokens={args.tokens} hidden={args.hidden_size} vocab={args.vocab_size} "
|
||||
f"dtype={_DTYPE} accum_dtype={_ACCUM_DTYPE}"
|
||||
f"dtype={_DTYPE} accum_dtype_requested={requested} ({accum_note})"
|
||||
)
|
||||
header = (
|
||||
f"{'backend':<10} {'chunk':>8} {'eff_chunk':>10} {'eff_accum':>10} "
|
||||
f"{'median_ms':>12} {'peak_MiB':>12} {'loss':>12}"
|
||||
)
|
||||
header = f"{'backend':<10} {'chunk':>8} {'eff_chunk':>10} {'median_ms':>12} {'peak_MiB':>12} {'loss':>12}"
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
|
||||
for name in backends:
|
||||
fn = _load_backend(name)
|
||||
eff_accum = _effective_accum(name, requested)
|
||||
for chunk_size in args.chunk_sizes:
|
||||
effective = min(chunk_size, args.tokens)
|
||||
x, w, target = _build_inputs(args.tokens, args.hidden_size, args.vocab_size, device, seed=args.seed)
|
||||
x.requires_grad_(True)
|
||||
w.requires_grad_(True)
|
||||
# Representative loss (recomputed cleanly, no graph retained into timing).
|
||||
loss_val = _loss_value(fn, x, w, target, chunk_size)
|
||||
median_ms, peak_mb = _time_backend(fn, x, w, target, chunk_size, args.warmup, args.iters)
|
||||
print(f"{name:<10} {chunk_size:>8} {effective:>10} {median_ms:>12.4f} {peak_mb:>12.2f} {loss_val:>12.5f}")
|
||||
loss_val = _loss_value(fn, x, w, target, chunk_size, accum_dtype=accum_value)
|
||||
median_ms, peak_mb = _time_backend(
|
||||
fn, x, w, target, chunk_size, args.warmup, args.iters, accum_dtype=accum_value
|
||||
)
|
||||
print(
|
||||
f"{name:<10} {chunk_size:>8} {effective:>10} {eff_accum:>10} "
|
||||
f"{median_ms:>12.4f} {peak_mb:>12.2f} {loss_val:>12.5f}"
|
||||
)
|
||||
# Drop this shape's tensors before the next backend/chunk allocates.
|
||||
del x, w, target
|
||||
|
||||
@@ -301,6 +355,17 @@ def _build_parser():
|
||||
"--warmup", type=_non_negative_int, default=5, help="warmup iterations excluded from timing (default: 5)"
|
||||
)
|
||||
parser.add_argument("--iters", type=_positive_int, default=20, help="timed iterations, must be >= 1 (default: 20)")
|
||||
parser.add_argument(
|
||||
"--accum-dtype",
|
||||
choices=sorted(_ACCUM_DTYPE_CHOICES),
|
||||
default="fp32",
|
||||
help=(
|
||||
"cross-chunk dW accumulator policy: 'fp32' (default, backward-compatible) passes "
|
||||
"torch.float32 to every backend so accumulation precision matches; 'default' passes None "
|
||||
"so each backend uses its own policy — with BF16 weights triton/cutedsl accumulate dW in "
|
||||
"BF16 while cutile always keeps an FP32 accumulator (default: fp32)"
|
||||
),
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=0, help="RNG seed for the shared inputs (default: 0)")
|
||||
return parser
|
||||
|
||||
|
||||
@@ -105,7 +105,10 @@ def fused_linear_cross_entropy_forward(
|
||||
|
||||
grad_input = torch.zeros_like(_input, device=device)
|
||||
|
||||
# we use fp32 for loss and gradients accumulator
|
||||
# loss_1d is always fp32. The weight/bias grad accumulators are NOT always fp32:
|
||||
# with accum_dtype=None they inherit the parameter dtype (e.g. bf16/fp16), and only
|
||||
# take accum_dtype when one is explicitly requested. The weight-grad projection below
|
||||
# therefore has a dtype-matched low-precision path in addition to the fp32 fast path.
|
||||
if input_requires_grad:
|
||||
if accum_dtype is None:
|
||||
grad_weight = torch.zeros_like(weight, device=device) if weight_needs_grad else None
|
||||
@@ -290,13 +293,17 @@ def fused_linear_cross_entropy_forward(
|
||||
|
||||
if grad_weight is not None and input_requires_grad:
|
||||
grad_logits_t = grad_logits_chunk.t()
|
||||
grad_logits_is_half = grad_logits_t.dtype in (torch.float16, torch.bfloat16)
|
||||
on_cuda_sm80 = (
|
||||
grad_weight.device.type == "cuda" and torch.cuda.get_device_capability(grad_weight.device)[0] >= 8
|
||||
)
|
||||
if (
|
||||
_ADDMM_SUPPORTS_OUT_DTYPE
|
||||
and grad_weight.device.type == "cuda"
|
||||
and torch.cuda.get_device_capability(grad_weight.device)[0] >= 8
|
||||
and on_cuda_sm80
|
||||
and grad_weight.dtype == torch.float32
|
||||
and grad_logits_t.dtype in (torch.float16, torch.bfloat16)
|
||||
and grad_logits_is_half
|
||||
):
|
||||
# FP32 accumulator (accum_dtype=torch.float32, or fp32 params under AMP).
|
||||
# Unlike torch.mm, torch.addmm's out_dtype path does not participate in
|
||||
# autocast operand casting, so under AMP (fp32 params, no bias) _input_chunk
|
||||
# can stay fp32 while grad_logits is the autocast dtype. addmm requires mat1
|
||||
@@ -311,7 +318,26 @@ def fused_linear_cross_entropy_forward(
|
||||
out_dtype=torch.float32,
|
||||
out=grad_weight,
|
||||
)
|
||||
elif on_cuda_sm80 and grad_logits_is_half 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); this mirrors the CuTe backend's direct bf16 addmm and
|
||||
# avoids the legacy path's parameter-sized bf16->fp32 temporary + cast per chunk.
|
||||
# In-place out ops are not autocast, so -- as torch.mm's autocast used to do --
|
||||
# align _input_chunk (which may be promoted fp32 under bf16 AMP) to grad_logits.
|
||||
input_chunk = _input_chunk
|
||||
if input_chunk.dtype != grad_logits_t.dtype:
|
||||
input_chunk = input_chunk.to(grad_logits_t.dtype)
|
||||
torch.addmm(
|
||||
grad_weight,
|
||||
grad_logits_t,
|
||||
input_chunk,
|
||||
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 += torch.mm(grad_logits_chunk.t(), _input_chunk).float()
|
||||
|
||||
if bias is not None and input_requires_grad:
|
||||
|
||||
@@ -0,0 +1,628 @@
|
||||
"""Regression tests for the weight-gradient ``addmm`` projection in
|
||||
``fused_linear_cross_entropy_forward``.
|
||||
|
||||
These lock in the low-precision ``grad_weight`` accumulation path: when
|
||||
``accum_dtype=None`` and the parameters are bf16/fp16, the per-chunk weight
|
||||
gradient must be accumulated *directly* into ``grad_weight`` via
|
||||
``torch.addmm(out=grad_weight)`` (mirroring the CuTe backend's direct bf16
|
||||
addmm), instead of the legacy ``torch.mm(...).float()`` which materialises a
|
||||
parameter-sized bf16->fp32 temporary + cast every chunk.
|
||||
|
||||
The tests assert on the actual aten dispatch (``TorchDispatchMode``) *and* on
|
||||
numerical parity against a full-precision autograd reference. The whole file is
|
||||
GPU-gated by ``test/ops/conftest.py``'s autouse ``_require_cuda_for_ops_tests``
|
||||
fixture, and the accumulation path itself requires a CUDA SM80+ device.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
|
||||
import liger_kernel.ops.fused_linear_cross_entropy as flce_ops
|
||||
|
||||
from liger_kernel.ops.fused_linear_cross_entropy import fused_linear_cross_entropy_forward
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fast-path capability gates
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# The whole file is CUDA-gated by ``conftest.py``'s autouse fixture, but *being
|
||||
# on CUDA* is not sufficient for the ``addmm(out=)`` / ``addmm.dtype_out`` fast
|
||||
# paths. Production (``fused_linear_cross_entropy_forward``) gates them on:
|
||||
#
|
||||
# * ``on_cuda_sm80`` -- ``get_device_capability()[0] >= 8`` (SM80+; this is
|
||||
# also exactly where bf16 GEMM hardware becomes available), and
|
||||
# * ``_ADDMM_SUPPORTS_OUT_DTYPE`` -- torch >= 2.8, additionally required for
|
||||
# the fp32-accumulator ``addmm.dtype_out`` overload.
|
||||
#
|
||||
# On valid-but-older setups (real SM70/SM75, or the base ``torch>=2.1.2`` pin
|
||||
# through torch<2.8) production correctly falls back to ``torch.mm(...).float()``.
|
||||
# Only tests that *assert the fast path was taken* must be gated with the same
|
||||
# predicate; parity/fallback tests stay ungated so they still run everywhere.
|
||||
|
||||
|
||||
def _cuda_capability_major() -> int:
|
||||
if not torch.cuda.is_available():
|
||||
return -1
|
||||
return torch.cuda.get_device_capability()[0]
|
||||
|
||||
|
||||
# Reused across every fast-path-asserting test (mirrors production's on_cuda_sm80).
|
||||
requires_sm80 = pytest.mark.skipif(
|
||||
_cuda_capability_major() < 8,
|
||||
reason="addmm(out=) fast path requires CUDA SM80+ (get_device_capability()[0] >= 8)",
|
||||
)
|
||||
|
||||
# Additionally required by the fp32-accumulator addmm.dtype_out fast path.
|
||||
requires_addmm_out_dtype = pytest.mark.skipif(
|
||||
not flce_ops._ADDMM_SUPPORTS_OUT_DTYPE,
|
||||
reason="addmm out_dtype overload (fp32 accumulator fast path) requires torch >= 2.8",
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dispatch recorder
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_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)."""
|
||||
|
||||
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.to_copies = [] # list of (in_dtype, out_dtype, shape)
|
||||
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
kwargs = kwargs or {}
|
||||
name = str(func)
|
||||
if name in _ADDMM_OPS:
|
||||
out = kwargs.get("out")
|
||||
shape = tuple(out.shape) if out is not None else None
|
||||
if name == "aten.addmm.out":
|
||||
self.addmm_out_shapes.append(shape)
|
||||
else:
|
||||
self.addmm_dtype_out_shapes.append(shape)
|
||||
elif name == "aten.mm.default":
|
||||
self.mm_default += 1
|
||||
elif name == "aten._to_copy.default":
|
||||
src = args[0] if args else None
|
||||
if isinstance(src, torch.Tensor):
|
||||
self.to_copies.append((src.dtype, kwargs.get("dtype"), tuple(src.shape)))
|
||||
return func(*args, **kwargs)
|
||||
|
||||
def has_addmm_out_into(self, shape) -> bool:
|
||||
return tuple(shape) in self.addmm_out_shapes
|
||||
|
||||
def has_addmm_dtype_out_into(self, shape) -> bool:
|
||||
return tuple(shape) in self.addmm_dtype_out_shapes
|
||||
|
||||
def has_param_upcast(self, shape, low_dtypes=(torch.bfloat16, torch.float16)) -> bool:
|
||||
"""True iff some ``_to_copy`` upcast a *parameter-shaped* low-precision
|
||||
tensor to fp32 (the legacy ``.float()`` temporary we want to eliminate)."""
|
||||
tgt = tuple(shape)
|
||||
return any(src in low_dtypes and dst == torch.float32 and sh == tgt for (src, dst, sh) in self.to_copies)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reference + driver helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
DEVICE = "cuda"
|
||||
|
||||
|
||||
def _reference_grads(
|
||||
_input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
target: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
ce_weight: Optional[torch.Tensor] = None,
|
||||
ignore_index: int = -100,
|
||||
label_smoothing: float = 0.0,
|
||||
reduction: str = "mean",
|
||||
token_grad_output: Optional[torch.Tensor] = None,
|
||||
):
|
||||
"""Full-precision autograd reference for (loss, grad_input, grad_weight, grad_bias)."""
|
||||
xi = _input.detach().float().requires_grad_(True)
|
||||
wi = weight.detach().float().requires_grad_(weight.requires_grad)
|
||||
bi = None
|
||||
if bias is not None:
|
||||
bi = bias.detach().float().requires_grad_(bias.requires_grad)
|
||||
logits = xi @ wi.t()
|
||||
if bi is not None:
|
||||
logits = logits + bi
|
||||
cw = ce_weight.detach().float() if ce_weight is not None else None
|
||||
loss = F.cross_entropy(
|
||||
logits,
|
||||
target,
|
||||
weight=cw,
|
||||
ignore_index=ignore_index,
|
||||
label_smoothing=label_smoothing,
|
||||
reduction=reduction,
|
||||
)
|
||||
if reduction == "none":
|
||||
assert token_grad_output is not None
|
||||
loss.backward(token_grad_output.float())
|
||||
else:
|
||||
loss.backward()
|
||||
return (
|
||||
loss.detach(),
|
||||
xi.grad.detach(),
|
||||
wi.grad.detach() if wi.grad is not None else None,
|
||||
bi.grad.detach() if (bi is not None and bi.grad is not None) else None,
|
||||
)
|
||||
|
||||
|
||||
def _run_flce(
|
||||
_input,
|
||||
weight,
|
||||
target,
|
||||
*,
|
||||
bias=None,
|
||||
ce_weight=None,
|
||||
ignore_index=-100,
|
||||
label_smoothing=0.0,
|
||||
reduction="mean",
|
||||
accum_dtype=None,
|
||||
chunk_size=None,
|
||||
use_token_scaling=False,
|
||||
token_grad_output=None,
|
||||
recorder: Optional[_OpRecorder] = None,
|
||||
autocast_dtype=None,
|
||||
):
|
||||
"""Invoke the forward the same way the autograd Function does: under
|
||||
``no_grad`` (Function.forward runs with grad disabled), optionally under
|
||||
autocast, optionally under a dispatch recorder."""
|
||||
|
||||
def _call():
|
||||
return fused_linear_cross_entropy_forward(
|
||||
_input,
|
||||
weight,
|
||||
target,
|
||||
ce_weight=ce_weight,
|
||||
bias=bias,
|
||||
ignore_index=ignore_index,
|
||||
label_smoothing=label_smoothing,
|
||||
reduction=reduction,
|
||||
accum_dtype=accum_dtype,
|
||||
use_token_scaling=use_token_scaling,
|
||||
token_grad_output=token_grad_output,
|
||||
chunk_size=chunk_size,
|
||||
)
|
||||
|
||||
import contextlib
|
||||
|
||||
stack = contextlib.ExitStack()
|
||||
with stack:
|
||||
stack.enter_context(torch.no_grad())
|
||||
if autocast_dtype is not None:
|
||||
stack.enter_context(torch.autocast("cuda", dtype=autocast_dtype))
|
||||
if recorder is not None:
|
||||
stack.enter_context(recorder)
|
||||
loss, z_loss, token_acc, pred, grad_input, grad_weight, grad_bias = _call()
|
||||
return loss, grad_input, grad_weight, grad_bias
|
||||
|
||||
|
||||
_TOLS = {
|
||||
torch.float16: dict(atol=5e-2, rtol=1e-2),
|
||||
torch.bfloat16: dict(atol=1e-1, rtol=2e-2),
|
||||
torch.float32: dict(atol=1e-4, rtol=1e-4),
|
||||
}
|
||||
|
||||
|
||||
def _mk_inputs(
|
||||
BT,
|
||||
V,
|
||||
H,
|
||||
dtype,
|
||||
*,
|
||||
bias,
|
||||
ce_weight,
|
||||
seed=0,
|
||||
input_dtype=None,
|
||||
weight_requires_grad=True,
|
||||
input_requires_grad=True,
|
||||
all_ignore=False,
|
||||
first_chunk_ignore=False,
|
||||
chunk_size=None,
|
||||
):
|
||||
g = torch.Generator(device=DEVICE).manual_seed(seed)
|
||||
input_dtype = input_dtype or dtype
|
||||
_input = torch.randn(BT, H, device=DEVICE, dtype=input_dtype, generator=g) * 0.1
|
||||
_input.requires_grad_(input_requires_grad)
|
||||
weight = torch.randn(V, H, device=DEVICE, dtype=dtype, generator=g) * 0.1
|
||||
weight.requires_grad_(weight_requires_grad)
|
||||
b = None
|
||||
if bias:
|
||||
b = (torch.randn(V, device=DEVICE, dtype=dtype, generator=g) * 0.1).requires_grad_(True)
|
||||
cw = None
|
||||
if ce_weight:
|
||||
cw = torch.rand(V, device=DEVICE, dtype=torch.float32, generator=g) + 0.1
|
||||
target = torch.randint(0, V, (BT,), device=DEVICE, generator=g, dtype=torch.long)
|
||||
if all_ignore:
|
||||
target[:] = -100
|
||||
elif first_chunk_ignore and chunk_size is not None:
|
||||
target[:chunk_size] = -100
|
||||
return _input, weight, b, cw, target
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. Dispatch: the low-precision default path uses addmm(out=) and no param upcast
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@requires_sm80
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("bias", [True, False])
|
||||
def test_lowprec_default_uses_addmm_out_no_param_upcast(dtype, bias):
|
||||
BT, V, H, cs = 40, 64, 32, 16 # 3 chunks (16,16,8) -> multichunk + tail
|
||||
_input, weight, b, cw, target = _mk_inputs(BT, V, H, dtype, bias=bias, ce_weight=False, chunk_size=cs)
|
||||
|
||||
rec = _OpRecorder()
|
||||
loss, gi, gw, gb = _run_flce(_input, weight, target, bias=b, chunk_size=cs, accum_dtype=None, recorder=rec)
|
||||
|
||||
# The accumulator inherits the (low-precision) parameter dtype ...
|
||||
assert gw.dtype == weight.dtype
|
||||
# ... and is written by addmm(out=grad_weight) into a parameter-shaped buffer.
|
||||
assert rec.has_addmm_out_into(weight.shape), (
|
||||
f"expected aten.addmm.out into {tuple(weight.shape)}, saw {rec.addmm_out_shapes}"
|
||||
)
|
||||
# The fp32 out_dtype overload must NOT be used for a low-precision accumulator.
|
||||
assert not rec.has_addmm_dtype_out_into(weight.shape)
|
||||
# And crucially: no parameter-sized bf16/fp16 -> fp32 temporary (the legacy .float()).
|
||||
assert not rec.has_param_upcast(weight.shape), f"unexpected parameter-shaped upcast to fp32: {rec.to_copies}"
|
||||
|
||||
# Parity vs full-precision autograd reference.
|
||||
tol = _TOLS[dtype]
|
||||
_, ref_gi, ref_gw, ref_gb = _reference_grads(_input, weight, target, bias=b)
|
||||
assert torch.isfinite(gw).all()
|
||||
torch.testing.assert_close(gw.float(), ref_gw, **tol)
|
||||
torch.testing.assert_close(gi.float(), ref_gi, **tol)
|
||||
if bias:
|
||||
torch.testing.assert_close(gb.float(), ref_gb, **tol)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. Numerical parity across the combinatorial space (small shapes)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
|
||||
@pytest.mark.parametrize("reduction", ["mean", "sum"])
|
||||
@pytest.mark.parametrize("bias", [True, False])
|
||||
@pytest.mark.parametrize("ce_weight", [True, False])
|
||||
@pytest.mark.parametrize("label_smoothing", [0.0, 0.1])
|
||||
def test_parity_combinations(dtype, reduction, bias, ce_weight, label_smoothing):
|
||||
BT, V, H, cs = 40, 64, 32, 16
|
||||
_input, weight, b, cw, target = _mk_inputs(BT, V, H, dtype, bias=bias, ce_weight=ce_weight, chunk_size=cs, seed=1)
|
||||
_, gi, gw, gb = _run_flce(
|
||||
_input,
|
||||
weight,
|
||||
target,
|
||||
bias=b,
|
||||
ce_weight=cw,
|
||||
reduction=reduction,
|
||||
label_smoothing=label_smoothing,
|
||||
chunk_size=cs,
|
||||
)
|
||||
_, ref_gi, ref_gw, ref_gb = _reference_grads(
|
||||
_input,
|
||||
weight,
|
||||
target,
|
||||
bias=b,
|
||||
ce_weight=cw,
|
||||
reduction=reduction,
|
||||
label_smoothing=label_smoothing,
|
||||
)
|
||||
tol = _TOLS[dtype]
|
||||
assert torch.isfinite(gw).all() and torch.isfinite(gi).all()
|
||||
torch.testing.assert_close(gw.float(), ref_gw, **tol)
|
||||
torch.testing.assert_close(gi.float(), ref_gi, **tol)
|
||||
if bias:
|
||||
torch.testing.assert_close(gb.float(), ref_gb, **tol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
def test_parity_reduction_none_nonuniform_upstream(dtype):
|
||||
BT, V, H, cs = 24, 48, 32, 8 # 3 chunks (8,8,8)
|
||||
_input, weight, b, cw, target = _mk_inputs(BT, V, H, dtype, bias=True, ce_weight=False, chunk_size=cs, seed=2)
|
||||
g = torch.Generator(device=DEVICE).manual_seed(7)
|
||||
token_grad_output = torch.rand(BT, device=DEVICE, dtype=torch.float32, generator=g) + 0.05
|
||||
_, gi, gw, gb = _run_flce(
|
||||
_input,
|
||||
weight,
|
||||
target,
|
||||
bias=b,
|
||||
reduction="none",
|
||||
token_grad_output=token_grad_output,
|
||||
chunk_size=cs,
|
||||
)
|
||||
_, ref_gi, ref_gw, ref_gb = _reference_grads(
|
||||
_input,
|
||||
weight,
|
||||
target,
|
||||
bias=b,
|
||||
reduction="none",
|
||||
token_grad_output=token_grad_output,
|
||||
)
|
||||
tol = _TOLS[dtype]
|
||||
torch.testing.assert_close(gw.float(), ref_gw, **tol)
|
||||
torch.testing.assert_close(gi.float(), ref_gi, **tol)
|
||||
torch.testing.assert_close(gb.float(), ref_gb, **tol)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. Ignore-index edge cases initialise correctly (zero-init preserved)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
def test_first_chunk_all_ignored(dtype):
|
||||
BT, V, H, cs = 40, 64, 32, 16
|
||||
_input, weight, b, cw, target = _mk_inputs(
|
||||
BT, V, H, dtype, bias=True, ce_weight=False, chunk_size=cs, seed=3, first_chunk_ignore=True
|
||||
)
|
||||
_, gi, gw, gb = _run_flce(_input, weight, target, bias=b, chunk_size=cs)
|
||||
_, ref_gi, ref_gw, ref_gb = _reference_grads(_input, weight, target, bias=b)
|
||||
tol = _TOLS[dtype]
|
||||
assert torch.isfinite(gw).all()
|
||||
# First chunk was fully ignored -> its rows contribute zero grad_input.
|
||||
assert torch.count_nonzero(gi[:cs]) == 0
|
||||
torch.testing.assert_close(gw.float(), ref_gw, **tol)
|
||||
torch.testing.assert_close(gi.float(), ref_gi, **tol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
def test_all_tokens_ignored_zero_grad(dtype):
|
||||
BT, V, H, cs = 32, 64, 32, 16
|
||||
_input, weight, b, cw, target = _mk_inputs(
|
||||
BT, V, H, dtype, bias=True, ce_weight=False, chunk_size=cs, seed=4, all_ignore=True
|
||||
)
|
||||
_, gi, gw, gb = _run_flce(_input, weight, target, bias=b, chunk_size=cs)
|
||||
# Every token ignored -> all grads are exactly the zero-initialised buffers.
|
||||
assert torch.isfinite(gw).all()
|
||||
assert torch.count_nonzero(gw) == 0
|
||||
assert torch.count_nonzero(gi) == 0
|
||||
assert torch.count_nonzero(gb) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. Frozen weight + trainable input -> no grad_weight, no addmm into weight
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_frozen_weight_trainable_input():
|
||||
BT, V, H, cs = 40, 64, 32, 16
|
||||
_input, weight, b, cw, target = _mk_inputs(
|
||||
BT,
|
||||
V,
|
||||
H,
|
||||
torch.bfloat16,
|
||||
bias=False,
|
||||
ce_weight=False,
|
||||
chunk_size=cs,
|
||||
seed=5,
|
||||
weight_requires_grad=False,
|
||||
input_requires_grad=True,
|
||||
)
|
||||
rec = _OpRecorder()
|
||||
_, gi, gw, gb = _run_flce(_input, weight, target, chunk_size=cs, recorder=rec)
|
||||
assert gw is None
|
||||
assert not rec.has_addmm_out_into(weight.shape)
|
||||
assert not rec.has_addmm_dtype_out_into(weight.shape)
|
||||
# grad_input must still be produced and correct.
|
||||
_, ref_gi, _, _ = _reference_grads(_input, weight, target)
|
||||
torch.testing.assert_close(gi.float(), ref_gi, **_TOLS[torch.bfloat16])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. Non-contiguous input / weight
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
def test_noncontiguous_input_and_weight(dtype):
|
||||
BT, V, H, cs = 40, 64, 32, 16
|
||||
g = torch.Generator(device=DEVICE).manual_seed(6)
|
||||
_input = (torch.randn(H, BT, device=DEVICE, dtype=dtype, generator=g) * 0.1).t().requires_grad_(True)
|
||||
weight = (torch.randn(H, V, device=DEVICE, dtype=dtype, generator=g) * 0.1).t().requires_grad_(True)
|
||||
assert not _input.is_contiguous() and not weight.is_contiguous()
|
||||
target = torch.randint(0, V, (BT,), device=DEVICE, generator=g, dtype=torch.long)
|
||||
_, gi, gw, gb = _run_flce(_input, weight, target, chunk_size=cs)
|
||||
_, ref_gi, ref_gw, _ = _reference_grads(_input, weight, target)
|
||||
tol = _TOLS[dtype]
|
||||
assert torch.isfinite(gw).all()
|
||||
torch.testing.assert_close(gw.float(), ref_gw, **tol)
|
||||
torch.testing.assert_close(gi.float(), ref_gi, **tol)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 6. Explicit fp32 accumulator stays on the out_dtype fast path
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@requires_sm80
|
||||
@requires_addmm_out_dtype
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
def test_fp32_accum_uses_dtype_out_fastpath(dtype):
|
||||
BT, V, H, cs = 40, 64, 32, 16
|
||||
_input, weight, b, cw, target = _mk_inputs(BT, V, H, dtype, bias=False, ce_weight=False, chunk_size=cs, seed=8)
|
||||
rec = _OpRecorder()
|
||||
_, gi, gw, gb = _run_flce(_input, weight, target, accum_dtype=torch.float32, chunk_size=cs, recorder=rec)
|
||||
# Accumulator is fp32; the returned grad is cast back to the weight dtype.
|
||||
assert gw.dtype == weight.dtype
|
||||
# The fp32 out_dtype overload must be used, and the plain out overload must not.
|
||||
assert rec.has_addmm_dtype_out_into(weight.shape), (
|
||||
f"expected aten.addmm.dtype_out into {tuple(weight.shape)}, saw {rec.addmm_dtype_out_shapes}"
|
||||
)
|
||||
assert not rec.has_addmm_out_into(weight.shape)
|
||||
# No parameter-sized bf16/fp16 -> fp32 legacy temporary.
|
||||
assert not rec.has_param_upcast(weight.shape)
|
||||
_, ref_gi, ref_gw, _ = _reference_grads(_input, weight, target)
|
||||
torch.testing.assert_close(gw.float(), ref_gw, atol=1e-2, rtol=1e-2)
|
||||
torch.testing.assert_close(gi.float(), ref_gi, **_TOLS[dtype])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 7. Pre-2.8 torch (no out_dtype support) -> legacy fallback, never the out_dtype op
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@requires_sm80
|
||||
@pytest.mark.parametrize("accum_dtype", [None, torch.float32])
|
||||
def test_pre28_falls_back_to_legacy_mm(monkeypatch, accum_dtype):
|
||||
BT, V, H, cs = 40, 64, 32, 16
|
||||
# Under bf16 params + fp32 accumulator, the fast path needs out_dtype support.
|
||||
_input, weight, b, cw, target = _mk_inputs(
|
||||
BT, V, H, torch.bfloat16, bias=False, ce_weight=False, chunk_size=cs, seed=9
|
||||
)
|
||||
monkeypatch.setattr(flce_ops, "_ADDMM_SUPPORTS_OUT_DTYPE", False)
|
||||
rec = _OpRecorder()
|
||||
_, gi, gw, gb = _run_flce(_input, weight, target, accum_dtype=accum_dtype, chunk_size=cs, recorder=rec)
|
||||
# The out_dtype overload must never be reached when the flag is False.
|
||||
assert not rec.has_addmm_dtype_out_into(weight.shape)
|
||||
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.has_param_upcast(weight.shape)
|
||||
assert not rec.has_addmm_out_into(weight.shape)
|
||||
else:
|
||||
# bf16 accumulator does not depend on the out_dtype flag; the dtype-matched
|
||||
# low-precision addmm(out=) path is still taken (no legacy upcast).
|
||||
assert rec.has_addmm_out_into(weight.shape)
|
||||
assert not rec.has_param_upcast(weight.shape)
|
||||
_, ref_gi, ref_gw, _ = _reference_grads(_input, weight, target)
|
||||
torch.testing.assert_close(gw.float(), ref_gw, **_TOLS[torch.bfloat16])
|
||||
torch.testing.assert_close(gi.float(), ref_gi, **_TOLS[torch.bfloat16])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 8. AMP: in-place out ops are not autocast, so _input_chunk must be aligned
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@requires_sm80
|
||||
def test_amp_bf16_params_fp32_input_uses_lowprec_addmm():
|
||||
"""BF16 params + FP32 input under bf16 autocast: grad_weight is bf16 and
|
||||
grad_logits is bf16, so the dtype-matched low-precision addmm(out=) path runs.
|
||||
The fp32 _input_chunk must be cast to bf16 for the (non-autocast) out op."""
|
||||
BT, V, H, cs = 40, 64, 32, 16
|
||||
g = torch.Generator(device=DEVICE).manual_seed(10)
|
||||
_input = (torch.randn(BT, H, device=DEVICE, dtype=torch.float32, generator=g) * 0.1).requires_grad_(True)
|
||||
weight = (torch.randn(V, H, device=DEVICE, dtype=torch.bfloat16, generator=g) * 0.1).requires_grad_(True)
|
||||
target = torch.randint(0, V, (BT,), device=DEVICE, generator=g, dtype=torch.long)
|
||||
|
||||
rec = _OpRecorder()
|
||||
_, gi, gw, gb = _run_flce(_input, weight, target, chunk_size=cs, recorder=rec, autocast_dtype=torch.bfloat16)
|
||||
assert gw.dtype == torch.bfloat16
|
||||
assert rec.has_addmm_out_into(weight.shape)
|
||||
assert not rec.has_addmm_dtype_out_into(weight.shape)
|
||||
assert not rec.has_param_upcast(weight.shape)
|
||||
_, ref_gi, ref_gw, _ = _reference_grads(_input, weight, target)
|
||||
tol = _TOLS[torch.bfloat16]
|
||||
torch.testing.assert_close(gw.float(), ref_gw, **tol)
|
||||
torch.testing.assert_close(gi.float(), ref_gi, **tol)
|
||||
|
||||
|
||||
@requires_sm80
|
||||
@requires_addmm_out_dtype
|
||||
def test_amp_fp32_params_no_bias_uses_dtype_out_fastpath():
|
||||
"""FP32 params, no bias, under bf16 autocast: grad_weight is fp32 and
|
||||
grad_logits is bf16, so the out_dtype fast path runs (with _input_chunk
|
||||
left fp32 and aligned to bf16 for addmm)."""
|
||||
BT, V, H, cs = 40, 64, 32, 16
|
||||
g = torch.Generator(device=DEVICE).manual_seed(11)
|
||||
_input = (torch.randn(BT, H, device=DEVICE, dtype=torch.float32, generator=g) * 0.1).requires_grad_(True)
|
||||
weight = (torch.randn(V, H, device=DEVICE, dtype=torch.float32, generator=g) * 0.1).requires_grad_(True)
|
||||
target = torch.randint(0, V, (BT,), device=DEVICE, generator=g, dtype=torch.long)
|
||||
|
||||
rec = _OpRecorder()
|
||||
_, gi, gw, gb = _run_flce(_input, weight, target, chunk_size=cs, recorder=rec, autocast_dtype=torch.bfloat16)
|
||||
assert gw.dtype == torch.float32
|
||||
assert rec.has_addmm_dtype_out_into(weight.shape)
|
||||
assert not rec.has_addmm_out_into(weight.shape)
|
||||
_, ref_gi, ref_gw, _ = _reference_grads(_input, weight, target)
|
||||
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)
|
||||
|
||||
|
||||
def test_amp_fp32_params_fp32_bias_promotes_and_falls_back():
|
||||
"""FP32 params + FP32 bias under bf16 autocast: the (bf16) logits + (fp32)
|
||||
bias add promotes grad_logits to fp32, so neither addmm fast path applies and
|
||||
the legacy torch.mm(...).float() path must be used (preserving prior behavior).
|
||||
This guards against accidentally running an fp32 addmm for promoted logits."""
|
||||
BT, V, H, cs = 40, 64, 32, 16
|
||||
g = torch.Generator(device=DEVICE).manual_seed(12)
|
||||
_input = (torch.randn(BT, H, device=DEVICE, dtype=torch.float32, generator=g) * 0.1).requires_grad_(True)
|
||||
weight = (torch.randn(V, H, device=DEVICE, dtype=torch.float32, generator=g) * 0.1).requires_grad_(True)
|
||||
bias = (torch.randn(V, device=DEVICE, dtype=torch.float32, generator=g) * 0.1).requires_grad_(True)
|
||||
target = torch.randint(0, V, (BT,), device=DEVICE, generator=g, dtype=torch.long)
|
||||
|
||||
rec = _OpRecorder()
|
||||
_, gi, gw, gb = _run_flce(
|
||||
_input, weight, target, bias=bias, chunk_size=cs, recorder=rec, autocast_dtype=torch.bfloat16
|
||||
)
|
||||
# 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
|
||||
_, ref_gi, ref_gw, ref_gb = _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)
|
||||
torch.testing.assert_close(gb.float(), ref_gb, atol=1e-2, rtol=1e-2)
|
||||
|
||||
|
||||
@requires_sm80
|
||||
def test_amp_use_token_scaling_lowprec():
|
||||
"""Scaled-loss (use_token_scaling) under bf16 params still routes through the
|
||||
dtype-matched low-precision addmm(out=) path and stays numerically finite."""
|
||||
BT, V, H, cs = 24, 48, 32, 8
|
||||
_input, weight, b, cw, target = _mk_inputs(
|
||||
BT, V, H, torch.bfloat16, bias=True, ce_weight=False, chunk_size=cs, seed=13
|
||||
)
|
||||
rec = _OpRecorder()
|
||||
_, gi, gw, gb = _run_flce(_input, weight, target, bias=b, chunk_size=cs, use_token_scaling=True, recorder=rec)
|
||||
assert gw.dtype == torch.bfloat16
|
||||
assert rec.has_addmm_out_into(weight.shape)
|
||||
assert not rec.has_param_upcast(weight.shape)
|
||||
assert torch.isfinite(gw).all() and torch.isfinite(gi).all()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 9. Capability guard: sub-SM80 reported capability forces the legacy fallback
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sm_capability_guard_forces_legacy_fallback(monkeypatch):
|
||||
"""Guard regression -- intentionally *not* SM80-gated. Even on real SM80+
|
||||
hardware, production gates the addmm fast paths on ``get_device_capability()
|
||||
[0] >= 8``; if the reported capability is < 8 the fast paths must not run.
|
||||
Monkeypatch the capability to (7, 0) and assert the legacy
|
||||
``torch.mm(...).float()`` fallback is taken (bf16 params, which would
|
||||
otherwise hit the low-precision addmm(out=) path)."""
|
||||
BT, V, H, cs = 40, 64, 32, 16
|
||||
_input, weight, b, cw, target = _mk_inputs(
|
||||
BT, V, H, torch.bfloat16, bias=False, ce_weight=False, chunk_size=cs, seed=14
|
||||
)
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (7, 0))
|
||||
rec = _OpRecorder()
|
||||
_, gi, gw, gb = _run_flce(_input, weight, target, chunk_size=cs, recorder=rec)
|
||||
# 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.has_param_upcast(weight.shape)
|
||||
_, ref_gi, ref_gw, _ = _reference_grads(_input, weight, target)
|
||||
torch.testing.assert_close(gw.float(), ref_gw, **_TOLS[torch.bfloat16])
|
||||
torch.testing.assert_close(gi.float(), ref_gi, **_TOLS[torch.bfloat16])
|
||||
@@ -202,6 +202,19 @@ class TestArgValidation:
|
||||
with pytest.raises(SystemExit):
|
||||
self._parse(["--warmup", "-1"])
|
||||
|
||||
def test_accum_dtype_defaults_to_fp32(self):
|
||||
# Backward compatible: omitting the flag keeps the FP32-matched accumulator.
|
||||
assert self._parse([]).accum_dtype == "fp32"
|
||||
|
||||
@pytest.mark.parametrize("mode", ["fp32", "default"])
|
||||
def test_accum_dtype_choices_accepted(self, mode):
|
||||
assert self._parse(["--accum-dtype", mode]).accum_dtype == mode
|
||||
|
||||
@pytest.mark.parametrize("bad", ["bf16", "none", "float32", "auto"])
|
||||
def test_accum_dtype_rejects_unknown(self, bad):
|
||||
with pytest.raises(SystemExit):
|
||||
self._parse(["--accum-dtype", bad])
|
||||
|
||||
|
||||
class TestTimeBackendLifecycle:
|
||||
def _patch_cuda(self, monkeypatch, on_baseline=None):
|
||||
@@ -278,6 +291,7 @@ class TestRunFailFast:
|
||||
"warmup": 0,
|
||||
"iters": 1,
|
||||
"seed": 0,
|
||||
"accum_dtype": "fp32",
|
||||
},
|
||||
)()
|
||||
|
||||
@@ -314,6 +328,132 @@ class TestRunFailFast:
|
||||
assert "eff_chunk" in out
|
||||
|
||||
|
||||
class _SlotRecFn(torch.autograd.Function):
|
||||
"""Records the value passed at FLCE Function slot 11 (``accum_dtype``)."""
|
||||
|
||||
seen = []
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, x, w, target, *rest):
|
||||
# rest = (bias, ce_weight, ignore_index, lse_square_scale, label_smoothing,
|
||||
# reduction, softcap, return_z_loss, accum_dtype, ...); slot 11 overall
|
||||
# is accum_dtype, i.e. rest[8].
|
||||
_SlotRecFn.seen.append(rest[8])
|
||||
ctx.save_for_backward(x, w)
|
||||
loss = (x.float() @ w.float().t()).sum()
|
||||
return loss, None, None, None
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_loss, *rest):
|
||||
x, w = ctx.saved_tensors
|
||||
return (torch.ones_like(x) * grad_loss, torch.ones_like(w) * grad_loss) + (None,) * 16
|
||||
|
||||
|
||||
class TestAccumDtypeThreading:
|
||||
def test_effective_accum_labels(self):
|
||||
# fp32 forces FP32 everywhere; default keeps triton/cutedsl BF16 but cutile FP32.
|
||||
for name in ("triton", "cutedsl", "cutile"):
|
||||
assert bench._effective_accum(name, "fp32") == "fp32"
|
||||
assert bench._effective_accum("triton", "default") == "bf16"
|
||||
assert bench._effective_accum("cutedsl", "default") == "bf16"
|
||||
assert bench._effective_accum("cutile", "default") == "fp32"
|
||||
|
||||
def test_apply_forwards_dtype_to_slot_11(self):
|
||||
x, w, target = _cpu_inputs()
|
||||
_SlotRecFn.seen = []
|
||||
bench._apply(_SlotRecFn, x, w, target, None, accum_dtype=torch.float32)
|
||||
bench._apply(_SlotRecFn, x, w, target, None, accum_dtype=None)
|
||||
assert _SlotRecFn.seen == [torch.float32, None]
|
||||
|
||||
def test_apply_default_is_fp32(self):
|
||||
# Drop-in default preserves the earlier hard-coded FP32 accumulator.
|
||||
x, w, target = _cpu_inputs()
|
||||
_SlotRecFn.seen = []
|
||||
bench._apply(_SlotRecFn, x, w, target, None)
|
||||
assert _SlotRecFn.seen == [torch.float32]
|
||||
|
||||
def _patch_cuda(self, monkeypatch):
|
||||
_FakeEvent._clock = 0.0
|
||||
monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None)
|
||||
monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", lambda *a, **k: None)
|
||||
monkeypatch.setattr(torch.cuda, "memory_allocated", lambda *a, **k: 1000)
|
||||
monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda *a, **k: 5000)
|
||||
monkeypatch.setattr(torch.cuda, "Event", lambda *a, **k: _FakeEvent())
|
||||
|
||||
@pytest.mark.parametrize("accum", [torch.float32, None])
|
||||
def test_both_modes_roundtrip_through_iterations(self, monkeypatch, accum):
|
||||
# Every timed/warmup/probe iteration threads the same accum dtype to slot 11.
|
||||
x, w, target = _cpu_inputs()
|
||||
self._patch_cuda(monkeypatch)
|
||||
_SlotRecFn.seen = []
|
||||
bench._time_backend(_SlotRecFn, x, w, target, None, warmup=2, iters=3, accum_dtype=accum)
|
||||
# 1 compile + 2 warmup + 1 memory-probe + 3 timed = 7 applies, all identical.
|
||||
assert len(_SlotRecFn.seen) == 7
|
||||
assert all(v is accum for v in _SlotRecFn.seen)
|
||||
|
||||
def _run_args(self, backends, accum_dtype):
|
||||
return type(
|
||||
"Args",
|
||||
(),
|
||||
{
|
||||
"backends": backends,
|
||||
"tokens": 8,
|
||||
"hidden_size": 4,
|
||||
"vocab_size": 6,
|
||||
"chunk_sizes": [4],
|
||||
"warmup": 0,
|
||||
"iters": 1,
|
||||
"seed": 0,
|
||||
"accum_dtype": accum_dtype,
|
||||
},
|
||||
)()
|
||||
|
||||
def _patch_run(self, monkeypatch):
|
||||
monkeypatch.setattr(bench, "_capability", lambda: (10, 0))
|
||||
monkeypatch.setattr(bench, "_check_backend_supported", lambda *a, **k: None)
|
||||
monkeypatch.setattr(bench, "_load_backend", lambda name: _CpuFLCE)
|
||||
monkeypatch.setattr(bench, "_build_inputs", lambda *a, **k: _cpu_inputs())
|
||||
monkeypatch.setattr(torch.cuda, "get_device_name", lambda *a, **k: "MockGPU")
|
||||
self._patch_cuda(monkeypatch)
|
||||
|
||||
def test_default_mode_prints_requested_and_per_backend_effective(self, monkeypatch, capsys):
|
||||
self._patch_run(monkeypatch)
|
||||
bench._run(self._run_args(["triton", "cutedsl", "cutile"], "default"))
|
||||
out = capsys.readouterr().out
|
||||
# Header advertises the requested mode and refuses to call it precision-matched.
|
||||
assert "accum_dtype_requested=default" in out
|
||||
assert "NOT precision matched" in out
|
||||
assert "eff_accum" in out
|
||||
rows = {
|
||||
line.split()[0]: line
|
||||
for line in out.splitlines()
|
||||
if line.split() and line.split()[0] in bench._MODULE_PATHS
|
||||
}
|
||||
# cutile keeps FP32 even under default; triton/cutedsl drop to BF16.
|
||||
assert "fp32" in rows["cutile"].split()
|
||||
assert "bf16" in rows["triton"].split()
|
||||
assert "bf16" in rows["cutedsl"].split()
|
||||
|
||||
def test_fp32_mode_labels_precision_matched(self, monkeypatch, capsys):
|
||||
self._patch_run(monkeypatch)
|
||||
bench._run(self._run_args(["triton", "cutile"], "fp32"))
|
||||
out = capsys.readouterr().out
|
||||
assert "accum_dtype_requested=fp32" in out
|
||||
assert "precision matched" in out
|
||||
assert "NOT precision matched" not in out
|
||||
for line in out.splitlines():
|
||||
if line.split() and line.split()[0] in bench._MODULE_PATHS:
|
||||
assert "fp32" in line.split()
|
||||
|
||||
def test_run_does_not_mutate_module_globals(self, monkeypatch):
|
||||
self._patch_run(monkeypatch)
|
||||
before = bench._ACCUM_DTYPE
|
||||
bench._run(self._run_args(["triton"], "default"))
|
||||
bench._run(self._run_args(["triton"], "fp32"))
|
||||
# Modes flow purely as parameters; no state leak between runs.
|
||||
assert bench._ACCUM_DTYPE is before is torch.float32
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# B. GPU projection-geometry traces (whole-pipeline, not CE-only isolation)
|
||||
# ===========================================================================
|
||||
@@ -523,6 +663,7 @@ def test_benchmark_run_all3_smoke():
|
||||
"warmup": 0,
|
||||
"iters": 2,
|
||||
"seed": 0,
|
||||
"accum_dtype": "fp32",
|
||||
},
|
||||
)()
|
||||
bench._run(args) # must not raise
|
||||
|
||||
Reference in New Issue
Block a user