From 9ef3dfa665b21b2fa14adef3fa821b424d7821ba Mon Sep 17 00:00:00 2001 From: Wang Zupeng Date: Thu, 10 Sep 2026 00:27:03 +0000 Subject: [PATCH] [PyTorch] Add timestep-conditioned AdaptiveLayerNorm to op fuser Signed-off-by: Wang Zupeng --- benchmarks/benchmark_adaptive_layer_norm.py | 231 +++++++++ docs/api/pytorch.rst | 2 + docs/examples/op_fuser/op_fuser.rst | 45 ++ qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/test_adaptive_layer_norm.py | 450 ++++++++++++++++++ .../common/triton/adaptive_layer_norm.py | 157 ++++++ .../pytorch/ops/basic/__init__.py | 1 + .../pytorch/ops/basic/adaptive_layer_norm.py | 155 ++++++ .../pytorch/triton/adaptive_layer_norm.py | 252 ++++++++++ 9 files changed, 1294 insertions(+) create mode 100644 benchmarks/benchmark_adaptive_layer_norm.py create mode 100644 tests/pytorch/test_adaptive_layer_norm.py create mode 100644 transformer_engine/common/triton/adaptive_layer_norm.py create mode 100644 transformer_engine/pytorch/ops/basic/adaptive_layer_norm.py create mode 100644 transformer_engine/pytorch/triton/adaptive_layer_norm.py diff --git a/benchmarks/benchmark_adaptive_layer_norm.py b/benchmarks/benchmark_adaptive_layer_norm.py new file mode 100644 index 0000000000..a8e681fa20 --- /dev/null +++ b/benchmarks/benchmark_adaptive_layer_norm.py @@ -0,0 +1,231 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Benchmark adaptive LayerNorm and an optional diffusion-transformer MLP. + +Reports CUDA graph replay latency for eager PyTorch, torch.compile, and TE. +The references accumulate normalization and modulation in float32, matching +AdaptiveLayerNorm. Shapes describe a single block, not a complete model. +""" + +import argparse +import json +import statistics +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +import transformer_engine.pytorch as te + + +def reference(x, scale, shift, eps, batch_dim): + """Normalize and modulate in float32, with one final output cast.""" + condition_shape = [1] * x.ndim + condition_shape[batch_dim] = x.shape[batch_dim] + condition_shape[-1] = x.shape[-1] + scale = scale.reshape(condition_shape) + shift = shift.reshape(condition_shape) + normalized = F.layer_norm(x.float(), (x.shape[-1],), eps=eps) + return (normalized * (1.0 + scale.float()) + shift.float()).to(x.dtype) + + +def graph_time(function, repeats): + """Measure a fixed CUDA graph, excluding compilation and graph capture.""" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(5): + function() + torch.cuda.current_stream().wait_stream(stream) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + outputs = function() + samples = [] + for _ in range(5): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(repeats): + graph.replay() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1000 / repeats) + # Keep the capture's output allocations live until all replays finish. + del outputs + return statistics.median(samples) + + +def host_time(function, repeats): + """Measure amortized execution including Python dispatch and synchronization.""" + samples = [] + for _ in range(5): + torch.cuda.synchronize() + start = time.perf_counter() + for _ in range(repeats): + function() + torch.cuda.synchronize() + samples.append((time.perf_counter() - start) * 1e6 / repeats) + return statistics.median(samples) + + +def peak_memory(function): + """Measure extra PyTorch-managed allocation above persistent inputs/caches.""" + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + outputs = function() + torch.cuda.synchronize() + peak = torch.cuda.max_memory_allocated() - before + del outputs + return peak + + +def main(): + """Run correctness checks and CUDA graph benchmarks.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--sequence", type=int, default=1024) + parser.add_argument("--hidden", type=int, default=1536) + parser.add_argument("--batch-dim", type=int, choices=(0, 1), default=1) + parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="bfloat16") + parser.add_argument( + "--ffn", type=int, default=0, help="Add an MLP with this intermediate size." + ) + parser.add_argument("--repeats", type=int, default=100) + parser.add_argument("--skip-compile", action="store_true") + parser.add_argument("--reverse-order", action="store_true") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + torch.manual_seed(1234) + dtype = getattr(torch, args.dtype) + shape = [args.sequence, args.sequence, args.hidden] + shape[args.batch_dim] = args.batch + condition_shape = [args.batch, args.hidden] + x = torch.randn(shape, device="cuda", dtype=dtype, requires_grad=True) + scale = torch.randn(condition_shape, device="cuda", requires_grad=True) + shift = torch.randn(condition_shape, device="cuda", requires_grad=True) + dy = torch.randn_like(x) + eps = 1e-6 + op = te.ops.AdaptiveLayerNorm(args.hidden, eps=eps, batch_dim=args.batch_dim) + + suffix = None + if args.ffn: + suffix = te.ops.Sequential( + te.ops.Linear(args.hidden, args.ffn, dtype=dtype), + te.ops.GELU(), + te.ops.Linear(args.ffn, args.hidden, dtype=dtype), + ) + + def eager(a, b, c): + return reference(a, b, c, eps, args.batch_dim) + + methods = {"pytorch_eager": eager, "transformer_engine": op} + if not args.skip_compile: + methods["pytorch_compile"] = torch.compile(eager, fullgraph=True) + if args.reverse_order: + methods = dict(reversed(methods.items())) + + # Compare against the same mathematical reference before measuring. + expected = eager(x, scale, shift) + expected_grads = torch.autograd.grad(expected, (x, scale, shift), dy) + expected = expected.detach() + results = [] + for name, norm in methods.items(): + actual = norm(x, scale, shift) + gradients = torch.autograd.grad(actual, (x, scale, shift), dy) + atol, rtol = (2e-2, 2e-2) if dtype == torch.bfloat16 else (2e-3, 2e-3) + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + for actual_grad, expected_grad in zip(gradients, expected_grads): + # Conditions accumulate over sequence positions in float32. + torch.testing.assert_close(actual_grad, expected_grad, atol=atol, rtol=rtol) + + # Release validation graphs before creating AccumulateGrad nodes on the + # capture stream. Retaining them would introduce a default-stream edge. + del actual + parameters = () if suffix is None else tuple(suffix.parameters()) + if suffix is None: + model = norm + elif name == "transformer_engine": + model = te.ops.Sequential(norm, suffix[0], suffix[1], suffix[2]) + else: + # Identical TE MLP operations/weights isolate the normalization change. + def model(a, b, c, normalization=norm): + return suffix(normalization(a, b, c)) + + block_gradient_checks = [] + if suffix is not None: + block_reference = suffix(eager(x, scale, shift)) + block_actual = model(x, scale, shift) + block_inputs = (x, scale, shift) + parameters + block_reference_grads = torch.autograd.grad(block_reference, block_inputs, dy) + block_actual_grads = torch.autograd.grad(block_actual, block_inputs, dy) + torch.testing.assert_close(block_actual, block_reference, atol=atol, rtol=rtol) + for actual_grad, expected_grad in zip(block_actual_grads, block_reference_grads): + # Reduced MLP gradients can be much larger than activations. + # Use one activation-dtype epsilon at the gradient RMS as the + # absolute scale, retaining the pointwise relative criterion. + ref_rms = expected_grad.float().square().mean().sqrt().item() + grad_atol = max(atol, torch.finfo(dtype).eps * ref_rms) + torch.testing.assert_close(actual_grad, expected_grad, atol=grad_atol, rtol=rtol) + error = (actual_grad.float() - expected_grad.float()).abs() + block_gradient_checks.append( + { + "shape": list(actual_grad.shape), + "reference_rms": ref_rms, + "error_rms": error.square().mean().sqrt().item(), + "max_error": error.max().item(), + "atol": grad_atol, + "rtol": rtol, + } + ) + + del block_reference, block_actual + + def forward(model_fn=model): + with torch.no_grad(): + return model_fn(x, scale, shift) + + def forward_backward(model_fn=model, model_params=parameters): + out = model_fn(x, scale, shift) + return torch.autograd.grad(out, (x, scale, shift) + model_params, dy) + + results.append( + { + "method": name, + "forward_us": graph_time(forward, args.repeats), + "forward_backward_us": graph_time(forward_backward, args.repeats), + "forward_host_us": host_time(forward, args.repeats), + "forward_backward_host_us": host_time(forward_backward, args.repeats), + "forward_backward_peak_bytes": peak_memory(forward_backward), + "correctness": "passed", + "block_gradient_checks": block_gradient_checks, + } + ) + print(json.dumps(results[-1]), flush=True) + report = { + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "shape": shape, + "condition_shape": condition_shape, + "dtype": args.dtype, + "condition_dtype": "float32", + "ffn": args.ffn, + "method_order": list(methods), + "measurement": ( + "median of 5 batches, microseconds per iteration; *_us uses CUDA graphs, " + "*_host_us includes Python dispatch" + ), + "results": results, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index a6afa2d0cc..301dcdb83f 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -230,6 +230,8 @@ Operation fuser .. autoapiclass:: transformer_engine.pytorch.ops.L2Normalization +.. autoapiclass:: transformer_engine.pytorch.ops.AdaptiveLayerNorm + .. autoapiclass:: transformer_engine.pytorch.ops.LayerNorm .. autoapiclass:: transformer_engine.pytorch.ops.MakeExtraOutput diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst index a6a500f20e..daf95084ff 100644 --- a/docs/examples/op_fuser/op_fuser.rst +++ b/docs/examples/op_fuser/op_fuser.rst @@ -81,6 +81,51 @@ Thus, using the operation fuser simply involves constructing Operations that match ``LayerNormMLP`` module. Note that different fusions have been applied in the forward and backward passes. +Adaptive normalization for diffusion transformers +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +``AdaptiveLayerNorm`` applies normalization followed by a per-sample +timestep-conditioned scale and shift. Unlike ``LayerNorm``, it has no +learnable affine parameters: the caller supplies both conditions, and the +backward pass returns their gradients. It can be composed with linear and +activation operations to construct an adaptive LayerNorm MLP: + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + + batch, sequence, hidden, ffn = 2, 1024, 1536, 6144 + mlp = te.ops.Sequential( + te.ops.AdaptiveLayerNorm(hidden, eps=1e-6), + te.ops.Linear(hidden, ffn, dtype=torch.bfloat16), + te.ops.GELU(), + te.ops.Linear(ffn, hidden, dtype=torch.bfloat16), + ) + x = torch.randn( + batch, sequence, hidden, device="cuda", + dtype=torch.bfloat16, requires_grad=True, + ) + # These tensors can also be outputs of a timestep embedding network. + scale = torch.randn(batch, 1, hidden, device="cuda", requires_grad=True) + shift = torch.randn_like(scale, requires_grad=True) + y = mlp(x, scale, shift) + y.float().square().mean().backward() + +For sequence-first input ``[S, B, H]``, set ``batch_dim=1`` and pass +conditions with shape ``[1, B, H]`` or ``[B, H]``. Conditions are shared +across the other dimensions and do not need to be expanded to the activation +shape. Non-contiguous conditions, such as views returned by ``chunk``, are +supported. + +Normalization and ``1 + scale`` modulation use float32 arithmetic, even when +the input or conditions use bfloat16 or float16. The output has the input +dtype. This avoids rounding small scales away by adding one in bfloat16. +The normalization and modulation run in one Triton kernel; condition gradients +use deterministic reductions with a bounded workspace. A subsequent linear +operation retains the operation fuser's matching patterns, but this operation +does not fuse normalization with GEMM or output quantization. + Quantization ^^^^^^^^^^^^ diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index a78a99d7f9..8e046c518e 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -52,6 +52,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_qk_norm.xml $TE_ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_optimizer.xml $TE_PATH/tests/pytorch/test_fused_optimizer.py || test_fail "test_fused_optimizer.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_multi_tensor.xml $TE_PATH/tests/pytorch/test_multi_tensor.py || test_fail "test_multi_tensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/test_fusible_ops.py || test_fail "test_fusible_ops.py" +NVIDIA_TF32_OVERRIDE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_adaptive_layer_norm.xml $TE_PATH/tests/pytorch/test_adaptive_layer_norm.py || test_fail "test_adaptive_layer_norm.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_selective_activation_checkpoint.xml $TE_PATH/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py || test_fail "test_selective_activation_checkpoint.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_distributed_weight.xml $TE_PATH/tests/pytorch/test_distributed_weight.py || test_fail "test_distributed_weight.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_override.xml $TE_PATH/tests/pytorch/test_backward_override.py || test_fail "test_backward_override.py" diff --git a/tests/pytorch/test_adaptive_layer_norm.py b/tests/pytorch/test_adaptive_layer_norm.py new file mode 100644 index 0000000000..32fa5677dd --- /dev/null +++ b/tests/pytorch/test_adaptive_layer_norm.py @@ -0,0 +1,450 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Adaptive LayerNorm numerics and composition with fusible operations. + +Run with NVIDIA_TF32_OVERRIDE=0 to compare TE and PyTorch GEMMs in full FP32, +as configured in qa/L0_pytorch_unittest/test.sh. +""" + +from __future__ import annotations + +import copy + +import pytest +import torch +import torch.nn.functional as F + +import transformer_engine.pytorch as te +import transformer_engine.pytorch.ops as te_ops +from transformer_engine.common.recipe import DelayedScaling, Float8CurrentScaling + + +_DTYPES = (torch.float32, torch.float16, torch.bfloat16) + + +@pytest.fixture(autouse=True) +def _seed_rng(): + torch.manual_seed(1234) + + +def _tolerances(dtype): + # The reference accumulates in FP64; the operation accumulates in FP32. + if dtype == torch.float16: + return {"rtol": 2e-3, "atol": 2e-3} + if dtype == torch.bfloat16: + return {"rtol": 1.6e-2, "atol": 2e-2} + return {"rtol": 2e-5, "atol": 2e-5} + + +def _assert_close(actual, expected): + torch.testing.assert_close(actual, expected.to(actual.dtype), **_tolerances(actual.dtype)) + + +def _condition_view(condition, input_shape, batch_dim): + if condition.ndim == 2 and len(input_shape) > 2: + shape = [1] * len(input_shape) + shape[batch_dim] = input_shape[batch_dim] + shape[-1] = input_shape[-1] + return condition.reshape(shape) + return condition + + +def _reference(x, scale, shift, *, eps=1e-5, batch_dim=0): + """Mathematical reference with FP64 statistics and modulation.""" + x = x.double() + scale = _condition_view(scale.double(), x.shape, batch_dim) + shift = _condition_view(shift.double(), x.shape, batch_dim) + normalized = F.layer_norm(x, (x.shape[-1],), eps=eps) + return normalized * (1 + scale) + shift + + +def _make_inputs(shape, batch_dim, dtype, *, cond_dtype=None, broadcast=False): + if cond_dtype is None: + cond_dtype = dtype + x = torch.randn(shape, device="cuda", dtype=dtype).requires_grad_() + batch_size, hidden_size = shape[batch_dim], shape[-1] + scale = torch.randn(batch_size, hidden_size, device="cuda", dtype=cond_dtype) + shift = torch.randn_like(scale) + if broadcast and len(shape) > 2: + cond_shape = [1] * len(shape) + cond_shape[batch_dim] = batch_size + cond_shape[-1] = hidden_size + scale = scale.reshape(cond_shape) + shift = shift.reshape(cond_shape) + return x, scale.requires_grad_(), shift.requires_grad_() + + +def _check_against_reference(inputs, *, batch_dim=0, eps=1e-5, op=None): + x, scale, shift = inputs + reference_inputs = tuple( + tensor.detach().double().requires_grad_(tensor.requires_grad) for tensor in inputs + ) + if op is None: + op = te_ops.AdaptiveLayerNorm(x.shape[-1], eps=eps, batch_dim=batch_dim) + actual = op(x, scale, shift) + expected = _reference(*reference_inputs, eps=eps, batch_dim=batch_dim) + assert actual.shape == x.shape + assert actual.dtype == x.dtype + _assert_close(actual, expected) + if any(tensor.requires_grad for tensor in inputs): + grad = torch.randn_like(actual) + actual.backward(grad) + expected.backward(grad.double()) + for tensor, ref_tensor in zip(inputs, reference_inputs): + if tensor.requires_grad: + assert tensor.grad is not None + assert tensor.grad.shape == tensor.shape + assert tensor.grad.dtype == tensor.dtype + _assert_close(tensor.grad, ref_tensor.grad) + else: + assert tensor.grad is None + else: + assert not actual.requires_grad + return actual + + +# Cover distinct layout/reduction cases, including non-power-of-two hidden +# sizes and the kernel's smallest/largest advertised hidden dimensions. +_LAYOUTS = ( + ((3, 1), 0, False), + ((3, 31), 0, False), + ((2, 7, 128), 0, False), + ((2, 7, 128), 0, True), + ((7, 2, 128), 1, False), + ((7, 2, 128), 1, True), + ((2, 3, 5, 63), 0, True), + ((3, 2, 5, 63), 1, False), + ((3, 5, 2, 63), 2, True), + ((2, 7, 1536), 0, False), + ((7, 2, 5120), 1, True), + ((2, 3, 16384), 0, False), +) + + +@pytest.mark.parametrize("dtype", _DTYPES) +@pytest.mark.parametrize("shape,batch_dim,broadcast", _LAYOUTS) +def test_forward_backward(shape, batch_dim, broadcast, dtype): + inputs = _make_inputs(shape, batch_dim, dtype, broadcast=broadcast) + _check_against_reference(inputs, batch_dim=batch_dim) + if shape[-1] == 1: + assert torch.count_nonzero(inputs[0].grad).item() == 0 + assert torch.count_nonzero(inputs[1].grad).item() == 0 + + +@pytest.mark.parametrize( + "dtype,cond_dtype", + ( + (torch.float16, torch.float32), + (torch.bfloat16, torch.float32), + (torch.float32, torch.float16), + (torch.float32, torch.bfloat16), + ), +) +@pytest.mark.parametrize("batch_dim", (0, 1)) +def test_mixed_condition_dtype(dtype, cond_dtype, batch_dim): + shape = (3, 11, 256) if batch_dim == 0 else (11, 3, 256) + inputs = _make_inputs(shape, batch_dim, dtype, cond_dtype=cond_dtype, broadcast=True) + _check_against_reference(inputs, batch_dim=batch_dim) + + +@pytest.mark.parametrize( + "requires_grad", + ( + (True, False, False), + (False, True, True), + (False, True, False), + (False, False, True), + (False, False, False), + ), +) +def test_gradient_requirements(requires_grad): + inputs = _make_inputs((3, 11, 64), 0, torch.float32) + for tensor, required in zip(inputs, requires_grad): + tensor.requires_grad_(required) + _check_against_reference(inputs) + + +@pytest.mark.parametrize("batch_dim", (0, 1)) +@pytest.mark.parametrize("condition_layout", ("chunk", "strided")) +def test_noncontiguous_inputs(batch_dim, condition_layout): + shape = (3, 11, 64) if batch_dim == 0 else (11, 3, 64) + # Non-contiguous hidden dimension, not just a transpose of token axes. + x = torch.randn(*shape[:-1], 128, device="cuda")[..., ::2] + x = x.detach().requires_grad_() + if condition_layout == "chunk": + packed = torch.randn(3, 128, device="cuda") + scale, shift = packed.chunk(2, dim=-1) + else: + packed = torch.randn(3, 256, device="cuda") + scale, shift = packed[:, :128:2], packed[:, 128::2] + scale = scale.detach().requires_grad_() + shift = shift.detach().requires_grad_() + assert not x.is_contiguous() + assert not scale.is_contiguous() + assert not shift.is_contiguous() + _check_against_reference((x, scale, shift), batch_dim=batch_dim) + + +@pytest.mark.parametrize("eps", (1e-6, 1e-3)) +def test_nearly_constant_input(eps): + inputs = _make_inputs((2, 7, 96), 0, torch.float32) + with torch.no_grad(): + inputs[0].mul_(1e-4).add_(1) + dy = torch.randn_like(inputs[0]) + output = te_ops.AdaptiveLayerNorm(96, eps=eps)(*inputs) + actual = (output, *torch.autograd.grad(output, inputs, dy)) + + references = [] + for dtype in (torch.float64, torch.float32): + x, scale, shift = (tensor.detach().to(dtype).requires_grad_() for tensor in inputs) + y = F.layer_norm(x, (96,), eps=eps) + y = y * (1 + scale[:, None, :]) + shift[:, None, :] + references.append((y, *torch.autograd.grad(y, (x, scale, shift), dy.to(dtype)))) + + # The small variance amplifies FP32 mean rounding around one. Compare to + # FP64, allowing the error scale of native PyTorch's FP32 LayerNorm rather + # than requiring FP64 statistics from either FP32 implementation. + for result, exact, native in zip(actual, references[0], references[1]): + native_error = (native.double() - exact).abs().max() + rounding_floor = 32 * torch.finfo(torch.float32).eps * exact.abs().max().clamp_min(1) + atol = torch.maximum(2 * native_error, rounding_floor).item() + torch.testing.assert_close(result, exact.to(result.dtype), rtol=2e-5, atol=atol) + + +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +def test_small_zero_centered_scale(dtype): + # Forming (1 + scale) in the input dtype would round it back to one. + # Keep x in FP32 so this lost modulation remains visible in the output. + x = torch.tensor([[[-1.0, 1.0, -1.0, 1.0]]], device="cuda", requires_grad=True) + value = 2**-12 if dtype == torch.float16 else 2**-10 + scale = torch.full((1, 4), value, device="cuda", dtype=dtype, requires_grad=True) + shift = torch.zeros_like(scale, requires_grad=True) + actual = _check_against_reference((x, scale, shift)) + rounded_scale_result = F.layer_norm(x, (4,)) * (1 + scale).float().unsqueeze(1) + assert (actual - rounded_scale_result).abs().max().item() > value / 2 + + +@pytest.mark.parametrize("dtype", _DTYPES) +@pytest.mark.parametrize("batch_dim", (0, 1)) +def test_sequential_upstream_and_downstream_gradients(dtype, batch_dim, monkeypatch): + """The parameter-free op must preserve gradients through upstream Linear.""" + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + first = te_ops.Linear(32, 64, device="cuda", dtype=dtype) + last = te_ops.Linear(64, 16, device="cuda", dtype=dtype) + adaptive = te_ops.AdaptiveLayerNorm(64, batch_dim=batch_dim) + model = te_ops.Sequential(first, adaptive, last) + shape = (3, 7, 32) if batch_dim == 0 else (7, 3, 32) + x = torch.randn(shape, device="cuda", dtype=dtype, requires_grad=True) + scale = torch.randn(3, 64, device="cuda", dtype=torch.float32, requires_grad=True) + shift = torch.randn_like(scale, requires_grad=True) + tensors = (x, scale, shift, first.weight, first.bias, last.weight, last.bias) + rx, rs, rb = (tensor.detach().clone().requires_grad_() for tensor in (x, scale, shift)) + ref_first, ref_last = copy.deepcopy(first), copy.deepcopy(last) + reference_tensors = ( + rx, + rs, + rb, + ref_first.weight, + ref_first.bias, + ref_last.weight, + ref_last.bias, + ) + # Keep the GEMM implementation and rounding boundaries identical. TE and + # PyTorch BF16 GEMMs may differ by one ULP, amplified by later weight grads. + # Normalization/modulation still use the independent FP64 reference. + expected = ref_first(rx) + expected = _reference(expected, rs, rb, batch_dim=batch_dim).to(dtype) + expected = ref_last(expected) + actual = model(x, scale, shift) + _assert_close(actual, expected) + grad = torch.randn_like(actual) + actual.backward(grad) + expected.backward(grad) + for tensor, reference_tensor in zip(tensors, reference_tensors): + assert tensor.grad is not None + assert reference_tensor.grad is not None + # BF16 linear outputs are rounded before normalization in both paths. + # Dcondition is FP32 but inherits the activation precision of the MLP. + tolerances = _tolerances(dtype) + torch.testing.assert_close(tensor.grad, reference_tensor.grad, **tolerances) + + +def test_conditions_change_between_calls(): + op = te_ops.AdaptiveLayerNorm(64) + assert list(op.parameters()) == [] + for batch_size, tokens in ((2, 7), (3, 11)): + inputs = _make_inputs((batch_size, tokens, 64), 0, torch.float32) + _check_against_reference(inputs, op=op) + assert list(op.parameters()) == [] + + +@pytest.mark.parametrize( + "shape,batch_dim,broadcast", + ( + ((0, 7, 64), 0, False), + ((3, 0, 64), 0, True), + ((0, 3, 64), 1, False), + ((7, 0, 64), 1, True), + ), +) +def test_empty_input(shape, batch_dim, broadcast): + inputs = _make_inputs(shape, batch_dim, torch.float32, broadcast=broadcast) + actual = te_ops.AdaptiveLayerNorm(64, batch_dim=batch_dim)(*inputs) + assert actual.shape == shape + actual.sum().backward() + for tensor in inputs: + assert tensor.grad is not None + assert tensor.grad.shape == tensor.shape + assert torch.count_nonzero(tensor.grad).item() == 0 + + +@pytest.mark.parametrize("hidden_size", (0, -1, 16385)) +def test_invalid_hidden_size(hidden_size): + with pytest.raises((TypeError, ValueError)): + te_ops.AdaptiveLayerNorm(hidden_size) + + +@pytest.mark.parametrize("eps", (-1e-5, float("nan"), float("inf"))) +def test_invalid_eps(eps): + with pytest.raises((TypeError, ValueError)): + te_ops.AdaptiveLayerNorm(64, eps=eps) + + +@pytest.mark.parametrize("batch_dim", (2, 3, -4)) +def test_invalid_batch_axis(batch_dim): + inputs = _make_inputs((3, 7, 64), 0, torch.float32) + with pytest.raises((TypeError, ValueError)): + te_ops.AdaptiveLayerNorm(64, batch_dim=batch_dim)(*inputs) + + +@pytest.mark.parametrize( + "input_shape,scale_shape,shift_shape", + ( + ((64,), (1, 64), (1, 64)), + ((3, 7, 32), (3, 64), (3, 64)), + ((3, 7, 64), (2, 64), (3, 64)), + ((3, 7, 64), (3, 64), (2, 64)), + ((3, 7, 64), (3, 32), (3, 64)), + ((3, 7, 64), (3, 7, 64), (3, 1, 64)), + ((3, 7, 64), (64,), (3, 64)), + ((3, 7, 64), (3, 1, 1, 64), (3, 64)), + ), +) +def test_invalid_tensor_shapes(input_shape, scale_shape, shift_shape): + tensors = tuple( + torch.randn(shape, device="cuda") for shape in (input_shape, scale_shape, shift_shape) + ) + with pytest.raises((TypeError, ValueError)): + te_ops.AdaptiveLayerNorm(64)(*tensors) + + +@pytest.mark.parametrize("argument", (0, 1, 2)) +def test_invalid_tensor_dtype(argument): + inputs = list(_make_inputs((3, 7, 64), 0, torch.float32)) + inputs[argument] = inputs[argument].detach().to(torch.int32) + with pytest.raises((TypeError, ValueError)): + te_ops.AdaptiveLayerNorm(64)(*inputs) + + +@pytest.mark.parametrize("dtype", _DTYPES) +def test_autocast_preserves_input_dtype(dtype): + inputs = _make_inputs((2, 7, 64), 0, dtype, cond_dtype=torch.float32) + with torch.autocast("cuda", dtype=torch.bfloat16): + _check_against_reference(inputs) + + +@pytest.mark.parametrize("dtype", (torch.float32, torch.bfloat16)) +def test_condition_projection_gradients(dtype, monkeypatch): + """Condition gradients must reach the trainable timestep projection.""" + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + hidden_size = 64 + projection = te_ops.Linear(16, 2 * hidden_size, device="cuda", dtype=dtype) + conditioner = te_ops.Sequential(projection) + op = te_ops.AdaptiveLayerNorm(hidden_size, batch_dim=1) + # Neither data input requires a gradient: projection parameters are the + # only reason the adaptive operation participates in backward. + x = torch.randn(7, 3, hidden_size, device="cuda", dtype=dtype) + timestep = torch.randn(3, 16, device="cuda", dtype=dtype) + scale, shift = conditioner(timestep).chunk(2, dim=-1) + actual = op(x, scale, shift) + ref_weight = projection.weight.detach().clone().requires_grad_() + ref_bias = projection.bias.detach().clone().requires_grad_() + ref_scale, ref_shift = F.linear(timestep, ref_weight, ref_bias).chunk(2, dim=-1) + expected = _reference(x, ref_scale, ref_shift, batch_dim=1).to(dtype) + _assert_close(actual, expected) + grad = torch.randn_like(actual) + actual.backward(grad) + expected.backward(grad) + assert x.grad is None + assert timestep.grad is None + for param, ref_param in ( + (projection.weight, ref_weight), + (projection.bias, ref_bias), + ): + assert param.grad is not None + _assert_close(param.grad, ref_param.grad) + + +@pytest.mark.parametrize("batch_dim", (0, 1)) +@pytest.mark.parametrize("dtype", (torch.float32, torch.bfloat16)) +def test_long_sequence_reduction_is_repeatable(batch_dim, dtype): + """Exercise split condition-gradient reduction across multiple token tiles.""" + shape = (2, 257, 65) if batch_dim == 0 else (257, 2, 65) + inputs = _make_inputs(shape, batch_dim, dtype, cond_dtype=torch.float32) + reference_inputs = tuple(tensor.detach().double().requires_grad_() for tensor in inputs) + expected = _reference(*reference_inputs, batch_dim=batch_dim) + grad = torch.randn(shape, device="cuda", dtype=dtype) + expected.backward(grad.double()) + op = te_ops.AdaptiveLayerNorm(65, batch_dim=batch_dim) + previous = None + for _ in range(2): + actual = op(*inputs) + actual.backward(grad) + _assert_close(actual, expected) + for tensor, ref_tensor in zip(inputs, reference_inputs): + _assert_close(tensor.grad, ref_tensor.grad) + results = (actual.detach().clone(), *(tensor.grad.clone() for tensor in inputs)) + if previous is not None: + for result, old_result in zip(results, previous): + assert torch.equal(result, old_result) + previous = results + for tensor in inputs: + tensor.grad = None + + +@pytest.mark.parametrize("recipe_type", (DelayedScaling, Float8CurrentScaling)) +def test_fp8_sequential_matches_separate_operations(recipe_type): + """Preserve all gradients when FP8 linears surround the unquantized op.""" + available, reason = te.is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + dtype = torch.bfloat16 + operations = ( + te_ops.Linear(64, 64, device="cuda", dtype=dtype), + te_ops.AdaptiveLayerNorm(64, batch_dim=1), + te_ops.Linear(64, 64, device="cuda", dtype=dtype), + ) + separate = tuple(copy.deepcopy(op) for op in operations) + model = te_ops.Sequential(*operations) + inputs = _make_inputs((8, 2, 64), 1, dtype, cond_dtype=torch.float32) + reference_inputs = tuple(tensor.detach().clone().requires_grad_() for tensor in inputs) + with te.autocast(recipe=recipe_type()): + actual = model(*inputs) + with te.autocast(recipe=recipe_type()): + x, scale, shift = reference_inputs + expected = separate[2](separate[1](separate[0](x), scale, shift)) + grad = torch.randn_like(actual) + actual.backward(grad) + expected.backward(grad) + assert actual.dtype == dtype + _assert_close(actual, expected) + parameters = tuple(model.parameters()) + reference_parameters = tuple(param for op in separate for param in op.parameters()) + for tensor, reference_tensor in zip( + inputs + parameters, reference_inputs + reference_parameters + ): + assert tensor.grad is not None + assert reference_tensor.grad is not None + torch.testing.assert_close(tensor.grad, reference_tensor.grad, **_tolerances(dtype)) diff --git a/transformer_engine/common/triton/adaptive_layer_norm.py b/transformer_engine/common/triton/adaptive_layer_norm.py new file mode 100644 index 0000000000..f72e439fe1 --- /dev/null +++ b/transformer_engine/common/triton/adaptive_layer_norm.py @@ -0,0 +1,157 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Triton kernels for adaptive layer normalization.""" + +import triton +import triton.language as tl + + +@triton.jit +def _adaptive_layernorm_fwd_kernel( + x_ptr, + scale_ptr, + shift_ptr, + y_ptr, + mean_ptr, + rstd_ptr, + BATCH_SIZE: tl.constexpr, + BATCH_STRIDE: tl.constexpr, + HIDDEN_SIZE: tl.constexpr, + EPS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Normalize one token and apply its sample's scale and shift.""" + # Large diffusion activations can exceed 2**31 elements. + row = tl.program_id(0).to(tl.int64) + sample = (row // BATCH_STRIDE) % BATCH_SIZE + cols = tl.arange(0, BLOCK_SIZE) + mask = cols < HIDDEN_SIZE + x = tl.load(x_ptr + row * HIDDEN_SIZE + cols, mask, other=0).to(tl.float32) + mean = tl.sum(x, axis=0) / HIDDEN_SIZE + centered = tl.where(mask, x - mean, 0.0) + variance = tl.sum(centered * centered, axis=0) / HIDDEN_SIZE + rstd = tl.rsqrt(variance + EPS) + scale = tl.load(scale_ptr + sample * HIDDEN_SIZE + cols, mask, other=0).to(tl.float32) + shift = tl.load(shift_ptr + sample * HIDDEN_SIZE + cols, mask, other=0).to(tl.float32) + output = centered * rstd * (1.0 + scale) + shift + tl.store(y_ptr + row * HIDDEN_SIZE + cols, output, mask) + tl.store(mean_ptr + row, mean) + tl.store(rstd_ptr + row, rstd) + + +@triton.jit +def _adaptive_layernorm_dx_kernel( + dy_ptr, + x_ptr, + scale_ptr, + mean_ptr, + rstd_ptr, + dx_ptr, + BATCH_SIZE: tl.constexpr, + BATCH_STRIDE: tl.constexpr, + HIDDEN_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Compute the input gradient independently for each token.""" + # Large diffusion activations can exceed 2**31 elements. + row = tl.program_id(0).to(tl.int64) + sample = (row // BATCH_STRIDE) % BATCH_SIZE + cols = tl.arange(0, BLOCK_SIZE) + mask = cols < HIDDEN_SIZE + x = tl.load(x_ptr + row * HIDDEN_SIZE + cols, mask, other=0).to(tl.float32) + dy = tl.load(dy_ptr + row * HIDDEN_SIZE + cols, mask, other=0).to(tl.float32) + scale = tl.load(scale_ptr + sample * HIDDEN_SIZE + cols, mask, other=0).to(tl.float32) + mean = tl.load(mean_ptr + row) + rstd = tl.load(rstd_ptr + row) + normalized = tl.where(mask, (x - mean) * rstd, 0.0) + grad_normalized = dy * (1.0 + scale) + grad_mean = tl.sum(grad_normalized, axis=0) / HIDDEN_SIZE + grad_projection = tl.sum(grad_normalized * normalized, axis=0) / HIDDEN_SIZE + dx = (grad_normalized - grad_mean - normalized * grad_projection) * rstd + tl.store(dx_ptr + row * HIDDEN_SIZE + cols, dx, mask) + + +@triton.jit +def _adaptive_layernorm_condition_grads_kernel( + dy_ptr, + x_ptr, + mean_ptr, + rstd_ptr, + dscale_ptr, + dshift_ptr, + SEQUENCE_LENGTH: tl.constexpr, + BATCH_SIZE: tl.constexpr, + BATCH_STRIDE: tl.constexpr, + HIDDEN_SIZE: tl.constexpr, + NUM_SPLITS: tl.constexpr, + ROWS_PER_SPLIT: tl.constexpr, + COMPUTE_DSCALE: tl.constexpr, + COMPUTE_DSHIFT: tl.constexpr, + BLOCK_ROWS: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + """Compute disjoint sequence partials without atomic additions.""" + sample = tl.program_id(0).to(tl.int64) + split = tl.program_id(1) + col_block = tl.program_id(2) + cols = col_block * BLOCK_COLS + tl.arange(0, BLOCK_COLS) + col_mask = cols < HIDDEN_SIZE + row_offsets = tl.arange(0, BLOCK_ROWS) + scale_acc = tl.zeros((BLOCK_COLS,), dtype=tl.float32) + shift_acc = tl.zeros((BLOCK_COLS,), dtype=tl.float32) + split_start = split.to(tl.int64) * ROWS_PER_SPLIT + split_end = tl.minimum(split_start + ROWS_PER_SPLIT, SEQUENCE_LENGTH) + for row_start in range(0, tl.cdiv(ROWS_PER_SPLIT, BLOCK_ROWS)): + rows = split_start + row_start * BLOCK_ROWS + row_offsets + row_mask = rows < split_end + mask = row_mask[:, None] & col_mask[None, :] + input_rows = ( + (rows // BATCH_STRIDE) * (BATCH_SIZE * BATCH_STRIDE) + + sample * BATCH_STRIDE + + rows % BATCH_STRIDE + ) + offsets = input_rows[:, None] * HIDDEN_SIZE + cols[None, :] + dy = tl.load(dy_ptr + offsets, mask, other=0).to(tl.float32) + if COMPUTE_DSCALE: + x = tl.load(x_ptr + offsets, mask, other=0).to(tl.float32) + mean = tl.load(mean_ptr + input_rows, row_mask, other=0) + rstd = tl.load(rstd_ptr + input_rows, row_mask, other=0) + normalized = (x - mean[:, None]) * rstd[:, None] + scale_acc += tl.sum(dy * normalized, axis=0) + if COMPUTE_DSHIFT: + shift_acc += tl.sum(dy, axis=0) + offsets = (sample * NUM_SPLITS + split) * HIDDEN_SIZE + cols + if COMPUTE_DSCALE: + tl.store(dscale_ptr + offsets, scale_acc, col_mask) + if COMPUTE_DSHIFT: + tl.store(dshift_ptr + offsets, shift_acc, col_mask) + + +@triton.jit +def _adaptive_layernorm_reduce_condition_grads_kernel( + partial_dscale_ptr, + partial_dshift_ptr, + dscale_ptr, + dshift_ptr, + HIDDEN_SIZE: tl.constexpr, + NUM_SPLITS: tl.constexpr, + COMPUTE_DSCALE: tl.constexpr, + COMPUTE_DSHIFT: tl.constexpr, + BLOCK_SPLITS: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + """Combine sequence partials in a fixed reduction order.""" + sample = tl.program_id(0).to(tl.int64) + cols = tl.program_id(1) * BLOCK_COLS + tl.arange(0, BLOCK_COLS) + splits = tl.arange(0, BLOCK_SPLITS) + mask = (splits[:, None] < NUM_SPLITS) & (cols[None, :] < HIDDEN_SIZE) + offsets = (sample * NUM_SPLITS + splits[:, None]) * HIDDEN_SIZE + cols[None, :] + output_offsets = sample * HIDDEN_SIZE + cols + if COMPUTE_DSCALE: + partial = tl.load(partial_dscale_ptr + offsets, mask, other=0) + tl.store(dscale_ptr + output_offsets, tl.sum(partial, axis=0), cols < HIDDEN_SIZE) + if COMPUTE_DSHIFT: + partial = tl.load(partial_dshift_ptr + offsets, mask, other=0) + tl.store(dshift_ptr + output_offsets, tl.sum(partial, axis=0), cols < HIDDEN_SIZE) diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index dae15330c1..a8b5ac7f8e 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -18,6 +18,7 @@ SReGLU, SiLU, ) +from .adaptive_layer_norm import AdaptiveLayerNorm from .add_extra_input import AddExtraInput from .all_gather import AllGather from .all_reduce import AllReduce diff --git a/transformer_engine/pytorch/ops/basic/adaptive_layer_norm.py b/transformer_engine/pytorch/ops/basic/adaptive_layer_norm.py new file mode 100644 index 0000000000..7724324140 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/adaptive_layer_norm.py @@ -0,0 +1,155 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible operation for timestep-conditioned layer normalization.""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +import math +from typing import Any, Optional + +import torch + +from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload +from ...tensor import Quantizer +from ...triton.adaptive_layer_norm import adaptive_layernorm_bwd, adaptive_layernorm_fwd +from .._common import maybe_dequantize +from ..op import BasicOperation, OperationContext + + +class AdaptiveLayerNorm(BasicOperation): + r"""Layer normalization with a per-sample scale and shift. + + Normalizes the last dimension, then applies a timestep-conditioned affine + transform, as used in diffusion transformers: + + .. math:: + + y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \varepsilon}} + (1 + \mathrm{scale}) + \mathrm{shift} + + This operation has no parameters. Scale and shift are extra tensor inputs, + and gradients propagate to both. Call it as ``op(x, scale, shift)``, or + pass the same extra inputs to a containing ``ops.Sequential``. + + The input has at least two dimensions. Conditions have shape + ``(batch_size, hidden_size)``, or the same rank as the input with all + dimensions except the batch and hidden dimensions equal to one. For + example, batch-first input ``[B, S, H]`` accepts ``[B, H]`` or + ``[B, 1, H]``; sequence-first input ``[S, B, H]`` with + ``batch_dim=1`` accepts ``[B, H]`` or ``[1, B, H]``. + + CUDA float32, float16, and bfloat16 tensors are supported. Conditions can + have a different dtype from the input. Normalization, the addition + ``1 + scale``, and modulation are computed in float32 before casting + the result to the input dtype. In particular, small bfloat16 scales are + not rounded away by adding one in bfloat16. Gradients have the dtype of + their respective inputs. + + Parameters + ---------- + hidden_size : int + Size of the last input dimension, between 1 and 16384. + eps : float, default = 1e-5 + Non-negative value added to the variance for numerical stability. + batch_dim : int, default = 0 + Non-negative index of the batch dimension. All remaining dimensions + except the last dimension share the per-sample conditions. + + Notes + ----- + This operation produces an unquantized output. A following operation can + quantize it according to its quantization recipe. It does not fuse + normalization and GEMM into a single kernel. + """ + + num_extra_inputs: int = 2 + + def __init__( + self, + hidden_size: int, + *, + eps: float = 1e-5, + batch_dim: int = 0, + ) -> None: + super().__init__() + if not isinstance(hidden_size, int) or not 1 <= hidden_size <= 16384: + raise ValueError("hidden_size must be an integer between 1 and 16384.") + if not math.isfinite(eps) or eps < 0: + raise ValueError("eps must be finite and non-negative.") + if not isinstance(batch_dim, int) or batch_dim < 0: + raise ValueError("batch_dim must be a non-negative integer.") + self.hidden_size = hidden_size + self.eps = eps + self.batch_dim = batch_dim + + def op_forward(self, *args, **kwargs) -> None: + raise RuntimeError("AdaptiveLayerNorm uses fuser_forward for its two extra inputs.") + + def op_backward(self, *args, **kwargs) -> None: + raise RuntimeError("AdaptiveLayerNorm uses fuser_backward for its two extra inputs.") + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: + ctx = basic_op_ctxs[0] + scale, shift = basic_op_extra_inputs[0] + if ctx.requires_grad: + ctx.compute_dscale = scale.requires_grad + ctx.compute_dshift = shift.requires_grad + if input_.ndim < 2 or input_.shape[-1] != self.hidden_size: + raise ValueError( + f"Input must have at least two dimensions and last dimension {self.hidden_size}, " + f"got {tuple(input_.shape)}." + ) + x = maybe_dequantize(input_).contiguous() + scale = maybe_dequantize(scale).contiguous() + shift = maybe_dequantize(shift).contiguous() + output, mean, rstd = adaptive_layernorm_fwd( + x, scale, shift, self.eps, batch_dim=self.batch_dim + ) + if ctx.requires_grad: + if is_cpu_offload_enabled(): + mark_activation_offload(x, scale, mean, rstd) + ctx.save_for_backward(x, scale, mean, rstd) + ctx.shift_shape = tuple(shift.shape) + ctx.shift_dtype = shift.dtype + return output, [()] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + ctx = basic_op_ctxs[0] + x, scale, mean, rstd = ctx.saved_tensors + # Intermediate tensors inside the fuser need not have requires_grad set, + # even when their producers require gradients. + dx, dscale, dshift = adaptive_layernorm_bwd( + maybe_dequantize(grad_output).contiguous(), + x, + scale, + mean, + rstd, + shift_shape=ctx.shift_shape, + shift_dtype=ctx.shift_dtype, + compute_dscale=ctx.compute_dscale, + compute_dshift=ctx.compute_dshift, + batch_dim=self.batch_dim, + ) + return dx, [()], [(dscale, dshift)] diff --git a/transformer_engine/pytorch/triton/adaptive_layer_norm.py b/transformer_engine/pytorch/triton/adaptive_layer_norm.py new file mode 100644 index 0000000000..48d642e802 --- /dev/null +++ b/transformer_engine/pytorch/triton/adaptive_layer_norm.py @@ -0,0 +1,252 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PyTorch launchers for adaptive layer normalization Triton kernels.""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch +import triton + +from transformer_engine.common.triton.adaptive_layer_norm import ( + _adaptive_layernorm_fwd_kernel, + _adaptive_layernorm_dx_kernel, + _adaptive_layernorm_condition_grads_kernel, + _adaptive_layernorm_reduce_condition_grads_kernel, +) + + +def _check_input(x: torch.Tensor, batch_dim: int) -> tuple[int, int, int, int]: + """Obtain batch size, tokens per sample, hidden size, and batch stride.""" + if x.device.type != "cuda": + raise ValueError("Adaptive layer normalization requires CUDA tensors.") + if x.dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise TypeError(f"Unsupported input dtype for adaptive layer normalization: {x.dtype}.") + if x.ndim < 2: + raise ValueError("Input must have shape (batch, ..., hidden_size).") + if not 0 <= batch_dim < x.ndim - 1: + raise ValueError("batch_dim must identify a non-normalized input dimension.") + batch_size, hidden_size = x.shape[batch_dim], x.shape[-1] + if not 1 <= hidden_size <= 16384: + raise ValueError("Adaptive layer normalization supports hidden sizes from 1 to 16384.") + sequence_length = math.prod(size for dim, size in enumerate(x.shape[:-1]) if dim != batch_dim) + batch_stride = math.prod(x.shape[batch_dim + 1 : -1]) + return batch_size, sequence_length, hidden_size, batch_stride + + +def _check_condition( + condition: torch.Tensor, + x: torch.Tensor, + name: str, + batch_dim: int, +) -> None: + """Check a per-sample condition, including an optional singleton sequence shape.""" + compact_shape = (x.shape[batch_dim], x.shape[-1]) + expanded_shape = tuple( + size if dim in (batch_dim, x.ndim - 1) else 1 for dim, size in enumerate(x.shape) + ) + if tuple(condition.shape) not in (compact_shape, expanded_shape): + raise ValueError( + f"{name} must have shape {compact_shape} or {expanded_shape}, " + f"got {tuple(condition.shape)}." + ) + if condition.device != x.device: + raise ValueError(f"{name} must be on the same CUDA device as the input.") + if condition.dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise TypeError(f"Unsupported {name} dtype: {condition.dtype}.") + + +def adaptive_layernorm_fwd( + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, + *, + batch_dim: int = 0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Apply adaptive layer normalization and return output, mean, and reciprocal stddev. + + The last input dimension is normalized; batch_dim identifies the batch. + Conditions have shape (batch, hidden_size), or match the input rank with + singleton dimensions except for batch_dim and the last dimension. + Normalization and modulation use float32 arithmetic and the output has the + input dtype. Mean and reciprocal stddev are flattened float32 + tensors, with one value per input row. + + These launchers do not register an autograd formula; callers must use + adaptive_layernorm_bwd to propagate gradients. + """ + batch_size, sequence_length, hidden_size, batch_stride = _check_input(x, batch_dim) + _check_condition(scale, x, "scale", batch_dim) + _check_condition(shift, x, "shift", batch_dim) + if not math.isfinite(eps) or eps < 0: + raise ValueError("eps must be finite and non-negative.") + x = x.contiguous() + scale = scale.contiguous() + shift = shift.contiguous() + output = torch.empty(x.shape, device=x.device, dtype=x.dtype) + mean = torch.empty(batch_size * sequence_length, device=x.device, dtype=torch.float32) + rstd = torch.empty_like(mean) + if batch_size * sequence_length == 0: + return output, mean, rstd + block_size = triton.next_power_of_2(hidden_size) + num_warps = min(8, max(1, block_size // 256)) + with torch.cuda.device(x.device): + _adaptive_layernorm_fwd_kernel[(batch_size * sequence_length,)]( + x, + scale, + shift, + output, + mean, + rstd, + BATCH_SIZE=batch_size, + BATCH_STRIDE=batch_stride, + HIDDEN_SIZE=hidden_size, + EPS=eps, + BLOCK_SIZE=block_size, + num_warps=num_warps, + ) + return output, mean, rstd + + +def adaptive_layernorm_bwd( + dy: torch.Tensor, + x: torch.Tensor, + scale: torch.Tensor, + mean: torch.Tensor, + rstd: torch.Tensor, + *, + shift_shape: tuple[int, ...], + shift_dtype: torch.dtype, + compute_dscale: bool = True, + compute_dshift: bool = True, + batch_dim: int = 0, +) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """Compute input and optional condition gradients with deterministic reductions. + + batch_dim must match the forward call. shift_shape and shift_dtype describe + the forward shift input. Its values are not needed for the backward pass. + Each condition gradient has the shape and dtype of its corresponding input. Split sequence reductions use a + bounded float32 workspace and never use atomic additions. + """ + batch_size, sequence_length, hidden_size, batch_stride = _check_input(x, batch_dim) + _check_condition(scale, x, "scale", batch_dim) + if dy.shape != x.shape or dy.device != x.device: + raise ValueError("Output gradient must have the input shape and device.") + if dy.dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise TypeError(f"Unsupported output gradient dtype: {dy.dtype}.") + expected_stats_shape = (batch_size * sequence_length,) + for name, stats in (("mean", mean), ("rstd", rstd)): + if ( + tuple(stats.shape) != expected_stats_shape + or stats.dtype != torch.float32 + or stats.device != x.device + or not stats.is_contiguous() + ): + raise ValueError(f"{name} must be a contiguous float32 row-statistics tensor.") + compact_shape = (batch_size, hidden_size) + expanded_shape = tuple( + size if dim in (batch_dim, x.ndim - 1) else 1 for dim, size in enumerate(x.shape) + ) + if tuple(shift_shape) not in (compact_shape, expanded_shape): + raise ValueError("Shift gradient shape must describe a per-sample condition.") + if shift_dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise TypeError(f"Unsupported shift gradient dtype: {shift_dtype}.") + x = x.contiguous() + dy = dy.contiguous() + scale = scale.contiguous() + dx = torch.empty(x.shape, device=x.device, dtype=x.dtype) + dscale = ( + torch.empty(scale.shape, device=x.device, dtype=scale.dtype) if compute_dscale else None + ) + dshift = ( + torch.empty(shift_shape, device=x.device, dtype=shift_dtype) if compute_dshift else None + ) + if batch_size * sequence_length == 0: + if dscale is not None: + dscale.zero_() + if dshift is not None: + dshift.zero_() + return dx, dscale, dshift + + block_size = triton.next_power_of_2(hidden_size) + num_warps = min(8, max(1, block_size // 256)) + with torch.cuda.device(x.device): + _adaptive_layernorm_dx_kernel[(batch_size * sequence_length,)]( + dy, + x, + scale, + mean, + rstd, + dx, + BATCH_SIZE=batch_size, + BATCH_STRIDE=batch_stride, + HIDDEN_SIZE=hidden_size, + BLOCK_SIZE=block_size, + num_warps=num_warps, + # H=1 has identically zero dX. Contracting its equal terms into + # separate FMAs can leave a residual amplified by reciprocal stddev. + enable_fp_fusion=hidden_size != 1, + ) + if compute_dscale or compute_dshift: + num_splits = min(32, triton.cdiv(sequence_length, 128)) + rows_per_split = triton.cdiv(sequence_length, num_splits) + # With one split the partial kernel stores directly in the outputs. + # Unused pointers alias dy and are removed by constexpr specialization. + partial_shape = (batch_size, num_splits, hidden_size) + partial_dscale = dy + partial_dshift = dy + if compute_dscale: + partial_dscale = ( + dscale + if num_splits == 1 + else torch.empty(partial_shape, device=x.device, dtype=torch.float32) + ) + if compute_dshift: + partial_dshift = ( + dshift + if num_splits == 1 + else torch.empty(partial_shape, device=x.device, dtype=torch.float32) + ) + _adaptive_layernorm_condition_grads_kernel[ + (batch_size, num_splits, triton.cdiv(hidden_size, 128)) + ]( + dy, + x, + mean, + rstd, + partial_dscale, + partial_dshift, + SEQUENCE_LENGTH=sequence_length, + BATCH_SIZE=batch_size, + BATCH_STRIDE=batch_stride, + HIDDEN_SIZE=hidden_size, + NUM_SPLITS=num_splits, + ROWS_PER_SPLIT=rows_per_split, + COMPUTE_DSCALE=compute_dscale, + COMPUTE_DSHIFT=compute_dshift, + BLOCK_ROWS=32, + BLOCK_COLS=128, + num_warps=4, + ) + if num_splits > 1: + _adaptive_layernorm_reduce_condition_grads_kernel[ + (batch_size, triton.cdiv(hidden_size, 128)) + ]( + partial_dscale, + partial_dshift, + dscale if dscale is not None else dy, + dshift if dshift is not None else dy, + HIDDEN_SIZE=hidden_size, + NUM_SPLITS=num_splits, + COMPUTE_DSCALE=compute_dscale, + COMPUTE_DSHIFT=compute_dshift, + BLOCK_SPLITS=triton.next_power_of_2(num_splits), + BLOCK_COLS=128, + num_warps=4, + ) + return dx, dscale, dshift